gistfile1.hs

{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE Rank2Types #-}
{-# LANGUAGE MultiParamTypeClasses #-}

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 Types 

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)

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")) ]

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)

-- * Core Compos Stuff

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
  entry :: alg f -> t -> f t
  ext :: (t -> f t) -> alg f -> alg f

instance Compos LangAlg Stm where
    compos = composStm; entry = aStm; ext f a = a { aStm = f }
instance Compos LangAlg Exp where
    compos = composExp; entry = aExp; ext f a = a { aExp = f }
instance Compos LangAlg Var where
    compos = composVar; entry = aVar; ext f a = a { aVar = f }
instance Compos LangAlg Typ where
    compos = composTyp; entry = aTyp; ext f a = a { aTyp = f }

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 #-}


genericLangAlg :: (forall t alg. Compos alg t => alg f -> t -> f t)
               -> (LangAlg f -> LangAlg f) -> LangAlg f
genericLangAlg op mod = self
  where
    self0 = LangAlg { aStm = op self
                    , aExp = op self
                    , aVar = op self
                    , aTyp = op self
                    }
    self = mod self0
{-# INLINE genericLangAlg #-}

-- * Helpers

{- Requires RankNTypes: 

mkFold :: b -> (b -> b -> b)
       -> (forall f. 
             (forall t alg. Compos alg t => alg f -> t -> f t)
             -> (alg f -> alg f) -> alg f)
       -> (alg (Const b) -> alg (Const b))
       -> alg (Const b)
mkFold z c genAlg mod = genAlg (opFold z c) mod

langFold' z c mod = mkFold z c genericLangAlg mod

-- -}

-- ** Folds

opFold :: Compos alg t => b -> (b -> b -> b) -> alg (Const b) -> t -> Const b t
opFold z c alg t = 
    compos (\_ -> Const z) (\(Const x) (Const y) -> Const (x `c` y)) alg t
{-# INLINE opFold #-}

langFold :: b -> (b -> b -> b)
         -> (LangAlg (Const b) -> LangAlg (Const b)) -> LangAlg (Const b)
langFold z c mod = genericLangAlg (opFold z c) mod
{-# INLINE langFold #-}

composFold :: Compos alg t => 
               b -> (b -> b -> b) -> alg (Const b) -> t -> b
composFold z c alg t = getConst $ opFold z c alg t
{-# INLINE composFold #-}

runFold :: Compos alg c => alg (Const b) -> c -> b
runFold alg c = getConst (entry alg c)
{-# INLINE runFold #-}

addFold :: Compos alg t => (t -> b) -> alg (Const b) -> alg (Const b)
addFold f = ext (Const . f)
{-# INLINE addFold #-}

-- *** Example

literals :: Compos LangAlg c => c -> [Int]
literals = runFold alg
  where
    alg = langFold [] (++) $ addFold goExp

    goExp (EInt n) = [n]
    goExp exp = composFold [] (++) alg exp

-- ** mapA

langMapA :: Applicative f => (LangAlg f -> LangAlg f) -> LangAlg f
langMapA mod     = genericLangAlg opA mod
{-# INLINE langMapA #-}

opA :: (Compos alg t, Applicative f) => alg f -> t -> f t
opA alg t = compos pure (<*>) alg t
{-# INLINE opA #-}

runMapA :: Compos alg c => alg f -> c -> f c
runMapA alg c = entry alg c
{-# INLINE runMapA #-}

composMapA :: (Compos alg t, Applicative f) => alg f -> t -> f t
composMapA alg c = opA alg c
{-# INLINE composMapA #-}

addMapA :: Compos alg t => (t -> f t) -> alg f -> alg f
addMapA f = ext f
{-# INLINE addMapA #-}

logAndRename :: Compos LangAlg c => c -> IO c
logAndRename = runMapA alg
  where
    alg = langMapA $ addMapA goVar

    goVar (V x) = do putStrLn $ "Renaming: " ++ x
                     return (V $ "_" ++ x)

-- ** map

langMap :: (LangAlg Identity -> LangAlg Identity) -> LangAlg Identity
langMap mod      = genericLangAlg op mod
{-# INLINE langMap #-}

op :: Compos alg t => alg Identity -> t -> Identity t
op alg t = opA alg t
{-# INLINE op #-}

runMap :: Compos alg c => alg Identity -> c -> c
runMap alg c = runIdentity $ entry alg c
{-# INLINE runMap #-}

composMap :: Compos alg t => alg Identity -> t -> t
composMap alg c = runIdentity $ op alg c
{-# INLINE composMap #-}

addMap :: Compos alg t => (t -> t) -> alg Identity -> alg Identity
addMap f = ext (Identity . f)
{-# INLINE addMap #-}

-- *** Example

constantFold :: Compos LangAlg c => c -> c
constantFold = runMap alg
  where
    alg = langMap $ addMap goExp

    goExp (EAdd x y) =
        case (constantFold x, constantFold y) of
          -- With enough INLINE pragmas calling 'constantFold' here instead
          -- of goExp is actually not a problem
          (EInt n, EInt m) -> EInt (n + m)
          (x', y') -> EAdd x' y'
    goExp e = composMap alg e


opA_ :: (Compos alg t, Applicative f) => 
        alg (Const (f ())) -> t -> Const (f ()) t
opA_ alg t = opFold (pure ()) (*>) alg t

langMapA_ mod    = genericLangAlg opA_

main = do 
  print (literals t1)
  t1' <- logAndRename t1
  print (literals (constantFold t1'))