AST walk

walk.hs

{-# LANGUAGE
    DeriveFunctor
  , DeriveFoldable
  , DeriveTraversable
  , StandaloneDeriving
  , ViewPatterns
  , PatternSynonyms
  #-}

-- AST independent

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

mapFix :: Functor f => (f (Fix g) -> g (Fix g)) -> Fix f -> Fix g
mapFix f = In . f . fmap (mapFix f) . out

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

-- Expression specific

data ExprF f
  = Num Int
  | Add f f
  | Mul f f
  deriving (Eq, Functor, Foldable, Traversable)

type Expr = Fix ExprF

num :: Int -> Expr
num = In . Num

add, mul :: Expr -> Expr -> Expr
add a b = In (Add a b)
mul a b = In (Mul a b)

pattern Add_ a b = Add (In a) (In b)
pattern Mul_ a b = Mul (In a) (In b)

simplify :: Expr -> Expr
simplify = mapFix $ \ex ->
  case ex of
    Add_ a (Num 0) -> a
    Add_ (Num 0) b -> b
    Mul_ a (Num 1) -> a
    Mul_ (Num 1) b -> b
    _              -> ex

eval :: Expr -> Int
eval = foldFix $ \ex ->
  case ex of
    Num i   -> i
    Add a b -> a + b
    Mul a b -> a * b

pp :: Expr -> String
pp = foldFix $ \ex ->
    case ex of
      Num i   -> show i
      Add a b -> "(" ++ a ++ " + " ++ b ++ ")"
      Mul a b -> "(" ++ a ++ " * " ++ b ++ ")"

main :: IO ()
main =
  do let ast = num 4 `add` (num 8 `mul` num 1) `add` num 0
     putStrLn (pp ast)
     let opt = simplify ast
     putStrLn (pp opt)
     let res = eval opt
     print res