uniplate / recursion schemes example

expr-stuff.cabal

-- Initial expr-stuff.cabal generated by cabal init.  For further 
-- documentation, see http://haskell.org/cabal/users-guide/

name:                expr-stuff
version:             0.1.0.0
-- synopsis:            
-- description:         
-- license:             
license-file:        LICENSE
author:              Ben Edwards
maintainer:          [email protected]
-- copyright:           
-- category:            
build-type:          Simple
-- extra-source-files:  
cabal-version:       >=1.10

executable expr-stuff
  main-is:             Main.hs
  -- other-modules:       
  -- other-extensions:    
  build-depends:       base >=4.7 && <4.8,
                       containers,
                       recursion-schemes,
                       transformers,
                       uniplate
  default-language:    Haskell2010

Main.hs

{-# LANGUAGE LambdaCase
  , DeriveFunctor
  , DeriveDataTypeable
  , TypeFamilies #-}
module Main where

import Control.Applicative
  ( pure
  , (<$>)
  , (<*>)
  )
import Control.Monad (mplus)
import Control.Monad.Trans.Reader
  ( Reader
  , asks
  , runReader
  )
import Data.Data
import Data.Functor.Foldable
import Data.Generics.Uniplate.Data (rewrite)
import Data.List (partition)
import Data.Map (Map)
import qualified Data.Map as M
import Data.Traversable (sequenceA)

type instance Base Exp = ExpF
type Env = Map String Int

main :: IO ()
main = let e = Add (I 4) (Add (I 5) (Sum [Var "x", I 3, I 4]))
           s = simplify e
           x = cata eval s 
           env = M.singleton "x" 4
        in print s >> (print $ runReader x env)

data Exp =
    I Int
  | Var String
  | Add Exp Exp
  | Sum [Exp]
  deriving (Eq, Show, Data, Typeable)

data ExpF a =
    IF Int
  | VarF String
  | AddF a a
  | SumF [a]
  deriving (Eq, Show, Functor)

instance Foldable Exp where
  project (I x) = IF x
  project (Var x) = VarF x
  project (Add x y) = AddF x y
  project (Sum xs) = SumF xs

eval :: ExpF (Reader Env Int) -> Reader Env Int
eval (IF x) = pure x
eval (VarF x) = asks $ maybe (error "oops") id . M.lookup x
eval (AddF x y) = (+) <$> x <*> y
eval (SumF xs) = foldr1 (+) <$> sequenceA xs

simplify :: Exp -> Exp
simplify = rewrite (\x -> add2sum x `mplus` constFold x)
  where
    add2sum (Add x (Sum xs)) = Just $ Sum (x:xs)
    add2sum (Add (Sum xs) x) = Just $ Sum (x:xs)
    add2sum (Add x y) = Just $ Sum [x,y]
    add2sum _ = Nothing
    constFold (Sum xs) | length is <= 1 = Nothing
                       | otherwise      = Just . Sum $ (g is):ts
      where
        (is, ts) = partition f xs
        f = \case { (I _) -> True; _ -> False }
        g = foldr1 (\(I x) (I y) -> I $ x + y)
    constFold _                    = Nothing