TTTAS-0.4.0: examples/CSE.hs

{-# OPTIONS -fglasgow-exts -XArrows #-}

module CSE where

import Language.AbstractSyntax.TTTAS
import Prelude hiding (lookup)
import Control.Arrow

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 TrafoCSE env = Trafo (Memo env) Expr

insertIfNew  ::  forall s a env . Expr a env 
             ->  TrafoCSE env s (Expr a s) (Ref a s)
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 
         ->  TrafoCSE env s (T env s) (Ref a s) 
app_cse (Var r)       = proc (T tenv_s) -> 
                           returnA -< tenv_s r

app_cse e@(IntVal i)  = proc _ -> 
                           insertIfNew e -< IntVal i 

app_cse e@(LessThan x y) 
                      = proc tt -> 
                          do  l <- app_cse x -< tt
                              r <- app_cse y -< tt
                              insertIfNew e -< LessThan (Var l) (Var r)
app_cse e@(BoolVal b) = proc _ -> 
                           insertIfNew e -< BoolVal b   
app_cse e@(Add x y)   = proc tt -> 
                          do  l <- app_cse x -< tt
                              r <- app_cse y -< tt
                              insertIfNew e -< Add (Var l) (Var r)
                             
app_cse e@(If x y z)  = proc tt -> 
                          do  b <- app_cse x -< tt
                              l <- app_cse y -< tt
                              r <- app_cse z -< tt
                              insertIfNew e -< If (Var b) (Var l) (Var r)
 
cse_env  ::  Env Expr env env'
         ->  TrafoCSE env s
                    (T env s) 
                    (Env Ref s env')

cse_env Empty       = proc _ -> returnA -< Empty 
cse_env (Ext es e) = proc tt ->
                       do  renv  <- cse_env es  -< tt
                           r     <- app_cse e   -< tt
                           returnA -< Ext renv r 

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

trafo :: Decls env -> TrafoCSE env s () (T env s)
trafo decls =  proc _ -> 
                      do  rec  let tt = refTransformer refs
                               refs <- cse_env decls -< tt
                          returnA -< tt

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