Easy syntax transformations with recursion schemes

Haskell's a great language for writing programming languages. Take your pick of the many tricks that make them easy to implement! Here's one: Factor out recursion from your syntax tree traversals to make them more succinct and composable.

(We're depending on a package called recursion-schemes that defines modules under Data.Functor.Foldable.)

import Data.Functor.Foldable (Base, Recursive (..), Corecursive (..), fold)
import Data.Functor.Foldable.TH (makeBaseFunctor)

Let's say we're representing our language expressions with a tidy little datatype Exp:

-- This language has variables, pairs, and lambdas - just
-- enough to get our point across.
-- Ignore the ! strictness annotations if they are unfamiliar.
data Exp =
    ExpVar !String
  | ExpPair Exp Exp
  | ExpLam !String Exp
  deriving stock (Show)

-- The simple expression `\y -> (x, y)`
-- A lambda with one bound and one free variable in the body
exampleExp :: Exp
exampleExp = ExpLam "y" (ExpPair (ExpVar "x") (ExpVar "y"))

Recall that a lambda introduces a binder (like y in our example). A variable occurrence is bound if an enclosing lambda binds its name (y), and free otherwise (x). To follow the rest of this post, it's enough to know that we we can identify these sets by traversing the tree.

As a first pass, we can identify all the free variables by manually recursing over the tree:

freeVars :: Exp -> [String]
freeVars = \case
  ExpVar v -> [v]
  ExpPair a b -> freeVars a ++ freeVars b
  ExpLam v a -> filter (/= v) (freeVars a)

>>> freeVars exampleExp
["x"]

It works, but we're repeating ourselves too much. It would be too easy to flub one of the recursive calls to freeVars with a type-compatible call to something else!

Factoring out the recursion will fix these issues at the cost of some boilerplate. Though we can define it all manually, let's use TemplateHaskell to derive an auxiliary type and some functions for this. Data.Functor.Foldable.TH gives us makeBaseFunctor to generate everything from our recursive type:

makeBaseFunctor ''Exp

This creates the type ExpF that we call the "base functor" of Exp. Here are the two types:

>>> :info Exp
type Exp :: *
data Exp = ExpVar !String | ExpPair Exp Exp | ExpLam !String Exp
instance Corecursive Exp
instance Recursive Exp
instance Show Exp
type instance Base Exp = ExpF

>>> :info ExpF
type ExpF :: * -> *
data ExpF r = ExpVarF ![Char] | ExpPairF r r | ExpLamF ![Char] r
instance Foldable ExpF
instance Functor ExpF
instance Traversable ExpF

Notice that ExpF is just Exp with holes punched in the self-recursive parts (and as expected with those holes it implements Functor/Foldable/Traversable). The Base type family relates the two data types, and Recursive and Corecursive define the conversions project and embed between them:

>>> :type project @Exp
project @Exp :: Recursive Exp => Exp -> Base Exp Exp
>>> :type embed @Exp
embed @Exp :: Corecursive Exp => Base Exp Exp -> Exp

Recall from above that Base Exp Exp is exactly ExpF Exp, so the pair of functions is an isomorpism between Exp and ExpF Exp. (In case it's unfamiliar, I'm using @Exp to fill in a type argument. It's not required, but it guides clearer type output here.)

We wrote our original freeVars function from a top-down perspective, explicitly calculating the free variables of child expressions. Now let’s think bottom-up:

-- Recall `freeVars :: Exp -> [String]`
-- Here we use the base functor and put the result type
-- into our child type position.
freeVarsF :: ExpF [String] -> [String]
freeVarsF = \case
  ExpVarF v -> [v]
  ExpPairF a b ->
    a ++ b  -- No call here, children already lists
  ExpLamF v a ->
    filter (/= v) a  -- Same here

Now we can use the fold combinator from Data.Functor.Foldable to run our bottom-up function all the way up our expression tree:

>>> :info fold
fold :: Recursive t => (Base t a -> a) -> t -> a

>>> :type fold @Exp freeVarsF
fold @Exp freeVarsF :: Exp -> [String]

>>> fold freeVarsF exampleExp
["x"]

Great, we did the same thing but with a lot of boilerplate! Was the juice worth the squeeze? Consider a second traversal, this one identifying binders:

bindersF :: ExpF [String] -> [String]
bindersF = \case
  ExpVarF _ -> []
  ExpPairF a b -> a ++ b
  ExpLamF v a -> [v] ++ a

Naively, we could fold a tree once to get free variables then again to get binders, but we can do better by fusing the traversals (with a small helper):

fuse :: Functor f =>
  (f x -> x) -> (f y -> y) -> (f (x, y) -> (x, y))
fuse f1 f2 fxy = (f1 (fmap fst fxy), f2 (fmap snd fxy))

We can calculate both in a single bottom-up traversal:

>>> fold (fuse freeVarsF bindersF) exampleExp
(["x"],["y"])

Additional remarks:

1. Consider effectF :: ExpF a -> m a for some Monad m. fold (sequence >=> effectF) :: Exp -> m a captures a bottom-up effectful traversal! effectF :: ExpF (m a) -> m a is also interesting with the use of MonadReader m and local to alter the environment of the child effect.
2. Take a look how the strictness annotations carry over from Exp to ExpF. By omitting them on the recursive positions, we allow folds to lazily demand results from below.
3. It's possible to derive more exotic structures that allow us to fold step by step with explicit stacks rather than working on the call stack.