Experiments in Generic Traversal of an AST

ExpTransform.hs

{-# LANGUAGE DeriveFunctor #-}
module Compiler.Transformations.Experiments where

import Control.Monad.State
import Text.Printf

data ExpF a = Var String
            | Lam [String] a
            | App a a deriving (Functor)

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

type Exp = Fix ExpF

var :: String -> Exp
var x = In (Var x)

lam :: [String] -> Exp -> Exp
lam bs e = In (Lam bs e)

app :: Exp -> Exp -> Exp
app e1 e2 = In (App e1 e2)

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

sample :: Exp
sample = lam ["x"] (lam ["y"] (lam ["z"] (app (app (var "+") (app (app (var "+") (var "x")) (var "y"))) (var "z"))))

type Carrier = State Int String

type ConvertFn a = Either (ExpF a) a -> Either (ExpF a) a

spaces :: Int -> String
spaces level =  [1..level] >> " "

printVar :: ConvertFn (State Int String)
printVar (Left (Var x)) = Right (do level <- get
                                    return (printf "\r\n%d%sVar(%s)" level (spaces level) x))

printVar x@_ = x

printLambda :: ConvertFn (State Int String)
printLambda (Left (Lam bs e)) = Right (do level <- get
                                          modify (+1)
                                          e' <- e
                                          modify (\s -> s - 1)
                                          return (printf "\r\n%d%sLam(%s, %s)" level (spaces level) (concat bs) e'))
printLambda x@_ = x

printApp :: ConvertFn (State Int String)
printApp (Left (App e1 e2)) = Right (do level <- get
                                        modify (+1)
                                        e1' <- e1
                                        e2' <- e2
                                        modify (\s -> s - 1)
                                        return (printf "\r\n%d%sApp(%s, %s)" level (spaces level) e1' e2'))
printApp x@_ = x

getRight :: Either a b -> b
getRight (Left _) = error "error"
getRight (Right x) = x

convert :: (Either (ExpF a) a -> Either (ExpF a) a) -> Exp -> a
convert g = cata(getRight . g . Left)

main :: IO()
main = putStrLn $ evalState (convert (printVar . printLambda . printApp) sample) 0