{-# LANGUAGE RankNTypes, NoMonomorphismRestriction, 
  ExistentialQuantification #-}

-- Code generation with control effects reaching beyond the
-- closest binder. Far-reaching let-insertion.
-- The implementation of the CPSA applicative hierarchy.

module TSCPST where

import TSCore
import Control.Applicative

-- A bizarre answer-type changing, sort of, CPS type
-- which happens to be applicative!
newtype CPSA w m a = 
    CPSA{unCPSA :: 
	    forall hw. (forall h1. m (h1->hw->a) -> m (h1->hw->w))
	    -> m (hw -> w)}

-- Instantiate h1 to ()
throw :: Functor m =>
	 (forall h1. m (h1->hw-> a) -> m (h1 -> hw->w)) -> m a -> m (hw->w)
throw k x = fmap ($ ()) $ k (fmap (\xv () _ -> xv) x)

-- How I HATE having to write this instance!!
instance (Applicative m) => Functor (CPSA w m) where
    fmap f m = pure f <*> m

instance Applicative m => Applicative (CPSA w m) where
    pure x = CPSA $ \k -> throw k (pure x)
    CPSA f <*> CPSA m = CPSA $ \k ->
      f (\fv -> fmap (\mr hf kw -> mr (hf,kw)) $ m (\mv -> 
       fmap (\kr hm (hf,kw) -> kr (hm,(hf,kw)) kw) 
         (k (comb <$> fv <*> mv))))
      where comb fvb mvb = \ (hm,(hf,kw)) _ -> (fvb hf kw) (mvb hm (hf,kw))


reset :: (Applicative m) => CPSA a m a -> CPSA w m a
reset m = CPSA $ \k -> throw k $ runCPSA m 

resetJ = J . reset . unJ

-- Instantiating hw to ()
runCPSA :: Applicative m => CPSA a m a -> m a
runCPSA m = fmap ($ ()) $ unCPSA m id

runJCPSA = J . runCPSA . unJ


exA2 = lam(\x -> lam(\y -> weaken (var x) +: var y))
exA2c = "\\x_0 -> \\x_1 -> (GHC.Num.+) x_0 x_1" == 
	(runCI . J . runCPSA . unJ $ exA2)


liftCA :: (Functor m, Monad m) => m a -> CPSA w m a
liftCA m = CPSA $ \k -> m >>= throw k . return

liftA1 :: (Functor m) => m a -> CPSA w m a
liftA1 m = CPSA $ \k -> throw k m

liftJA = J . liftA1 . unJ

trace :: String -> IO a -> IO a
trace msg m = putStrLn msg >> m


addt :: (SSym repr) => 
	J (CPSA w IO) (HV h repr) (Int -> Int -> Int)
addt = J $ liftCA $ trace "generating add" (unJ add)

ex0' = lam(\x -> addt $$ int 1 $$ var x)
ex0'c = fmap ("\\x_0 -> (GHC.Num.+) 1 x_0" ==) $ runC (runJCPSA ex0')


e1 :: (SSym repr) => J (CPSA w IO) (HV h repr) Int
e1 = addt $$ int 2 $$ int 3


glet0 = J (CPSA (\k -> 
	(unCPSA . unJ $ let_ e1 (\x -> 
   		 J (CPSA (\k0 -> 
		   throw k0 $
                   fmap (\kr -> J(\ (h1,(hw,h)) -> unJ (kr h1 hw) h))
				 (k (pure (\h1 _ -> J(\h -> unJ x (h1,())))))))
		       ))
        (\x -> fmap (\xv h1 hw -> J(\h -> unJ (xv h1 hw) (hw,h))) x)))


ex401 =  int 1 +: glet0
ex401c = fmap ("let z_0 = (GHC.Num.+) 2 3\n in (GHC.Num.+) 1 z_0" ==)
        (runC (runJCPSA ex401))

ex411 = resetJ $ lam (\x -> int 1 +: glet0)
ex411c = fmap ("let z_0 = (GHC.Num.+) 2 3\n in \\x_1 -> (GHC.Num.+) 1 z_0" ==)
        (runC (runJCPSA ex411))


genlet :: (SSym repr, SymLet repr, Applicative m) =>
     J (CPSA (HV hw repr w) m) (HV hw repr) a
     -> J (CPSA (HV hw repr w) m) (HV ha repr) a
genlet e = J (CPSA (\k -> 
	(unCPSA . unJ $ let_ (weaken e) (\x -> 
   		 J (CPSA (\k0 -> 
		   throw k0 $
                   fmap (\kr -> J(\ (h1,(hw,h)) -> unJ (kr h1 hw) h))
				 (k (pure (\h1 _ -> J(\h -> unJ x (h1,())))))))
		       ))
        (\x -> fmap (\xv h1 hw -> J(\h -> unJ (xv h1 hw) (hw,h))) x)
  ))

tt0:: (SSym repr, SymLet repr) =>
     J (CPSA (HV hw repr w) IO) (HV ha repr) Int
tt0 = int 1 +: genlet e1
tt0c = (runC (runJCPSA tt0))

-- the first example of nested genlet!

tt1:: (SSym repr, SymLet repr) =>
     J (CPSA (HV hw repr w) (CPSA (HV hw1 repr w1) IO)) (HV ha repr) Int
tt1 = int 1 +: genlet (liftJA glet0)

tt1c = fmap ("let z_0 = (GHC.Num.+) 2 3\n in let z_1 = z_0\n"++
	     "     in (GHC.Num.+) 1 z_1" ==)
       $ runC (runJCPSA (runJCPSA tt1))


-- The example from TSCPS.hs
ex49 = lam (\x -> resetJ (lam (\y ->
         int 1 +: genlet (var x +: genlet (int 3 +: int 4))
       +: genlet (int 5 +: int 6))))

ex49c = ("\\x_0 -> let z_1 = (GHC.Num.+) 3 4\n"++
	 "         in let z_2 = (GHC.Num.+) x_0 z_1\n"++
	 "             in let z_3 = (GHC.Num.+) 5 6\n"++
	 "                 in \\x_4 -> (GHC.Num.+) ((GHC.Num.+) 1 z_2) z_3") 
	== runCI (runJCPSA ex49)


-- Now we add two liftJA

-- The type is now lifted....
ex50  :: (SSym repr, Applicative m, SymLet repr, LamPure repr) =>
     J (CPSA w1 (CPSA (HV hw repr w) m)) (HV h repr) (Int -> a -> Int)

ex50 = lam (\x -> resetJ (lam (\y ->
         int 1 +: genlet (var x +: (liftJA $ genlet (int 3 +: int 4)))
       +: (liftJA $ genlet (int 5 +: int 6)))))


-- And two let appear on the outside the binder!
-- So, several let-insertion each crosses the different number of
-- binders (including the binders of the previous let)
ex50c = ("let z_0 = (GHC.Num.+) 3 4\n"++
	 " in let z_1 = (GHC.Num.+) 5 6\n"++
	 "     in \\x_2 -> let z_3 = (GHC.Num.+) x_2 z_0\n"++
	 "                 in \\x_4 -> (GHC.Num.+) ((GHC.Num.+) 1 z_3) z_1"==)
	 $ runCI (runJCPSA (runJCPSA ex50))