{-# LANGUAGE ExistentialQuantification, ScopedTypeVariables, RankNTypes, GADTs, Arrows  #-}

-- We can't use arrow notation because RebindableSyntax 
-- doesn't work well with Arrows (yet)

module CSE where

import Language.AbstractSyntax.TTTAS2
import Prelude hiding (lookup)
-- import Control.Arrow

ifThenElse c t e = case c of
                   True -> t
                   False -> e
data Expr  a env where
     Var       ::  Ref a env      -> Expr a env
     IntVal    ::  Int            -> Expr Int env
     BoolVal   ::  Bool           -> Expr Bool env 
     Cons      ::  Expr a env     -> Expr [a]  env 
               ->  Expr [a] env
     Nil       ::  Expr [a] env
     Add       ::  Expr Int env   -> Expr Int env         
               ->  Expr Int env
     LessThan  ::  Expr Int env   -> Expr Int env                 
               ->  Expr Bool env 
     If        ::  Expr Bool env  -> Expr a env          
               ->  Expr a env     ->  Expr a env

eval    :: Expr  a env   -> env ->  a
eval (Var r)         e   =  lookup r e
eval (IntVal i)      _   =  i
eval (BoolVal b)     _   =  b
eval (Add x y)       e   =  eval x e + eval y e
eval (Cons x y)      e   =  eval x e : eval y e
eval Nil             _   =  []
eval (LessThan x y)  e   =  eval x e < eval y e
eval (If x y z)      e   =  if    eval x e 
                            then  eval y e
                            else  eval z e

type Decls env  = Env Expr env env

data TDecls env = forall env' . TDecls  (Decls env') 
                                        (T env env')

a = Suc Zero
b = Zero
exampledecls :: Decls (((),Int),Int)
exampledecls = 
 Empty  `Ext` (IntVal 4) 
        `Ext` (Add  (Add (Var a) (IntVal 4)) 
                    (Add (Var a) (IntVal 4)))

resdecls :: TDecls (((),Int),Int)
resdecls = cse exampledecls


evalDecls   :: Decls env -> env
evalVar     :: Ref a env -> TDecls env -> a
evalVar var (TDecls ds (T tt))
 = lookup (tt var) (evalDecls ds)

value_a = evalVar a resdecls
value_b = evalVar b resdecls

evalDecls (ds :: Decls env) = 
            let result :: env
                result = evalD ds
                evalD :: Env Expr env def -> def
                evalD Empty = ()
                evalD (Ext ds e) = (evalD ds,eval e result)
            in result

resdecls2 :: TDecls (((),Int),Int)  
resdecls2 =  

  TDecls
  (  Empty  `Ext` (IntVal 4) 
            `Ext` (Add  (Var (Suc (Suc Zero))) 
                        (Var (Suc (Suc Zero))))
            `Ext` (Add  (Var (Suc Zero)) 
                        (Var (Suc Zero)))
  )
  (  T (\ref -> case ref of  
                  Zero      -> Zero 
                  Suc Zero  -> Suc (Suc Zero))
     ::  T (((), Int), Int) ((((), Int), Int), Int))

instance Show (Ref a env) where
    show x = "#" ++ (show $ refint x) 
refint :: Ref a env -> Int
refint Zero = 0
refint (Suc x) = 1 + refint x
instance Show (Expr a env) where
    show (Var r) = "v" ++ (show $ refint r)
    show (IntVal i) = show i
    show (BoolVal i) = show i
    show (Add x y) = "(" ++ show x ++ "+"++ show y ++ ")"
    show (LessThan x y) = "(" ++ show x ++ "<"++ show y ++ ")"
    show (If x y z) = "if" ++ show x ++ "then"++ show y ++ "else" ++ show z

equals  ::  Expr a env -> Expr b env ->  Maybe (Equal a b)
equals (Var r1) (Var r2)          = match r1 r2
equals (IntVal i1) (IntVal i2) 
                        | i1==i2  = Just Eq
equals (LessThan x1 y1) (LessThan x2 y2) 
       = do  Eq <- equals x1 x2 
             Eq <- equals y1 y2 
             return Eq 

equals (BoolVal b1) (BoolVal b2) 
                        | b1==b2  = Just Eq
equals (Add x1 y1) (Add x2 y2) = do
                                      Eq <- equals x1 x2
                                      Eq <- equals y1 y2 
                                      return Eq 
equals (If x1 y1 z1) (If x2 y2 z2)  
                                  = do
                                      Eq <- equals x1 x2
                                      Eq <- equals y1 y2 
                                      Eq <- equals z1 z2 
                                      return Eq 
equals _ _ = Nothing

newtype Memo env env'
  =  Memo 
     ( forall x  .   Expr x env 
                 ->  Maybe (Ref x env')
     )

emptyMemo  ::  Memo env ()
emptyMemo  =   Memo (const Nothing)

type TrafoCSE2 env = Trafo (Memo env) Expr

insertIfNew  ::  forall s a env . Expr a env 
                 ->  TrafoCSE2 env  (Expr a ) (Ref a )
insertIfNew e =   
  Trafo
  (\(Memo m :: Memo env env') -> case m e of
     Nothing -> extEnv  (extMemo e (Memo m)) 
     Just r  -> castSRef (Memo m) r 
  )
 
extMemo  :: Expr a env -> Memo env env' 
         -> Memo env (env',a)
extMemo e (Memo m) 
         = Memo  (\s -> case equals e s of
                                           Just Eq -> Just Zero
                                           Nothing -> fmap Suc (m s)
                 ) 

app_cse  ::  Expr a env 
         ->  TrafoCSE2 env  (T env) (Ref a) 

app_cse (Var r)       = arr (\(T tenv_s) -> tenv_s r)
app_cse e@(IntVal i)  = arr (const (IntVal i)) 
                      >>> insertIfNew e    
app_cse e@(BoolVal b) = arr (const (BoolVal b)) 
                      >>> insertIfNew e    
app_cse e@(Add x y)   = (app_cse x &&& app_cse y) 
                      >>> arr (\(P (l,r)) -> Add (Var l) (Var r)) 
                      >>> insertIfNew e

app_cse e@(LessThan x y) 
                       = (app_cse x &&& app_cse y) 
                      >>> arr (\(P (l,r)) -> LessThan (Var l) (Var r))
                      >>> insertIfNew e


app_cse e@(If x y z)  = (app_cse x &&& app_cse y &&& app_cse z) 
                      >>> arr (\(P (P (b,l),r)) -> If (Var b) (Var l) (Var r)) 
                      >>> insertIfNew e

 
newtype EEnv r e1 e2 = EEnv (Env r e2 e1)

cse_env  ::  Env Expr env env'
         ->  TrafoCSE2 env 
                    (T env ) 
                    (EEnv Ref  env')


cse_env Empty      = arr (const (EEnv Empty))  

cse_env (Ext es e) = (app_cse e &&& cse_env es) 
                    >>> arr  (\(P (r, (EEnv  renv)))
                                -> EEnv (Ext renv r)
                             )


refTransformer :: Env Ref s env -> T env s
refTransformer refs = T (\r -> lookupEnv r refs)


data EmptyEnv s = EmptyEnv

trafo :: Decls env -> TrafoCSE2 env EmptyEnv (T env )
trafo decls = loop $ arr (\(P (_,EEnv refs)) -> refTransformer refs) >>> (arr id &&& cse_env decls)  


cse :: forall env . Decls env -> TDecls env
cse decls
   = case  runTrafo (trafo decls) emptyMemo EmptyEnv of
           Result _ t env -> TDecls env t