Compos with mutually recursive data types

{-# OPTIONS_GHC -Wall -fno-warn-name-shadowing #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE Rank2Types #-}

import Control.Applicative

------------------------------------------------------------------------
newtype Identity a = Identity { runIdentity :: a }

instance Functor Identity where
  fmap f (Identity x) = Identity (f x)

instance Applicative Identity where
  pure x = Identity x
  Identity f <*> Identity x = Identity (f x)
------------------------------------------------------------------------


data Stm
  = SDecl Typ Var
  | SAss  Var Exp
  | SBlock [Stm]
  | SReturn Exp
  deriving (Eq, Show)

data Exp
  = EStm Stm
  | EAdd Exp Exp
  | EVar Var
  | EInt Int
  deriving (Eq, Show)

data Var = V String
  deriving (Eq, Show)

data Typ = T_int | T_float
  deriving (Eq, Show)

data LangAlg f = LangAlg
  { aStm :: Stm -> f Stm
  , aExp :: Exp -> f Exp
  , aVar :: Var -> f Var
  , aTyp :: Typ -> f Typ
  }

type LangFold a = LangAlg (Const a)

class Compos alg t where
  compos :: (forall a. a -> f a)
         -> (forall a b. f (a -> b) -> f a -> f b)
         -> alg f
         -> t
         -> f t

instance Compos LangAlg Stm where compos = composStm
instance Compos LangAlg Exp where compos = composExp
instance Compos LangAlg Var where compos = composVar
instance Compos LangAlg Typ where compos = composTyp

composStm :: (forall a. a -> f a)
          -> (forall a b. f (a -> b) -> f a -> f b)
          -> LangAlg f
          -> Stm -> f Stm
composStm pure ap alg stm = case stm of
  SDecl t v   -> pure SDecl `ap` aTyp alg t `ap` aVar alg v
  SAss v exp  -> pure SAss `ap` aVar alg v `ap` aExp alg exp
  SBlock stms -> pure SBlock `ap` mapMf (aStm alg) stms
  SReturn exp -> pure SReturn `ap` aExp alg exp
 where
   mapMf _ [] = pure []
   mapMf g (x:xs) = pure (:) `ap` g x `ap` mapMf g xs
   {-# INLINE mapMf #-}
{-# INLINE composStm #-}

composExp :: (forall a. a -> f a)
          -> (forall a b. f (a -> b) -> f a -> f b)
          -> LangAlg f
          -> Exp -> f Exp
composExp pure ap alg exp = case exp of
  EStm stm   -> pure EStm `ap` aStm alg stm
  EAdd e1 e2 -> pure EAdd `ap` aExp alg e1 `ap` aExp alg e2
  EVar v     -> pure EVar `ap` aVar alg v
  _ -> pure exp
{-# INLINE composExp #-}

composVar :: (forall a. a -> f a)
          -> (forall a b. f (a -> b) -> f a -> f b)
          -> LangAlg f
          -> Var -> f Var
composVar pure _ap _alg var = pure var
{-# INLINE composVar #-}

composTyp :: (forall a. a -> f a)
          -> (forall a b. f (a -> b) -> f a -> f b)
          -> LangAlg f
          -> Typ -> f Typ
composTyp pure _ap _alg typ = pure typ
{-# INLINE composTyp #-}

------------------------------------------------------------------------

-- - * Convenience Transformations

-- These combinators take a compos-style function for some type @c@ and
-- convert them into a monadic (or applicative) map or fold.
--
-- For example:
-- 
-- @
-- TODO:
-- @
--
-- Unfortunately, these two combinators require @RankNTypes@ (as opposed to
-- @Rank2Types@ which is needed by the @compos@ combinators).  If that is
-- not desired, consider writing manual wrappers.


composOpA :: (Compos alg t, Applicative f) => 
             alg f -> t -> f t
composOpA alg = compos pure (<*>) alg

runMapA :: alg f -> TypeSelector alg c -> c -> f c
runMapA alg sel c = sel alg c

composOp :: Compos alg t => alg Identity -> t -> Identity t
composOp alg = composOpA alg

runMap :: alg Identity -> TypeSelector alg c -> c -> c
runMap alg sel c = runIdentity (sel alg c)

composFold :: Compos alg t =>
              b -> (b -> b -> b)
           -> alg (Const b) -> t -> Const b t
composFold z c alg t =
    compos (\_ -> Const z)
           (\(Const x) (Const y) -> Const (x `c` y))
           alg
           t

runFold :: alg (Const b) -> TypeSelector alg c -> c -> b
runFold alg sel c = getConst (sel alg c)

composOpA_ :: (Applicative f, Compos alg t) =>
              alg (Const (f ())) -> t -> Const (f ()) t
composOpA_ = composFold (pure ()) (*>)

runMapA_ :: alg (Const (f ())) -> TypeSelector alg c -> c -> f ()
runMapA_ alg sel c = getConst (sel alg c)

type TypeSelector alg c = forall f. alg f -> c -> f c

genericLangMapA :: Applicative f => 
                  LangAlg f -> LangAlg f
genericLangMapA self =
    LangAlg { aStm = composOpA self
            , aExp = composOpA self
            , aVar = composOpA self
            , aTyp = composOpA self
            }
{-# INLINE genericLangMapA #-}

genericLangMap :: LangAlg Identity -> LangAlg Identity
genericLangMap self =
    LangAlg { aStm = composOp self
            , aExp = composOp self
            , aVar = composOp self
            , aTyp = composOp self
            }
{-# INLINE genericLangMap #-}

genericLangMapA_ :: Applicative f =>
                    LangAlg (Const (f ())) -> LangAlg (Const (f ()))
genericLangMapA_ self =
    LangAlg { aStm = composOpA_ self
            , aExp = composOpA_ self
            , aVar = composOpA_ self
            , aTyp = composOpA_ self
            }
{-# INLINE genericLangMapA_ #-}

result :: b -> Const b a
result x = Const x

genericLangFold :: b -> (b -> b -> b) -> LangFold b -> LangFold b
genericLangFold z c self =
    LangAlg { aStm = composFold z c self -- foldOp z c composStm self
            , aExp = composFold z c self -- foldOp z c composExp self
            , aVar = composFold z c self -- foldOp z c composVar self
            , aTyp = composFold z c self -- foldOp z c composTyp self
            }
{-# INLINE genericLangFold #-}

------------------------------------------------------------------------

t1 :: Stm
t1 = SBlock [ SDecl T_int (V "x")
            , SAss (V "x") (EAdd (EAdd (EInt 2) (EInt 4))
                                 (EAdd (EVar (V "x")) (EInt 3)))
            , SReturn (EVar (V "x")) ]

type LangSelector c = TypeSelector LangAlg c

literals :: LangSelector c -> c -> [Int]
literals sel = runFold literalsAlg sel
  where
    literalsAlg = (genericLangFold [] (++) literalsAlg) { aExp = go_Exp }

    go_Exp (EInt n) = result [n]
    go_Exp exp = composFold [] (++) literalsAlg exp

rename :: LangSelector c -> c -> c
rename sel c = runMap renameAlg sel c
  where
    renameAlg = genericLangMap renameAlg { aVar = go_Var }

    go_Var (V x) = pure (V $ "_" ++ x)

warnAssign :: LangSelector c -> c -> IO ()
warnAssign sel c = runMapA_ warnAssignAlg sel c
  where
    warnAssignAlg = genericLangMapA_ warnAssignAlg { aStm = go_Stm }

    go_Stm (SAss var _) = result (putStrLn $ "Assignment to: " ++ show var)
    go_Stm s            = composOpA_ warnAssignAlg s

symbols :: LangSelector c -> c -> [(Var, Typ)]
symbols = runFold alg
  where
    alg = genericLangFold [] (++) alg { aStm = go_Stm }

    go_Stm (SDecl typ var) = result [(var, typ)]
    go_Stm s = composFold [] (++) alg s

constFold :: LangSelector c -> c -> c
constFold = runMap alg
  where
    alg = genericLangMap alg { aExp = go_Exp }

    go_Exp (EAdd x y) =
        liftA2 cfold (go_Exp x) (go_Exp y)
    -- TODO: This doesn't work. Why?
    -- go_Exp (EAdd x y) =
    --    liftA2 cfold (aExp alg x) (aExp alg y)
    go_Exp exp = composOpA alg exp

    cfold (EInt n) (EInt m) = EInt (n + m)
    cfold x'       y'       = EAdd x' y'

main :: IO ()
main = 
    do 
       print (literals aStm t1 :: [Int])
       print (rename aStm t1)
       warnAssign aStm t1
       print (symbols aStm t1)
       print (constFold aStm t1)

--main = return ()