{-# OPTIONS -fglasgow-exts  #-}
module CSE2 where

  import Language.TTTAS


  data Expr1 a where
       IntVal1    :: Int                               -> Expr1 Int
       BoolVal1   :: Bool                              -> Expr1 Bool
       Add1       :: Expr1 Int -> Expr1 Int            -> Expr1 Int
       LessThan1  :: Expr1 Int -> Expr1 Int            -> Expr1 Bool
       If1        :: Expr1 Bool -> Expr1 a -> Expr1 a  -> Expr1 a  


  data Expr a env where
       Var       ::  Ref a env      -> Expr a env
       IntVal    ::  Int            -> Expr Int env
       BoolVal   ::  Bool           -> Expr Bool 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

  -----------------------------EASY VERSION (Expr1 -> Env)
  
  newtype MapExpr1 env2
    =  MapExpr1 
       ( forall x  .   Expr1 x 
                   ->  Maybe (Ref x env2)
       )
  
  emptyExpr1  ::  MapExpr1 env2 
  emptyExpr1  =   MapExpr1 (const Nothing)

  type TrafoCSE1 inp out = Trafo2 MapExpr1 Expr inp out

  data FinalExprs = forall env. FinalExprs (Env Expr env env)

  cse1 :: Expr1 t -> FinalExprs
  cse1 e
     = let result = runTrafo2 (app_cse1 e) emptyExpr1 undefined
       in case result of 
         Result _ _ envs -> FinalExprs envs


  app_cse1 :: Expr1 a -> TrafoCSE1 t (Ref a) 
  app_cse1 e@(IntVal1 i)  = arr2 (const (IntVal i)) 
                        >>>> insertExpr1 e    
  app_cse1 e@(BoolVal1 b) = arr2 (const (BoolVal b)) 
                        >>>> insertExpr1 e    
  app_cse1 e@(Add1 x y)   = (app_cse1 x &&&& app_cse1 y) 
                        >>>> arr2 (\(P (l,r)) -> Add (Var l) (Var r)) 
                        >>>> insertExpr1 e
  app_cse1 e@(LessThan1 x y) 
                          = (app_cse1 x &&&& app_cse1 y) 
                        >>>> arr2 (\(P (l,r)) -> LessThan (Var l) (Var r))
                        >>>> insertExpr1 e
  app_cse1 e@(If1 x y z)  = (app_cse1 x &&&& app_cse1 y &&&& app_cse1 z) 
                        >>>> arr2 (\(P (P (b,l),r)) -> If (Var b) (Var l) (Var r)) 
                        >>>> insertExpr1 e

  insertExpr1 :: Expr1 a -> TrafoCSE1 (Expr a) (Ref a)
  insertExpr1 e =   
    Trafo2
    (\(MapExpr1 m) -> case m e of
       Nothing  ->  
         case newERef1 e
         of Trafo2 step -> step (MapExpr1 m)
       Just r   -> TrafoE2 (MapExpr1 m) 
                        (\e (T t) env1 ->  (t r, T t, env1))
    )

  newERef1 ::  Expr1 a 
           ->  TrafoCSE1 (Expr a) (Ref a)
  newERef1 e =
    Trafo2 
    (\(MapExpr1 m :: MapExpr1 env1) ->
       let m2 = MapExpr1  (\s -> aux1 (matchExpr1 e s) m s)
       in TrafoE2  m2 
                 (\e (T t) env1 -> (  t Zero
                                   ,  T (t . Suc) 
                                   ,  Ext env1 e
    )            )                 )

  aux1 :: Maybe (Equal a x) -> (Expr1 x -> Maybe (Ref x env1)) 
       -> Expr1 x -> Maybe (Ref x (env1,a))
  aux1 (Just Eq) _ _  = Just Zero
  aux1 Nothing   m s  = fmap Suc (m s)

  matchExpr1  ::  Expr1 a -> Expr1 b 
              ->  Maybe (Equal a b)
  matchExpr1 (IntVal1 i1) (IntVal1 i2) 
                          | i1==i2  = Just Eq
  matchExpr1 (BoolVal1 b1) (BoolVal1 b2) 
                          | b1==b2  = Just Eq
  matchExpr1 (Add1 x1 y1) (Add1 x2 y2) 
                                    = do
                                        Eq <- matchExpr1 x1 x2
                                        Eq <- matchExpr1 y1 y2 
                                        return Eq 
  matchExpr1 (LessThan1 x1 y1) (LessThan1 x2 y2) 
                                    = do
                                        Eq <- matchExpr1 x1 x2
                                        Eq <- matchExpr1 y1 y2 
                                        return Eq 
  matchExpr1 (If1 x1 y1 z1) (If1 x2 y2 z2)  
                                    = do
                                        Eq <- matchExpr1 x1 x2
                                        Eq <- matchExpr1 y1 y2 
                                        Eq <- matchExpr1 z1 z2 
                                        return Eq 
  matchExpr1 _ _ = Nothing

  --- little test1
  n1 x = IntVal1 x
  env1 = Add1 (Add1 (n1 2) (n1 3)) (Add1 (n1 4) (Add1 (n1 2) (n1 3)))

  res1 = show $ cse1 env1


  ----------------------------MORE COMPLICATED VERSION (Env -> Env)



  newtype MapExpr env env2
    =  MapExpr 
       ( forall x  .   Expr x env 
                   ->  Maybe (Ref x env2)
       )

  emptyExpr  ::  MapExpr env env2 
  emptyExpr  =   MapExpr (const Nothing)

  initExpr  :: TrafoCSE env a b 
            ->  Trafo2 Unit Expr a b
  initExpr (Trafo2 st)
    = Trafo2  (\_ ->  case st emptyExpr of
                      TrafoE2 _  f -> TrafoE2 Unit f
            )

  type TrafoCSE env inp out = Trafo2 (MapExpr env) Expr inp out


  newtype Mapping old new 
           = Mapping (Env Ref new old) 

  map2trans :: Mapping env s -> T env s
  map2trans (Mapping env) 
     = T (\r -> (lookupEnv r env))



  cse :: Env Expr env env -> FinalExprs
  cse e
     = let result = runTrafo2  
                    ( loop2 $ second2 $
                      arr2 (\menv_s -> map2trans menv_s) >>>> cse_env e 
                    )
                    Unit       -- meta-data
                    undefined  -- input
       in case result of 
         Result _ _ env -> FinalExprs env


  cse_env  ::  Env Expr env env'
           ->  Trafo2  Unit 
                      Expr
                      (T env) 
                      (Mapping env')

  cse_env (Ext es e) = (initExpr (app_cse e) &&&& cse_env es) 
                    >>>> arr2  (\(P (r, (Mapping  renv)))
                                -> Mapping (Ext renv r)
                             )
  cse_env Empty       = arr2 (const (Mapping Empty))  


  app_cse :: Expr a env -> TrafoCSE env (T env) (Ref a) 
  app_cse (Var r)       = arr2 (\(T tenv_s) -> tenv_s r)
  app_cse e@(IntVal i)  = arr2 (const (IntVal i)) 
                      >>>> insertExpr e    
  app_cse e@(BoolVal b) = arr2 (const (BoolVal b)) 
                      >>>> insertExpr e    
  app_cse e@(Add x y)   = (app_cse x &&&& app_cse y) 
                      >>>> arr2 (\(P (l,r)) -> Add (Var l) (Var r)) 
                      >>>> insertExpr e
  app_cse e@(LessThan x y) 
                        = (app_cse x &&&& app_cse y) 
                      >>>> arr2 (\(P (l,r)) -> LessThan (Var l) (Var r))
                      >>>> insertExpr e
  app_cse e@(If x y z)  = (app_cse x &&&& app_cse y &&&& app_cse z) 
                      >>>> arr2 (\(P (P (b,l),r)) -> If (Var b) (Var l) (Var r)) 
                      >>>> insertExpr e

  insertExpr :: Expr a env -> TrafoCSE env (Expr a) (Ref a)
  insertExpr e =   
    Trafo2
    (\(MapExpr m) -> case m e of
       Nothing  ->  
         case newERef e
         of Trafo2 step -> step (MapExpr m)
       Just r   -> TrafoE2 (MapExpr m) 
                        (\e (T t) env1 ->  (t r, T t, env1))
    )

  newERef ::  Expr a env
          ->  TrafoCSE env (Expr a) (Ref a)
  newERef e =
    Trafo2 
    (\(MapExpr m :: MapExpr env env1) ->
       let m2 = MapExpr (\s -> aux (matchExpr e s) m s)
       in TrafoE2  m2 
                 (\e (T t) env1 -> (  t Zero
                                   ,  T (t . Suc) 
                                   ,  Ext env1 e
    )            )                 )

  aux :: Maybe (Equal a x) -> (Expr x env -> Maybe (Ref x env1)) 
      -> Expr x env -> Maybe (Ref x (env1,a))
  aux (Just Eq) _ _  = Just Zero
  aux Nothing   m s  = fmap Suc (m s)


  matchExpr  ::  Expr a env -> Expr b env 
             ->  Maybe (Equal a b)
  matchExpr (Var r1) (Var r2)       = match r1 r2
  matchExpr (IntVal i1) (IntVal i2) 
                          | i1==i2  = Just Eq
  matchExpr (BoolVal b1) (BoolVal b2) 
                          | b1==b2  = Just Eq
  matchExpr (Add x1 y1) (Add x2 y2) = do
                                        Eq <- matchExpr x1 x2
                                        Eq <- matchExpr y1 y2 
                                        return Eq 
  matchExpr (LessThan x1 y1) (LessThan x2 y2) 
                                    = do
                                        Eq <- matchExpr x1 x2
                                        Eq <- matchExpr y1 y2 
                                        return Eq 
  matchExpr (If x1 y1 z1) (If x2 y2 z2)  
                                    = do
                                        Eq <- matchExpr x1 x2
                                        Eq <- matchExpr y1 y2 
                                        Eq <- matchExpr z1 z2 
                                        return Eq 
  matchExpr _ _ = Nothing


  --- little test2
  vx = Var Zero
  n x = IntVal x
  env2 =  Empty `Ext` (n 2)
                `Ext` (Add (Add vx (n 3)) (Add (n 4) (Add vx (n 3))))
 
  res2 = show $ cse env2


  ------------------------- SHOW

  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) = show 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

  instance Show (Env Expr env1 env2) where
    show (Ext es e) = (show e) ++ "|" ++ show es
    show Empty = ""

  instance Show FinalExprs where
    show (FinalExprs env) = show env