{-# LANGUAGE Rank2Types, NoMonomorphismRestriction #-}

-- Various examples of Tagless-staged, in particular
-- staged eta (the example from our PEPM08 paper) and staged twice-eta

module TSEx where

import TSCore
import Control.Applicative


ex0 = lam(\x -> int 1 +: var x)
ex0c = "\\x_0 -> (GHC.Num.+) 1 x_0" == runCI ex0

exA2 :: (Applicative m, SSym repr, LamPure repr) =>
     J m (HV h repr) (Int -> Int -> Int)
exA2 = lam(\x -> lam(\y -> weaken (var x) +: var y))
exA2c = "\\x_0 -> \\x_1 -> (GHC.Num.+) x_0 x_1" == runCI exA2

-- Checking of printing and sharing

-- ex11 :: (Applicative m, SSym repr, LamPure repr) =>
--      J m (HV h repr) (Int -> Int -> Int)
ex11 = lam (\x -> lam (\y -> (weaken (var x) *: weaken (var x)) +:
		                (var y *: var y)))

ex11r = 25 == runRI ex11 3 4
ex11c = 
 "\\x_0 -> \\x_1 -> (GHC.Num.+) ((GHC.Num.*) x_0 x_0) ((GHC.Num.*) x_1 x_1)"
 == runCI ex11

-- Naming of variables: we indeed use levels rather than the global
-- counter
ex12 = "(\\x_0 -> x_0) (\\x_0 -> x_0)" 
       == runCI (lam (\x -> var x) $$ lam (\x -> var x))



-- The ef example from our PEPM08 paper
{-
ef = \z -> <\x -> ~z + x>
ef1 = ef <1>
ef2 = <\x y -> ~(ef <x *  y>)>

ef1 evaluates to <\x -> 1 + x> whereas ef2 evaluates to
<\x y -> \x' ->  x * y + x'>. In the latter result,
we need to distinguish two later-stage variables named x.
To maintain hygiene and \alpha-equivalence, we must lexically
link each use of a variable to its binding occurrence and
rename variables if their names clash.
-}


ef :: (Applicative m, SSym repr, LamPure repr) =>
      J m (HV h repr) Int -> J m (HV h repr) (Int -> Int)
ef z = lam (\x -> weaken z +: var x)


ef1 = ef (int 1)

-- ef2 :: (Applicative m, MLam m repr) =>
--      J m (HV h repr) (Int -> Int -> Int -> Int)
ef2 = lam (\x -> lam (\y -> ef (weaken (var x) *: var y)))

-- Check to see what ef1 and ef2 actually do.

-- instantiate repr to R

ef2r = 10 == runRI ef2 2 3 4
-- -- 10, which is 2*3 + 4 

ef1c = "\\x_0 -> (GHC.Num.+) 1 x_0" == runCI ef1

ef2c = "\\x_0 -> \\x_1 -> \\x_2 -> (GHC.Num.+) ((GHC.Num.*) x_0 x_1) x_2"
       == runCI ef2


-- Effectful code generation
-- First example: printing when generating pure functions
-- addt is the `version' of add that prints trace at the generation time

trace :: String -> J IO c a -> J IO c a
trace msg m = J $ putStrLn msg >> (unJ m)


addt :: (SSym repr) => 
	J IO (HV h repr) (Int -> Int -> Int)
addt = trace "generating add" add


ex1 :: (SSym repr, LamPure repr) => 
       J IO (HV h repr) (Int -> Int)
ex1 = lam (\x -> addt $$ var x $$ int 1)

-- plusM is printed only once!
ex1r = do
       f <- runR ex1
       putStrLn "Generation is complete"
       print (f 1)
       print (f 2)

ex1c = runC ex1


-- the placement of weaken below is type-directed

ex2 :: (SSym repr, LamPure repr) =>
       J IO (HV h repr) (Int -> Int -> Int)
ex2 = lam (\x -> lam (\y -> addt $$ weaken (var x) $$ (var y)))

-- plusM is printed only once!
ex2r = do
       f <- runR ex2
       putStrLn "Generation is complete"
       print (f 1 2)
       print (f 3 4)
{-
*TaglessStaged> ex2r
generating add
Generation is complete
3
7
-}

ex2c = runC ex2

{-
generating add
\x_0 -> \x_1 -> (GHC.Num.+) x_0 x_1
-}


-- Cannot leak out a variable
-- ex3 = lam (\x -> weaken (var x))



-- The eta example

-- All functions from code -> code have to have rank-2 type (because of the
-- parameter s marking the variables)
-- Therefore, we have to write signatures
eta :: (Applicative m, SSym repr, LamPure repr) => 
       (forall h1. J m (HV (h1, h) repr) a -> J m (HV (h1, h) repr) b)
       -> J m (HV h repr) (a->b)
eta f = lam (\x -> f (var x))

teta0 () = eta (\z -> int 1 +: z)
teta0c = "\\x_0 -> (GHC.Num.+) 1 x_0" 
	 == runCI (teta0 ())

teta1 () = lam (\y -> eta (\z -> z +: weaken (var y)))
teta1c = "\\x_0 -> \\x_1 -> (GHC.Num.+) x_1 x_0"
	 == runCI (teta1 ())
teta1r = 5 == runRI (teta1 ()) 2 3

-- Instantiating eta with an effectful function, which has the effect
-- at the generation time!
teta1RIO = do
	   teta <- runR $ lam (\y -> eta (\z -> addt $$ z $$ weaken (var y)))
	   putStrLn "Generation complete"
	   print $ teta 2 3
{-
*TaglessStaged> teta1RIO
generating add
Generation complete
5
-}



-- .< fun y -> fun w -> 
--              .~(eta (fun z -> .< .~z + y*w >.)) >.)

teta2 = lam (\y -> lam (\w -> eta (\z -> 
            z +: (weaken (weaken (var y)) *: weaken (var w)))))
teta2c = "\\x_0 -> \\x_1 -> \\x_2 -> (GHC.Num.+) x_2 ((GHC.Num.*) x_0 x_1)"
	 == runCI teta2

teta2r = 10 == runRI teta2 2 3 4  -- 10

-- Use generic weaking
teta2' = lam (\y -> lam (\w -> eta (\z -> 
            z +: (weakens (var y)) *: weakens (var w))))
teta2c' = "\\x_0 -> \\x_1 -> \\x_2 -> (GHC.Num.+) x_2 ((GHC.Num.*) x_0 x_1)"
	 == runCI teta2'


-- .< fun x u -> 
--              .~(eta (fun z -> .<fun y -> .~z + u*x*y >.)) >.

-- See teta3' below for the example of eliminating repeated weaken
teta3 = lam(\x -> lam(\u -> eta (\z ->
            lam(\y -> weaken z +: (w2 u *: w3 x *: var y))
              )))
 where w2 = weaken . weaken . var
       w3 = weaken . weaken . weaken . var
 

teta3c = 
   "\\x_0 -> \\x_1 -> \\x_2 -> \\x_3 -> " ++
      "(GHC.Num.+) x_2 ((GHC.Num.*) ((GHC.Num.*) x_1 x_0) x_3)"
      == runCI teta3
teta3r = 34 == runRI teta3 2 3 4 5  -- 34

teta3' = lam(\x -> lam(\u -> eta (\z ->
            lam(\y -> weakens z +: (vr u *: vr x *: vr y))
              )))
 where vr = weakens . var

teta3c' = 
   "\\x_0 -> \\x_1 -> \\x_2 -> \\x_3 -> " ++
      "(GHC.Num.+) x_2 ((GHC.Num.*) ((GHC.Num.*) x_1 x_0) x_3)"
      == runCI teta3'
 



-- Illustrating Hyland's Disjointness Property or its violation:
-- no binder ends up within a copy of itself if we stay within 
-- residual theory (of the \-calculus), i.e., if we do not contract 
-- newly created redexes.
-- Implement "(\x.x x) (\y.\z.y z) ->> \z1.\z2.z1 z2"
-- The function (\y.\z.y z) with the right staging annotations
-- is eta

-- Inferred type:
-- eta_refl  :: (Applicative m, SSym repr, LamPure repr) =>
--      J m (HV h repr) (a -> b) -> J m (HV h repr) (a -> b)

eta_refl = \y -> eta (\u -> weaken y $$ u)

-- etaeta  :: (Applicative m, SSym repr, LamPure repr) =>
--      J m (HV h repr) ((a -> b) -> a -> b)
etaeta = eta eta_refl

etaetac = "\\x_0 -> \\x_1 -> x_0 x_1"
	  == runCI etaeta

etaetaetac = "\\x_0 -> (\\x_1 -> \\x_2 -> x_1 x_2) x_0"
	  == runCI (eta_refl (eta eta_refl))






-- Second-order eta
-- We need to essentially curry and uncurry
eta2 :: (Applicative m, SSym repr, LamPure repr) => 
       (forall h1. J m (HV (h1, h) repr) a -> J m (HV (h1, h) repr) b
	           -> J m (HV (h1, h) repr) c)
       -> J m (HV h repr) (a->b->c)
eta2 f = lam (\x -> lam (\y -> 
			 unpack $ f (pack (weaken (var x))) (pack (var y))))
 where
 pack :: (Applicative m, SSym repr) => 
	 J m (HV (h1,(h2,h)) repr) a -> J m (HV ((h1,h2),h) repr) a
 pack = J . fmap (hmap (\ ((h1,h2),h) -> (h1,(h2,h)))) . unJ
 unpack :: (Applicative m, SSym repr) => 
	   J m (HV ((h1,h2),h) repr) a -> J m (HV (h1,(h2,h)) repr) a
 unpack = J . fmap (hmap (\ (h1,(h2,h)) -> ((h1,h2),h))) . unJ


teta22 = lam (\x -> lam (\y -> eta2 (\z1 z2 -> 
                  (z1 *: weakens (var x)) +:
                  (z2 *: weakens (var y)))))


teta22c = "\\x_0 -> \\x_1 -> \\x_2 -> \\x_3 -> "++
            "(GHC.Num.+) ((GHC.Num.*) x_2 x_0) ((GHC.Num.*) x_3 x_1)"
	    == runCI teta22

-- Attempting to write eta2 without pack/unpack
eta2' :: (Applicative m, SSym repr, LamPure repr) => 
       (forall h'. Extends h h' =>
	J m (HV h' repr) a -> J m (HV h' repr) b -> J m (HV h' repr) c)
       -> J m (HV h repr) (a->b->c)
eta2' f = lam (\x -> lam (\y -> f (weaken (var x)) (var y)))



teta2'0 = eta2' (\z1 z2 -> (z1 *: int 1) +: (z2 *: int 2))

teta2'0c =
    "\\x_0 -> \\x_1 -> (GHC.Num.+) ((GHC.Num.*) x_0 1) ((GHC.Num.*) x_1 2)"
    == runCI teta2'0

teta2'1 = lam (\y1 -> eta2' (\z1 z2 -> 
			     (z1 *: weakens (var y1)) +: (z2 *: int 2)))

teta2'1c = ("\\x_0 -> \\x_1 -> \\x_2 -> "++
	    "(GHC.Num.+) ((GHC.Num.*) x_1 x_0) ((GHC.Num.*) x_2 2)")
    == runCI teta2'1


-- Alas, here it fails: the compiler cannot make a simple 
-- inference of transitivity:
-- if hh extends (x1,(x2,h)) than hh certainly extends (x2,h)
{-

teta2'2 = lam (\y1 -> lam (\y2 -> eta2' (\z1 z2 -> 
			     (z1 *: weakens (var y1)) +:
                             (z2 *: weakens (var y2)))))

    Could not deduce (Extends (H repr s2 Int, h) h'1)
      from the context (Extends (H repr s3 Int, (H repr s2 Int, h)) h'1)
      arising from a use of `weakens'
                   at /home/oleg/papers/MetaFX/hybrid/TSEx.hs:265:15-30
-}

