CataCPSTransform.hs

{-# language FlexibleInstances, UndecidableInstances, TypeSynonymInstances #-}

import Prelude hiding (mapM)
import Control.Monad.State hiding (mapM)
import Control.Applicative
import Data.Traversable
import Data.Foldable
import Data.Monoid

data ASTLevel name self
    = Var name
    | App self [self]
    | Let [Binding name self] self
    | Lam (Maybe name) [name] self
    | Literal Literal
    -- | Seq [AST]

type Name
    = String

data Binding name ast
    = Binding name [name] ast
    deriving Show

data Literal
    = Integer Integer
    | Double  Double
    | String  String
    deriving Show

data LCLevel name self
    = Define name [name] self self
    | Call (Term name) [Term name]

data Term name = Name name | Constant String

data Fix f = In { out :: f (Fix f) }

--instance (Functor f, Show (f String)) => Show (Fix f) where
--    show = cata show

instance Show (LC String) where
    show (In (Define x args it rest)) = "let " ++ x ++ " " ++ unwords args ++ " = " ++ show it ++ " in " ++ show rest
    show (In (Call f xs))             = "(" ++ unwords (map show (f : xs)) ++ ")"

instance Show (Term String) where
    show (Name a) = a
    show (Constant x) = x

type AST name = Fix (ASTLevel name)
type LC  name = Fix (LCLevel  name)

instance Functor (LCLevel name) where
    fmap f (Define name args it rest) = Define name args (f it) (f rest)
    fmap _ (Call f xs) = Call f xs

instance Foldable (LCLevel name) where
    foldMap f (Define name args it rest) = f it <> f rest
    foldMap _ (Call f xs) = mempty

instance Functor (ASTLevel name) where
    fmap f (Let bs rest) = Let (fmap (fmap f) bs) (f rest)
    fmap f (App g xs)    = App (f g) (fmap f xs)
    fmap f (Lam n args body) = Lam n args (f body)
    fmap _ (Literal lit) = Literal lit

instance Functor (Binding name) where
    fmap f (Binding x args it) = Binding x args (f it)

instance Foldable (ASTLevel name) where
    foldMap f (Let bs rest) = foldMap (foldMap f) bs <> f rest
    foldMap f (App g xs)    = f g <> foldMap f xs
    foldMap f (Lam n args body) = f body
    foldMap _ _             = mempty

instance Foldable (Binding name) where
    foldMap f (Binding x args it) = f it

instance Traversable (ASTLevel name) where
    traverse f (Let bs rest) = Let <$> traverse (traverse f) bs <*> f rest
    traverse f (App g xs)    = App <$> f g <*> traverse f xs
    traverse f (Lam n args body) = Lam n args <$> f body
    traverse _ (Literal lit) = pure (Literal lit)
    traverse _ (Var name)    = pure (Var name)

instance Traversable (Binding name) where
    traverse f (Binding x args it) = Binding x args <$> f it

cata :: Functor f => (f a -> a) -> Fix f -> a
cata algebra = algebra . fmap (cata algebra) . out

ana :: Functor f => (a -> f a) -> a -> Fix f
ana coalgebra = In . fmap (ana coalgebra) . coalgebra

cataM :: (Traversable f, Monad m) => (f a -> m a) -> Fix f -> m a
cataM algebra = algebra <=< mapM (cataM algebra) . out

--ana :: Functor f => (a -> f a) -> a -> Fix f
--ana coalgebra = In . fmap (ana coalgebra) . coalgebra

term (In (Call _ [])) = True
term  _               = False

getTerm (In (Call t [])) = t

simplification :: ASTLevel String (LC String) -> State Int (LC String)
simplification (Var name) = return $ In (Call (Name name) [])
simplification (App f xs) = do
    f : xs            <- mapM  slice         (f : xs)
    (actions, f : xs) <- foldM glue (id, []) (f : xs)
    return $ actions (In (Call f xs))

simplification (Let bs lcrest) = do
    foldM sliceBinding lcrest bs

simplification (Lam mname args body) = do
    name <- maybe fresh return mname
    return $ In (Define name args body (In (Call (Name name) [])))

sliceBinding rest (Binding name args lcthing) =
    return $ In (Define name args lcthing rest)

----glue :: Eq name => (LC name -> LC name, [Term name]) -> (LC name -> LC name, Term name) -> (LC name -> LC name, [Term name])
glue (f, args) (df, arg) = return (f . df, args ++ [arg])

----slice :: Eq name => LC name -> (LC name -> LC name, Term name)
slice thing | term thing = return (id, getTerm thing)
slice thing = do
    name <- fresh
    return (In . Define name [] thing, Name name)

fresh = do
    int <- get
    modify (+ 1)
    return ("v" ++ show int)

test = (`runState` 0) . cataM simplification $
    In (Let
        [Binding "id" ["x"]
            (In (Var "x"))]

        (In (App
            (In (Let [Binding "bot" [] (In (Var "bot"))] (In (Var "bot"))))
            [In (Var "id")]
            )
        )
    )