-- We have to help the compiler out, using u and type assignment
-- We stress that we do NOT do any manipulation on values;
-- me merely specialize types to help the compiler draw the transitivity
-- inference
teta2'2 = lam (\y1 -> lam (\y2 -> let u = weaken (var y1) `asTypeOf` (var y2) in
			   eta2' (\z1 z2 -> 
			     (z1 *: weakens u) +:
                             (z2 *: weakens (var y2)))))

teta2'2c = ("\\x_0 -> \\x_1 -> \\x_2 -> \\x_3 -> "++
	    "(GHC.Num.+) ((GHC.Num.*) x_2 x_0) ((GHC.Num.*) x_3 x_1)")
	   == runCI teta2'2




-- Tests of the let-form

tl0 = let x = (trace "Let RHS" (int 1 +: int 2)) in x *: x

tl0r = fmap (9 ==) $ runR tl0
{-
Let RHS
Let RHS
True
-}

tl0c = fmap ("(GHC.Num.*) ((GHC.Num.+) 1 2) ((GHC.Num.+) 1 2)" == ) $
       runC tl0
{-
Let RHS
Let RHS
-}

tl1 = let_ (trace "Let RHS" (int 1 +: int 2)) $ \x -> var x *: var x

-- The message is printed only once! And our generation effect is applicative,
-- not a monad!
tl1r = fmap (9 ==) $ runR tl1
{-
Let RHS
6
-}


-- Again, the trace message is printed only once
tl1c = fmap ("let z_0 = (GHC.Num.+) 1 2\n in (GHC.Num.*) z_0 z_0"==) $ runC tl1
{-
Let RHS
let z_0 = 1
 in (GHC.Num.+) z_0 z_0
-}

-- How to deal with polymoprhic let?
-- The following code is essentially
--   let tl2a = lam (\x -> var x) in tl2a $$ tl2a
-- that handles polymorphism by inlining.

tl2a :: (SSym repr, LamPure repr, Applicative m) =>
     J m (HV h repr) (a -> a)
tl2a = lam (\x -> var x)

tl2 = tl2a $$ tl2a
tl2c = "(\\x_0 -> x_0) (\\x_0 -> x_0)" == runCI tl2

-- To do otherwise, we first need a polymorphic let in the object
-- language. We need new lam, for example, of the type
-- lamP :: (forall a s.
--  HV (H repr s A a, h) repr (A a) 
--       -> J m (HV (H repr s (A a), h) repr) (B a))
--     -> J m (HV h repr) (Arr A B)
-- where A and B are some type constructors (not necessarily functors:
-- perhaps newtype over type function).