Working with grids in Haskell (or any purely functional language, I’m sure) can be pretty painful, but it doesn’t have to be. The reason it’s so painful is that we’re tempted to manipulate grids using indices. Now maybe if we’re using a nice array library such as vector, then index-based manipulations aren’t so gross, but in this article, I’ll show you a way to manipulate grids that is based on good ol’ lists, and which generalizes to structures beyond grids.

A common problem involving grids is to construct a new one based on an old one by looking in a neighbourhood around each point in the grid. For example, Conway’s Game of Life, which is what’s called a cellular automaton, is such a problem. The idea is that we have a grid, and each cell is either dead or alive. Each iteration constructs a new grid from the current grid by applying the following process to each cell simultaneously:

• A dead cell becomes alive if it has exactly three live neighbours.
• A live cell remains alive if it has two or three live neighbours.
• Any cell dies or remains dead otherwise.

The challenge with this problem is that each cell is not completely independent: each cell needs to know something about its neighbours in order to update. The upshot is that we can’t just map over the grid to do an update. We need some kind of “context-aware” map. The ultimate goal of this article will be to arrive at exactly such an abstraction.

One dimension

Yes, Game of Life is flashy, but let’s start with a 1D cellular automaton called Rule 110. I’ll skip the details of this automaton and instead let you look at this beautiful gif that illustrates how the next generation is constructed.

In short, each cell looks at at itself and its left and right neighbours to decide what state it should be in (on or off) in the next generation.

A 1D grid is just a list, but remember, we have this requirement of some kind of context awareness. So let’s just add a notion of a “focus” to a list. We end up with a structure that I’ll call Z.

{-# LANGUAGE DeriveFunctor, DeriveFoldable, DeriveTraversable #-}

module Z where

import Data.List.NonEmpty ( NonEmpty(..) )
import qualified Data.List.NonEmpty as N

data Z a = Z [a] a [a] deriving (Functor, Foldable, Traversable)

fromList :: NonEmpty a -> Z a
fromList (x :| xs) = Z [] x xs

We can of course extract the focused element from the Z.

extractZ :: Z a -> a
extractZ (Z _ x _) = x

We can also move the focus around. In particular, we can move it to the left, or to the right, provided there’s still stuff remaining in the relevant list.

left :: Z a -> Maybe (Z a)
left (Z [] _ _ ) = Nothing
left (Z (l:ls) x rs) = Just (Z ls l (x:rs))

right :: Z a -> Maybe (Z a)
right (Z _ _ []) = Nothing
right (Z ls x (r:rs)) = Just (Z (x:ls) r rs)

Visually, here’s what moving left and right looks like, using | to isolate the focus.

Initially:   ... a  b |c| d e ...
Moving left: ... a |b| c  d e ...

The resulting structure has the same list inside of it; all that’s changed is our point of view. Now that isn’t particularly mindblowing. But here’s what is: what if we collected all the different possible focuses? In other words, we want a structure that has a version of this Z with a selected, and one with b selected, and one with c selected, and so on. Then, we can map over this structure in order to implement Rule 110.

Let’s call this “collect all the different versions” operation “duplicateZ”. We provide duplicateZ with a particular Z a, which is focusing on some a. This input Z a must appear in the output, since it’s one of the ways we can look at the list that underlies the Z a.

To represent the output of duplicateZ, let’s use the type Z itself! This choice might sound a bit arbitrary at first, but it has some very nice benefits. First, it explains why we called it duplicateZ; have a look at the type now:

duplicateZ :: Z a -> Z (Z a)

It’s the Z that gets duplicated!

To sort out the implementation of duplicateZ, let’s look at some properties we would like it to have:

1. The input Z a ought to be the focused element of the output, meaning that extract (duplicateZ z) gives us back the input z.
2. Moving left and right should commute with duplication, namely duplicateZ <$> left z should equal left (duplicateZ z) (and likewise for right).

To implement duplicateZ satisfying these properties, we need a way to collect all the left and right moves from the input. We’ll need a helper.

iterateMaybe :: (a -> Maybe a) -> a -> [a]
iterateMaybe f x = x : maybe [] (iterateMaybe f) (f x)

duplicateZ :: Z a -> Z (Z a)
duplicateZ z =
  Z (tail $ iterateMaybe left z) z (tail $ iterateMaybe right z)

Remember the property we expected to hold? Look closely at the implementation above to convince yourself that it holds. We can generalize the property a bit more: if we extract from all the positions generated by the duplication, we get back the original structure. Of course, since Z is also a Functor, we can use fmap to extract at every position.

fmap extractZ . duplicateZ = id

This property is one of the comonad laws. Yes, Z is what’s called a comonad! The ability to extract from and duplicate the structure in such a way that these operations ‘cancel out’ when composed is the essence of what it means to be a comonad. There is a further operation we can derive from these two, namely a function to apply a transformation at every possible focused position. We’ll call this operation extend.

class Functor w => Comonad w where
  duplicate :: w a -> w (w a)
  extract :: w a -> a

  extend :: (w a -> b) -> w a -> w b
  extend k = fmap k . duplicate

instance Comonad Z where
  duplicate = duplicateZ
  extract = extractZ

Extend is implemented by duplicating the structure (which collects all the focuses into a new structure) and then using fmap to apply the transformation everywhere. What’s fascinating about extend is that it looks almost like fmap. The difference is that the input to the function parameter gets a w a instead of merely an a. Concretely for Z, what this means is that the passed function can inspect what’s around the focus in order to compute the b. Tying this back to cellular automata, we can use extend to represent the process of simultaneously applying the rule of the automaton to every position in the strip (or grid, as we’ll see in 2D).

Now let’s implement the cellular automaton, Rule 110. We’ll write a function rule110 :: Z Bool -> Bool which, given a focused cell, decides whether that cell should be alive or dead in the next iteration. To handle the boundaries, we will consider that a cell outside the bounds of the strip are dead.

rule110 :: Z Bool -> Bool
rule110 z = case (get left, extract z, get right) of
  (False, False, False) -> False
  (False, False, True) -> True
  (False, True, False) -> True
  (False, True, True) -> True
  (True, False, False) -> False
  (True, False, True) -> True
  (True, True, False) -> True
  (True, True, True) -> False
  where
    get f = maybe False extract (f z)

This function determines whether a particular cell is dead or alive in the next iteration. Using extend, we can upgrade this to a function that computes the whole next strip.

step :: Z Bool -> Z Bool
step = extend rule110

And finally, we can upgrade this stepping function to one that computes, given an initial strip, the infinite list of its evolution.

steps :: Z Bool -> [Z Bool]
steps = iterate step

Overall, comonads give us an elegant language to express contextual computations. The computation of the cellular automaton Rule 110 is contextual: each cell’s future state depends not only on its current state, but also on the current state of both its neighbours.

Two dimensions

Let’s solve a real problem, with an eye towards efficiency. What problem could be more real than this year’s Advent of Code, Day 4? It’s a problem about grids that look like this.

..@@.@@@@.
@@@.@.@.@@
@@@@@.@.@@
@.@@@@..@.
@@.@@@@.@@
.@@@@@@@.@
.@.@.@.@@@
@.@@@.@@@@
.@@@@@@@@.
@.@.@@@.@.

In part 1, we want to count how many @ cells have fewer than 4 other @s in their eight-square neighbourhood.

We could – and I certainly did, in the past – construct a two-dimensional variation on what we did above. However, the resulting structure, being built entirely out of linked lists, is enormously inefficient. Not only that, but it will end up privileging certain movements: if we designed this grid in a row-major way, then moving to an adjacent row is instant, but moving to an adjacent column requires shifting every row. This is unacceptable.

Instead, let’s bust out the trusty vector library, with its \(O(1)\) indexing. This way, rather than represent the current focus directly in the structure of the data as we did with Z, we’ll store an index.

Sounds like I lied, right? Whatever happened to the “grids without indices”? As an unfortunate consequence of desiring efficiency, our Grid module will house a few index calculations, but the actual application logic that solves the problem will not. So I only half-lied.

module Grid where

data G a = G !(Vector (Vector a)) !Int !Int -- row and column indices
    deriving Functor

out :: G a -> Vector (Vector a)
out (G v _ _) = v

width, height :: G a -> Int
width (G v _ _) = V.length (v V.! 0)
height (G v _ _) = V.length v

Next, we need to write Comonad instance for G – this is the interesting part.

• extract :: G a -> a will use the underlying vectors’ indexing to obtain the focused element.
• duplicate :: G a -> G (G a) will construct a grid of grids. This sounds like it will absolutely explode the memory usage of the program, but in fact, the underlying grid we start with will be shared unchanged among all the new G objects we’ll create. These objects will merely differ in which element is focused, i.e. in what indices they store.
• extend :: (G a -> b) -> G a -> G b will function as a fusion of duplicate together with an fmap of the transformation.

instance Comonad G where
  extract (G v i j) = v V.! i V.! j

  duplicate g@(G v i j) = G v' i j where
    v' = V.generate (height g) $ \i ->
      V.generate (width g) $ \j ->
        G v i j

  extend f g@(G v i j) = G v' i j where
    v' = V.generate (height g) $ \i ->
      V.generate (width g) $ \j ->
        f (G v i j)

Although we could define extend f = fmap f . duplicate, I’m unsure of whether this will fuse away the intermediate grid. Better safe than sorry.

Finally, our Grid module will be complete with some functions for moving the focus around, either absolutely or relatively.

seek :: G a -> (Int, Int) -> Maybe (G a)
seek g@(G v _ _) (i, j)
  | 0 <= i && i < height g
  && 0 <= j && j < width g = Just (G v i j)
  | otherwise = Nothing

move :: G a -> (Int, Int) -> Maybe (G a)
move (G v i j) (di, dj) = seek (G v i j) (di + i, dj + j)

-- the 8 relative offsets for the neighbourhood of interest
dir8 = [(-1, -1), (-1, 0), (-1, 1), (0, 1), (1, 1), (1, 0), (1, -1), (0, -1)]

Now we’re ready to tackle solving part 1 of the coding challenge. Remember, we want to count the @s that have fewer than 4 @s in the 8-square neighbourhood around them. We encode this as a predicate G Cell -> Bool that decides whether the focused cell of the grid is an accessible @.

data Cell = Empty | Full
    deriving Enum -- to convert Empty to 0 and Full to 1

-- Is the focused cell an accessible roll of paper?
accessible :: G Cell -> Bool
accessible g
  | Full <- extract g = (< 4) . sum $ maybe 0 (fromEnum . extract) . move g <$> dir8
  | otherwise = False

• move g <$> dir8 moves in all 8 relative directions given by the list dir8. Since this could go out of bounds, each result in wrapped with Maybe.
• maybe 0 (fromEnum . extract) converts a single result into 1 only if the movement was in bounds and then focuses on an @.

With this predicate defined, we can finally express the solution to part 1.

part1 :: Vector (Vector Cell) -> Int
part1 v = go (G v 0 0) where
    go = sum . fmap sum . out . fmap fromEnum . extend accessible

• extend accessible gives us a grid of booleans each saying whether the cell at the corresponding position is an accessible @.
• fmap fromEnum changes the booleans into integers.
• sum . fmap sum . out throws out the focus and adds everything up.

While we’re here, we may as well do part2. It’s not much harder. We’re asked to repeatedly remove all accessible @s from the grid, counting how many we remove in total.

To solve this, let’s express a new contextual computation that builds on accessible. Rather than merely map to a boolean, we’ll map to a tuple consisting of a count (zero or one) and a new value for the cell under focus. I have no good name for this, so let’s use a greek letter.

rho g = if accessible g then (1, Empty) else (0, extract g)

Applying this everywhere in the grid via extend will give us a grid of tuples. We’ll unzip this to get a grid of ones and zeros – we add them up to get the count of removed @s – and a new grid with fewer @s in it.

Then, it suffices to iterate this process to produce an infinite stream of counts of removed @s in each iteration. This stream will eventually reach zero, staying there forever, so we’ll take the nonzero prefix and add that all up.

part2 :: Vector (Vector Cell) -> Int
part2 v = go (G v 0 0) where
    go = sum . takeWhile (> 0) . unfoldr (Just . phi) where
        phi = first (sum . fmap sum . out) . funzip . extend rho

funzip f = (fst <$> f, snd <$> f)

See, I promised there would be no indices in the solution, and I delivered! The resulting solution crucially uses extend to apply a contextually-aware transformation everywhere in the grid. All index manipulation is relegated to the Grid module’s jump function. Overall, I find that this approach to solving problems about grids to be delightfully declarative. Sometimes it’s really just a matter of finding the right abstraction.

P.S.

I originally did most of this development using Haskell’s quasi-dependent types. This was necessary in order to use a store comonad parameterized by a representable functor. Phew. I used this Store comonad because of its touted ability to memoize results. The generality of this technique cut against me, however, as it required me to find a type to serve as an index into the grid. (Int, Int) didn’t cut it, because of the further requirement that there be no “out of bounds” values in the index type. Well this is where the dependent types came in. I decided to use bounded natural numbers as indices. This did actually give a workable solution, but it was all quite disgusting.

It hit me a little later that I didn’t need to use this highly general machinery. I could just define a Comonad instance myself for a 2D vector + coordinates, so that’s the approach that ended up in this article.

First imagine that the OCaml type int is infinite. Now what happens when you refunctionalize an infinite datatype?

Let make a concrete program to work with. This will be a stateful generator that outputs the integers, starting at 0, and ending at an exclusive limit n.

type 'a gen = { next : unit -> 'a option }
let range n =
    let r = ref 0 in {
        next = fun () ->
            let i = !r in
            if i < n then
                Some (r := i+1; i)
            else
                None
    }

This program stores an integer in a reference cell, and branches on it to produce two different behaviours. Our goal is to replace the integer in the cell with a function in the cell that will just do the right thing (i.e. perform the right one of these two behaviours) without needing to branch on the integer to decide what to do.

We will need to store some initial function inside the reference. This function, when called, needs to emit 0 and also update its own reference cell to be the function that will emit 1. But this new function must also do more: when it’s called it also needs to update its own reference to be the function that emits 2, and so on.

So I lied: we’re not going to get rid of the integer, as we would expect from a standard r17n. We just won’t be storing it directly. We’ll be storing it indirectly, somehow.

To see how we’ll store this integer, let’s begin translating the function, leaving blanks for the parts we aren’t sure about yet.

let range n =
    let r = ref (fun () -> ???; 0) in
    { next = fun () -> !r () }

What must go in the question marks is our logic to update the reference cell to be the “next” function in the sequence. Let’s encapsulate this update logic into a function update. Crucially this function will take an integer as input, to identify the next function in the sequence. To identify the next function to store, the update function must decide whether the real sequence, of integers, is finished. In that case, it stores a function that simply outputs None. If the real sequence isn’t finished, then update must store a function that will contain a recursive call to update, but on the next integer.

let range n =
    let rec r = ref (fun () -> update 1; 0)
    and update i =
        if i < n then
            r := fun () -> update (i+1); Some i
        else
            r := fun () -> None
    in
    { next = fun () -> !r () }

There is a slight issue with this implementation: it’s incorrect on an edge case. How about range 0? This generator should start out by returning None, but our implementation would have any generator constructed from range n begin by emitting 0 unconditionally.

The insight we need to patch this up is to recognize that we can make the initial function stored in the reference cell begin by calling update 0. This will check whether 0 is in bounds before proceeding to store the appropriate function in the reference. Therefore, our new initial function, after calling update 0, can then dereference r to obtain the new handler, for the 0 case, and call it.

let range n =
    let rec r = ref (fun () -> update 0; !r ())
    and update i =
        if i < n then
            r := fun () -> update (i+1); Some i
        else
            r := fun () -> None
    in
    { next = fun () -> !r () }

Indeed, the integer is still alive and well, but it isn’t what we store in the reference cell anymore, at least not directly. Instead, we store the closure fun () -> update (i+1); Some i, which has captured the integer i.

Did our transformation in any way improve the program we wrote? No way. But you have to agree, there’s something mysterious and exciting about using a mutable variable to store a recursive function that updates the mutable variable that houses that function.

Getting back to the title of this post now, did we really refunctionalize an integer? I’d actually say we did. R17n would have us write down one function for each of the values in the first-order datatype we’re refunctionalizing. Doing this for a datatype with several billion values, like the integers, is practically infeasible. However, since there were just two behaviours, depending on the relationship between the index i and the limit n, there are really just two different functions we need to store. One of those is fun () -> None, and the other is fun () -> update (i+1); Some i. This latter code snippet, having a free variable i, actually determines a family of functions indexed by i. For a fixed limit n, we implemented the function update that maps each integer to a function that simulates the behaviour in the original program.

Taking a more general stance, when we attempt to refunctionalize a datatype with infinitely many values – or with so many values that we can approximate the count to be infinite – we end up needing to describe a systematic way of translating each of the values into a corresponding function. In other words, we implement a function to translate each value into its refunctionalized form.

Short answer: hard, but not as hard as I thought.

This is part 2 of a series of posts on the implementation of a dependently-typed lambda-calculus. In part 1, we discussed the need for normalization of terms, and used a powerful, extensible approach called Normalization by Evaluation to implement a normalization procedure. In this post, we’ll implement the typechecking procedure for our small calculus, and crucially use the normalization procedure to do.

Bidirectional typechecking

Bidirectional typechecking is a technique that exploits a duality within terms to reduce the amount of typing information required from the programmer. We divide the syntax of terms into normal terms and neutral terms, precisely as we did for values, giving a syntactic characterization of beta-normal terms.

• Normal terms are associated with constructors and are those terms whose typechecking is also type-directed. In other words, we check that a given normal term has a given type. Normal terms are also sometimes called checkable terms.
• Neutral terms are associated with eliminators and are those terms from which their type may be synthesized. In other words, we infer the type of a neutral term in a given context. Neutral terms are therefore sometimes called synthesizable or inferrable terms.

For example, in order to appropriately extend the context during typechecking, we classify a lambda abstraction as a normal term. Hence, its type is given, making the context extension straightforward to compute. In picking apart the given type of the lambda abstraction, we obtain the expected type of the abstraction’s body, so we equally classify the abstraction body as normal.

Dually, we expect a function application to be neutral. To infer the type of an application, we first infer the type of its subject – requiring that the function be neutral assures beta-normality – and insist that the type be a Pi-type (x:A) -> B. This in turn reveals the expected type A of the function’s argument, so we let the argument be normal.

Here’s a BNF grammar.

Normal terms  t, A, B  ::= λx.t | () | (x:A) -> B | ⊤ | ★ | s
Neutral terms       s  ::= x | s t

This exposition shows intuitively that typechecking beta-normal terms is in some way easier. We begin by typechecking a normal term against its given type. The term, being a stack of introduction forms, is structurally aligned with its type. In doing so, we can easily extend the context at binding sites. Typechecking switches modes into type inference, when the normal term switches to a neutral term. Information stored in the context comes into play upon encountering a variable, and furthermore when inferring the type of the subject of a function application.

We express this bimodal typechecking scheme formally using a pair of mutually defined judgments.

• G |- t <= A normal term t checks against given type A in context G.
• G |- s => A neutral term s synthesizes type A in context G.

These judgments are defined inductively by the following rules. Let’s start with some straightforward ones.

    G, x:A |- t <= B              G(x) = A
---------------------------    ------------
  G |- λx.t <= (x:A) -> B        G |- x => A


  G |- t => (x:A) -> B    G |- t' <= A
--------------------------------------    --------------
         G |- t t' => [t'/x]B              G |- () <= ⊤

  G |- A <= ★    G, x:A |- B <= ★
---------------------------------     --------------
       G |- (x:A) -> B <= ★              G |- ⊤ <= ★


--------------
  G |- ★ <= ★

Yes, at the bottom there is the type-in-type rule, meaning that the resulting system is unsound. I’ll explore later how we might introduce universes to address that.

 G |- s => B    G |- A ≡ B : ★
-------------------------------
         G |- s <= A

This rule captures mode-switching, from checking to inference. When we need to check a synthesizable term against a given type, we synthesize the type of that term before checking that the expected and actual type are equal as types G |- A ≡ B : ★. In practice, we’ll already have A and B be in normal form, so the equality check is merely syntactic!

Now it suffices to encode these rules into a pair of functions in OCaml. The hard case will be function application, as there’s a substitution there and we don’t yet know how to compute those.

let rec check (cG : ctx) (t : tm) (tA : tp) : unit =
    match tA, t with
    | Star, Top -> true
    | Star, Star -> true
    | Star, Pi ((x, tA), tB) ->
        check cG tA Star;
        check ((x, tA)::cG) tB Star
    | Top, Unit -> true
    | Pi ((_x, tA), tB), Lam (x, t) ->
        check ((x, tA)::cG) t tB
    | tA, s ->
        let tA' = synth cG s in
        (* and now we need to check that tA' = tA *)

and synth (cG : ctx) (t : tm) : tp =
    match t with
    | Var i -> List.nth cG i |> snd
    | App (s, t) ->
        begin match synth cG s with
        | Pi ((x, tA), tB) ->
            check cG t tA;
            (* and now we need to compute and return [t/x]tB *)
        end

In the code above, there are two TODOs to resolve. The first one is easy. We’ll insist that before the initial call to check, we have already normalized the given type tA. Then, we’ll need to insist that synth also return a type in normal form. Hence we resolve the first TODO by performing a syntactic equality check on tA and tA'. If we had used a completely nameless representation, it would suffice to use OCaml’s built-in equality tA' = tA, but on account of our compromise of keeping some names, we’ll actually need to implement a tm_eq function that ignores names. I’ll omit the code of that function.

The second TODO is more sensitive. Even if we assume that t and tB are in normal form, the result of the substitution [t/x]tB might not be.

Substitution

As a first approach, we might implement a function subst : (tm * ix) -> tm -> tm such that subst (t, x) tB computes the substitution [t/x]tB, then use the NbE procedure from part 1 of the series to normalize.

This first approach sucks for two reasons. First, implementing substitution on nameless terms requires implementing various shifting (weakening) helpers to account for those situations where the substitution crosses a binder such as a lambda abstraction. Second, we end up effectively traversing the term tB three times: once for the substitution, then twice during normalization (once for eval, and once for quote).

Funny enough, we already implemented a clever substitution procedure that addresses both those issues! Remember eval? It replaces all free variables in a term with the values held in an environment. The values in the environment crucially represent variables using de Bruijn levels, not indices, which gives us weakening for free. No need for any shifting.

It suffices to construct an appropriate environment for the substitution we wish to perform, use eval to perform it, and use quote to convert back to syntax. We therefore avoid the tricky business of implementing nameless substitutions entirely, and cut down on the number of term traversals required.

To build the environment we need, consider that t lives in context cG and tB lives in context (x, tA)::cG – those contexts tell us how many free variables appear in these terms. We need to convert the context cG into a ‘dummy environment’ eG that maps each variable back to itself. Then, we can evaluate t in the environment eG to get a value v to form the environment (x, v)::eG to finally evaluate tB. This will replace the variable x with v, as required!

To convert a context into such a ‘dummy environment’, let’s implement ctx2env.

let vvar l = VN (NVar l)

let rec ctx2env (cG : ctx) : env * int = match cG with
    | [] -> ([], 0)
    | (x, _)::cG ->
        let eG, n = ctx2env cG in
        ((x, vvar n)::eG, n+1)

In order to ‘count backwards’ to compute the correct levels, ctx2env also computes as a by-product the length of the context.

Next, let’s fill in the second TODO of the typechecker, following the gameplan I outlined above.

let rec check (cG : ctx) (t : tm) (tA : tp) : unit = ...
and synth (cG : ctx) (t : tm) : tp =
    match t with
    | Var i -> List.nth cG i
    | App (s, t) ->
        begin match synth cG s with
        | Pi ((x, tA), tB) ->
            check cG t tA;
            let eG, n = ctx2env cG in
            let v = eval eG t in
            let vB = eval ((x, v)::eG) tB in
            (* TODO: `quote vB` *)
        end

Quoting vB is a bit tricky. Recall from part 1 that quoting is type-directed, so we need to supply the type at which we’re quoting, and a typing environment that maps each free variable to its semantic type. We needed this typing information in order to perform correct eta-expansion.

The type at which we’re quoting is easy. Since we’re quoting a type, we’re quoting at type VStar. The typing environment, on the other hand, is harder to come by. We need to evaluate each of the types in cG. But then each entry in cG lives in the subcontext of remaining entries. Yuck. Let’s implement a helper ctx2tyenv that can at least leverage the eG we already computed.

let rec ctx2tyenv (cG : ctx) (eG : env) : tp_env =
    match cG, eG with
    | [], [] -> []
    | (x, tA)::cG, _::eG ->
        (* cG |- tA : ★
           and eG is the dummy env of cG *)
        let vA = eval eG tA in
        (x, vA) :: ctx2tyenv cG eG

Finally we can quote vB.

let rec check (cG : ctx) (t : tm) (tA : tp) : unit = ...
and synth (cG : ctx) (t : tm) : tp =
    match t with
    | Var i -> List.nth cG i
    | App (s, t) ->
        begin match synth cG s with
        | Pi ((x, tA), tB) ->
            check cG t tA;
            let eG, n = ctx2env cG in
            let v = eval eG t in
            let vB = eval ((x, v)::eG) tB in
            quote n (ctx2tyenv cG eG) VStar vB
        | _ -> failwith "ill-typed: application subject not a function"
        end

Unfortunately, this approach merely trades one inefficiency for another.

Stop and consider for a moment that applications are often nested. This means what the subterm s is likely to be yet another application. We end up repeatedly calling ctx2env and ctx2tyenv on the same context cG, meaning we repeatedly evaluate the types in that context.

Moreover, in the event of nested applications, notice that we quote vB only to get back another Pi-type on whose subterms we would then redundantly eval again!

Leaning on the semantics

Rather than repeatedly convert back and forth from syntax into semantics, we could instead work more closely with the semantics. For instance, recall that we imposed a precondition on check, that the given type would be in normal form. In that case, let’s not take the type as a syntactic tp, but rather as a semantic vtp, which we designed to capture only beta-normal forms anyway. Likewise, the context ought to only store types in normal form, too, so why not just use a tp_env instead of a ctx? And finally, if synth should output a type in normal form, again let’s output a vtp instead.

Crucially, the only one that remains as a raw, syntactic term is the term under consideration for checking or synthesis. What’s more, that’s the only term we can’t evaluate first, as we don’t yet know whether it’s well-typed!

Besides representing everything that ought to be in normal form as a value, we’ll also want to avoiding using ctx2env to generate the ‘dummy’ environment we need when calling eval during typechecking. To do so, the typechecker will also track a dummy environment. Whenever it goes under a binder, we will extend not only the typing environment but also the dummy environment. To extend the dummy environment, we need to generate a new variable whose level is the current length of the dummy environment. To avoid computing this length, we’ll also track it as we go.

let rec check (d : lvl) (e : env) (eG : tp_env) (t : tm) (vA : vtp) : unit =
    match vA, t with
    | VStar, Top -> ()
    | VStar, Star -> ()
    | VStar, Pi ((x, tA), tB) ->
        check d e eG tA VStar;
        let vA = eval e tA in
        check (d+1) ((x, vvar d)::e) ((x, vA)::eG) tB VStar
    | VTop, Unit -> ()
    | VPi ((_, vA), fB), Lam (x, t) ->
        check (d+1) ((x, vvar d)::e) ((x, vA)::eG) t (fB (vvar d))
    | vA, s ->
        let vA' = synth d eG s in
        let tA = quote d eG VStar vA in
        let tA' = quote d eG VStar vA' in
        if not (tm_eq tA tA') then failwith "type mismatch"

and synth (d : lvl) (e : env) (eG : tp_env) (t : tm) : vtp =
    match t with
    | Var i -> List.nth eG i |> snd
    | App (s, t) ->
        begin match synth d e eG s with
        | VPi ((x, vA), fB) ->
            check d e eG t vA;
            let v = eval e t in
            fB v
        | _ -> failwith "ill-typed: application subject not a function"
        end
    | _ -> failwith "cannot synthesize type of checkable term"

There are two things that remain a bit unpleasant about this implementation. In the last case of check, we need to check that vA equals vA', and we do so by quoting both and comparing as terms. It’s a bit wasteful to fully quote both values if it turns out they aren’t equal though. When they are equal, we still end up performing three traversals: once for each quote, and once more for tm_eq.

We can do better by implementing an equality procedure directly on values. Essentially, the procedure does the job of quoting (and hence eta-expanding) both values as long as they match, stopping early as soon as it detects that they don’t.

Since we need to maintain typing information along the way (to support eta-expansion), the implementation will follow the normal/neutral split into a pair of mutually recursion functions val_eq and neu_eq, with neu_eq synthesizing a type.

let rec val_eq (d : lvl) (eG : tp_env) (v1 : value) (v2 : value) (vA : vtp) : bool =
    match vA, v1, v2 with
    | VStar, VTop, VTop -> true
    | VStar, VStar, VStar -> true
    | VTop, VUnit, VUnit -> true
    | VStar, VPi ((x1, vA1), fB1), VPi ((x2, vA2), fB2) ->
        let v = vvar d in
        val_eq d eG vA1 vA2 VStar && val_eq (d+1) ((x1, vA1)::eG) (fB1 v) (fB2 v) VStar
    | VPi ((x, vA), fB), v1, v2 -> (* handles eta-expansion *)
        let v = vvar d in
        val_eq (d+1) ((x, vA)::eG) (apply v1 v) (apply v2 v) (fB v)
        (* using `apply` to either reduce an application of a lambda or to extend the stack of
           neutral terms *)
    | vA, VN n1, VN n2 ->
        begin match neu_eq d eG n1 n2 with
        | Some _ -> true
        | None -> false
        end
    | _ -> false

(* checking equality of neutral terms needs to compute a type as output, so we use `vtp option` as
an output type to represent "yes they're equal and here's their type" or "no they're not equal". *)
and neu_eq (d : lvl) (eG : tp_env) (n1 : neu) (n2 : neu) : vtp option
    match n1, n2 with
    | NVar l1, NVar l2 when l1 = l2 = Some (List.nth (lvl2ix d l1) eG)
    | NApp (n1, v1), NApp (n2, v2) ->
        begin match neu_eq d eG n1 n2 with
        | Some (VPi ((x, vA), fB)) ->
            if val_eq d eG v1 v2 vA then
                Some (fB v1)
            else
                None
        | Some _ -> failwith "impossible: inputs of neu_eq are ill-typed"
        | None -> None
        end
    | _ -> None

Now we can rewrite check to use this procedure.

let rec check (d : lvl) (eG : tp_env) (t : tm) (vA : vtp) : unit =
    match vA, t with
    (* ... *)
    | vA, s ->
        let vA' = synth d eG s in
        if not (val_eq d eG vA vA' VStar) then failwith "type mismatch"

Much better.

Then there’s just one thing left to reflect on in the implementation.

Checking and evaluating at once?

In our implementation of check and synth, there are two places where we typecheck a subterm right before we evaluate it, both involving Pi-types. This begs the question, could we not fuse the typechecking and evaluation procedures together, to at least handle those cases?

let (let*) x f = Option.bind
let rec eval_check (d : lvl) (eG : tp_env) (e : env) (t : tm) (vA : vtp) : value option =
    match vA, t with
    | VStar, Top -> Some VTop
    | VStar, Star -> Some VStar
    | VStar, Pi ((x, tA), tB) ->
        let* vA = eval_check d eG e tA VStar in
        let* _ = eval_check (d+1) ((x, vA)::eG) ((x, vvar d)::e) tB VStar in
        Some (VPi ((x, vA), fun v -> eval (v::e) tB))
    | VTop, Unit -> VUnit
    | VPi ((_, vA), fB), Lam (x, t) ->
        let* _ = eval_check (d+1) ((x, vA)::eG) ((x, vvar d)::e) t (fB (vvar d)) in
        Some (VLam (x, fun v -> eval (v::e) t))
    | vA, s ->
        let* v, vA' = eval_synth d eG e s in
        if val_eq d eG vA vA' VStar then
            Some v
        else
            None

and eval_synth (d : lvl) (eG : tp_env) (e : env) (t : tm) : (value * vtp) option =
    match t with
    | Var i -> Some (List.nth e i, List.nth eG i)
    | App (s, t) ->
        begin match eval_synth d eG e s with
        | Some (v1, VPi ((x, vA), fB)) ->
            let* v2 = eval_check d eG e t vA in
            Some (apply v1 v2, fB v2)
        | _ -> None (* application subject not a function *)
        end
    | _ -> None (* trying to synthesize from a checkable term *)

and apply v1 v2 = match v1 with
    | VN n -> VN (NApp (n, v2))
    | VLam (_, f) -> f v2
    | _ -> failwith "runtime type error: application subject not a function"

This is not an improvement. The point of typechecking is that it’s static – we don’t want to run the program when we typecheck it. Running the program typically takes significantly longer than merely typechecking it. In the presence of dependent types, we unfortunately have to run some parts of the program, to do normalization, but that doesn’t mean we should go farther and insist on running the entire program.

Conclusion

Of the surveyed implementations, the strategy that leans on the semantics avoids the back-and-forth between syntax and semantics, while also limiting evaluation to only those subterms required as dependencies in types.

Let’s recap that implementation in full, here.

type ix = int

type tm =
    (* terms *)
    | Lam of name * tm
    | Var of ix
    | App of tm * tm
    | Unit
    (* types *)
    | Pi of (name * tp) * tp
    | Top
    | Star
and tp = tm (* to understand a term as a type *)
and ctx = (name * tp) list

type value =
    | VLam of name * (value -> value)
    | VUnit
    | VPi of (name * vtp) * (value -> vtp)
    | VTop
    | VStar
    | VN of neu (* neutral terms embed as values *)

and neu =
    | NVar of lvl
    | NApp of neu * value

and lvl = int

type env = value list

let rec eval (e : env) (t : tm) : value =
    match t with
    (* types: *)
    | Top -> VTop
    | Star -> VStar
    | Pi ((x, tA), tB) -> VPi (eval e tA, (x, fun v -> eval (v::e) tB))
    (* terms: *)
    | Unit -> VUnit
    | Lam (x, t) -> VLam (x, fun v -> eval (v::e) t)
    | Var i -> List.nth e i
    | App (t1, t2) -> apply (eval e t1) (eval e t2)

and apply v1 v2 = match v1 with
    | VN n -> VN (NApp (n, v2))
    | VLam (_, f) -> f v2
    | _ -> failwith "type error"

type tp_env = env

let vvar l = VN (NVar l)

let rec val_eq (d : lvl) (eG : tp_env) (v1 : value) (v2 : value) (vA : vtp) : bool =
    match vA, v1, v2 with
    | VStar, VTop, VTop -> true
    | VStar, VStar, VStar -> true
    | VTop, VUnit, VUnit -> true
    | VStar, VPi ((x1, vA1), fB1), VPi ((x2, vA2), fB2) ->
        let v = vvar d in
        val_eq d eG vA1 vA2 VStar && val_eq (d+1) ((x1, vA1)::eG) (fB1 v) (fB2 v) VStar
    | VPi ((x, vA), fB), v1, v2 -> (* handles eta-expansion *)
        let v = vvar d in
        val_eq (d+1) ((x, vA)::eG) (apply v1 v) (apply v2 v) (fB v)
        (* using `apply` to either reduce an application of a lambda or to extend the stack of
           neutral terms *)
    | vA, VN n1, VN n2 ->
        begin match neu_eq d eG n1 n2 with
        | Some _ -> true
        | None -> false
        end
    | _ -> false

and neu_eq (d : lvl) (eG : tp_env) (n1 : neu) (n2 : neu) : vtp option
    match n1, n2 with
    | NVar l1, NVar l2 when l1 = l2 = Some (List.nth (lvl2ix d l1) eG)
    | NApp (n1, v1), NApp (n2, v2) ->
        begin match neu_eq d eG n1 n2 with
        | Some (VPi ((x, vA), fB)) ->
            if val_eq d eG v1 v2 vA then
                Some (fB v1)
            else
                None
        | Some _ -> failwith "impossible: inputs of neu_eq are ill-typed"
        | None -> None
        end
    | _ -> None

let rec check (d : lvl) (e : env) (eG : tp_env) (t : tm) (vA : vtp) : unit =
    match vA, t with
    | VStar, Top -> ()
    | VStar, Star -> ()
    | VStar, Pi ((x, tA), tB) ->
        check d e eG tA VStar;
        let vA = eval e tA in
        check (d+1) ((x, vvar d)::e) ((x, vA)::eG) tB VStar
    | VTop, Unit -> ()
    | VPi ((_, vA), fB), Lam (x, t) ->
        check (d+1) ((x, vvar d)::e) ((x, vA)::eG) t (fB (vvar d))
    | vA, s ->
        let vA' = synth d e eG s in
        if not (val_eq d eG vA vA' VStar) then failwith "type mismatch"

and synth (d : lvl) (e : env) (eG : tp_env) (t : tm) : vtp =
    match t with
    | Var i -> List.nth eG i |> snd
    | App (s, t) ->
        begin match synth d eG s with
        | VPi ((x, vA), fB) ->
            check d e eG t vA;
            let v = eval e t in
            fB v
        | _ -> failwith "ill-typed: application subject not a function"
        end
    | _ -> failwith "cannot synthesize type of checkable term"

To actually use this implementation, consider that a user will supply a term together with its type, both represented as syntax and assumed to be closed.

let check_user (t : tm) (tA : tp) : unit =
    check 0 [] [] tA VStar; (* check that the type is a valid type *)
    let vA = eval [] tA in (* evaluate it, to then guide checking of the user's program *)
    check 0 [] [] t vA (* check the user's program *)

Depending on the situation, after evaluation, we might simply evaluate the user’s program, or compile it.

This gives, in 115 lines of code, a fairly straightforward implementation of a dependently-typed lambda calculus. This implementation omits quote. After introducing equality directly on values, it became unnecessary to convert back into syntax to complete the implementation of the type checker. In a practical implementation, with real error messages, we would need it: in the case that type checking switches to synthesis, for instance, we check equality of values, and would like to tell the user what the expected and actual types are!

In the next post in this series, I’ll extend the calculus with some built-in inductive types – natural numbers and equality – and encode a proof by induction that addition is associative.

Short answer: hard, but not as hard as I thought.

Type systems in which the types (of functions, typically) may depend on terms are called dependently-typed. These systems exist on a spectrum, where restricted forms of dependent types exist in common functional programming languages such as OCaml, Haskell, and even TypeScript; the most complex forms of dependent types appear in proof assistants like Agda and Coq.

Dependent types complicate the typechecking process. Normally, a typechecker verifies, for instance, that the type of the argument in a function application is compatible with the type of the parameter of the function. In a simply-typed language, it suffices to check that the types are equal. In the dependently-typed setting, types may contain terms requiring evaluation, so the typechecker must normalize types before proceeding to check equality. For example, the typechecker must decide whether 3 + 5 (a term requiring evaluation) equals 8 in the course of deciding whether a term of type Vec (3 + 5) A can be passed to a function expecting an input of type Vec 8 A.

Let me be clear: proving that every well-typed term in a dependently-typed language can be normalized – this is called the normalization property – is very challenging. But that is not the goal of this post.

In this series of posts, I want to demonstrate how we can implement a typechecker for a small, dependently-typed lambda-calculus by relying on an elegant normalization technique called Normalization by Evaluation. normalization procedure. The plan for the series is the following.

• Part 1: core syntax, evaluation, normal vs neutral terms, normalization.
• Part 2: bidirectional typechecking, substitutions, semantic equality.
• Part 3: scaling up the language: induction on natural numbers, identity type.

The core calculus

If we allow terms to compute types and types to refer to terms, then the syntactic separation between these objects loses its relevance. Instead, let’s define a unified syntax that captures both, but nonetheless use the letter t to refer to terms understood as terms but the letters A and B to refer to terms understood as types.

Terms t, A, B ::= x | λx. t | t1 t2 | () | (x:A) -> B | ⊤ | ★

This syntax contains variables, lambda abstractions, applications, a constant () (unit), the dependent function type (x:A) -> B, the unit type ⊤, and the type of types ★. (To simplify, we’ll use the rule ★ : ★.)

With this syntax, we can write things we couldn’t in the simply-typed lambda calculus. For example, consider the type (x:⊤) -> (λx.⊤) x. From the point of view of the simply-typed lambda calculus, this type is very unusual: it contains a reducible expression (redex), namely the right-hand side of the arrow can be simplified to just ⊤ by beta-reduction.

This syntax also allows for quantification over types, allowing us to write polymorphic functions, too. For example, here’s the type of a (simply-typed) composition operator.

(A:★) -> (B:★) -> (C:★) -> (f : B -> C) -> (g : A -> B) -> (x:A) -> C

We encode the syntax of the core calculus as a recursive type.

type ix = int

type tm =
    (* terms *)
    | Lam of name * tm
    | Var of ix
    | App of tm * tm
    | Unit
    (* types *)
    | Pi of (name * tp) * tp
    | Top
    | Star
and tp = tm (* to understand a term as a type *)
and ctx = (name * tp) list

This syntax uses a nameless (de Bruijn) representation for variables in order to simplify the implementation of capture-avoiding substitution, but I do store names of variables in binders and contexts as a compromise to enable writing a pretty-printer that shows names.

Normalization by Evaluation

A term is in normal form when it cannot be further reduced. Suppose t and A are in normal form, and that A contains a free variable x. Will the substitution [t/x]A be in normal form?

Not necessarily! That’s the central challenge in implementing our typechecker, which we will surmount by implementing a normalization procedure.

Normalization by Evaluation is a semantic approach to normalization. It is made up of two related procedures.

1. eval is a bog-standard evaluation procedure, which interprets a lambda-term into a semantic space of values. This procedure is responsible for finding the beta-normal form of a lambda-term by eliminating all redexes from the term.
2. quote is its inverse, which reifies an element from the semantics (a value) back into a lambda-term. This procedure is responsible for eta-expanding terms at function types, enabling us to consider f and λx. f x as judgmentally equal at the type A -> B.

The game plan to implement this pair of procedures is first to define the semantic space, i.e. the set values; then to implement eval; and finally to implement quote.

Values

The value of a term is its semantics. Since we’re working in OCaml, we can give the semantics of a lambda-term at function type as an OCaml function. In other words, we design our evaluation procedure so that if t : A -> B, then eval(t) gives us an OCaml function from values of type A to values of type B. This is a good idea, operationally speaking, because we get to wait until the argument to the function is known (and evaluated) before proceeding to evaluate the function body. The same line of reasoning leads to interpreting a Pi-type as an OCaml function on types.

type value =
    | VLam of name * (value -> value)
    | VUnit
    | VPi of (name * vtp) * (value -> vtp)
    | VTop
    | VStar

and vtp = value

This syntax of values is good enough for closed terms, having no free variables, but it is insufficient for open terms. Concretely, we will run into a problem when trying to implement quote later on: how will we reconstruct a lambda-term from an OCaml function value -> value? We want to quote the body of the abstraction – that is the output of the function value -> value – but to access the body, we must apply this OCaml function to some value. What value can we use?

The resolution to this quandry lies in generalizing the syntax of values to explicitly account for open terms. We might think to add merely a constructor for variables, but this isn’t enough, as e.g. x (λx. y) is an application that’s in normal form.

Rather than merely add variables, we add a separate syntax of so-called neutral terms. These are variables and elimination forms applied to neutral terms. Conceptually, a neutral term is a stack of elimination forms that are ultimately blocked on a variable. In contrast, the syntax of values I gave above corresponds to the introduction forms of the lambda-calculus.

type value =
    | VLam of name * (value -> value)
    | VUnit
    | VPi of (name * vtp) * (value -> vtp)
    | VTop
    | VStar
    | VN of neu (* neutral terms embed as values *)

and neu =
    | NVar of lvl
    | NApp of neu * value

and vtp = value
and lvl = int

Notice that the representation I use for variables here mentions lvl – these are de Bruijn levels, the dual of de Bruijn indices. When variables use de Bruijn levels, weakening such terms becomes free, whereas the use of de Bruijn indices would require costly explicit shifts. This will end up simplifying the implementation of eval.

Finally, this combined syntax of values and neutral terms is specifically designed to represent only beta-normal forms. A value is a stack of introduction forms that might at some point “switch” to a neutral term, which is then a stack of elimination forms ending on a variable. Crucially impossible is to write an elimination form whose subject is an introduction form, i.e. a non-beta-normal term.

eval

Recall from kindergarten how to evaluate lambda-terms, and use your imagination to extend the procedure to types. We’ll use an environment-based approach to efficiently handle substitutions in a lazy fashion. In the presence of neutral terms, whenever we evaluate an elimination form, we need to explicitly check whether the subject of the elimination is normal or neutral to proceed accordingly.

• Elimination subject is normal. A reduction is therefore possible, so we must perform it.
• Elimination subject is neutral. Reduction is ultimately blocked on a variable, so we extend the stack of blocked eliminations.

The helper function apply, which interprets a function application, performs this check.

type env = value list

let rec eval (e : env) (t : tm) : value =
    match t with
    (* types: *)
    | Top -> VTop
    | Star -> VStar
    | Pi ((x, tA), tB) -> VPi ((x, eval e tA), fun v -> eval (v::e) tB)
    (* terms: *)
    | Unit -> VUnit
    | Lam (x, t) -> VLam (x, fun v -> eval (v::e) t)
    | Var i -> List.nth e i
    | App (t1, t2) -> apply (eval e t1) (eval e t2)

and apply v1 v2 = match v1 with
    | VN n -> VN (NApp (n, v2))
    | VLam (_, f) -> f v2
    | _ -> failwith "type error"

quote

Equipped with eval to perform reductions, interpreting a lambda-term into an OCaml semantics, we’re now ready to implement its dual, to finally arrive at a normalization procedure by composing the two. This dual procedure, called quote, transforms a semantic value back into syntax.

What’s conceptually tricky about quote is how we handle functions. Recall that Lam (x, t) evaluates to the (metalanguage) closure fun v -> eval (v::e) t where the (meta) free variable e is an environment providing values for the (object) free variables present in t. To quote this, we must first apply the closure to a value. Thankfully, we have neutral terms at our disposal: we generate a variable, corresponding to the bound variable of the abstraction, to use as an argument.

Since neutral variables use de Bruijn levels, the quote procedure tracks a current depth, incremented by one whenever it traverses a binder.

Furthermore, quote is responsible for finding an eta-normal form, by performing eta-expansion of neutral terms at function types. This requires typing information, which quote will ultimately receive from the typechecker in the form of a semantic type. Then, quote must perform type reconstruction to maintain typing information at every step: quote traverses the given value together with its (semantic) type.

To know the type of a variable when quoting encounters it, we equip the process with a map from variables to semantic types. This mapping is called a typing environment, in contrast with an (ordinary) context that maps variables to syntactic types and an (ordinary) environment that maps variables to values (semantic terms). Since we conflate types and terms in our setup, a typing environment is really nothing more than an ordinary environment.

type tp_env = env

Since the syntax of values is broken down into normal values and neutral values, we need a separate quote_neu to quote neutral terms. When working with neutral terms, the flow of typing information is reversed: rather than take as input a type to be broken down alongside the value to quote, quote_neu instead produces as output the type of the given neutral value. Recall that a neutral term is a stack of elimination forms ultimately blocked on a variable; that variable’s type will have been recorded in the context previously when quoting an abstraction, and the stacked elimination forms will break down that type.

let vvar l = VN (NVar l)
let lvl2ix d l = d - l - 1

let rec quote (d : lvl) (e : tp_env) (vA : vtp) (v : value) : tm =
    match vA, v with
    (* types: *)
    | VStar, VTop -> Top
    | VStar, VStar -> Star
    | VStar, VPi ((x, vA), fB)->
        Pi ((x, quote d e VStar vA), quote (d+1) ((x, vA)::e) VStar (fB (vvar d)))
    (* terms: *)
    | VTop, VUnit -> Unit
    | VPi ((_, vA), fB), v ->
        Lam (x, quote (d+1) ((x, vA)::e) (fB (vvar d)) (apply v (vvar d))
    | vA, VN n -> quote_neu d e n |> fst
    | _ -> failwith "quote: ill-typed"

and quote_neu (d : lvl) (e : tp_env) : neu -> tm * vtp = function
    | NVar l ->
        let i = lvl2ix d l in
        let vA = ix_lookup i e in
        (Var i, vA)
    | NApp (n, v) ->
        let t, vAB = quote_neu d e n in
        begin match vAB with
        | VPi ((x, vA), fB) ->
            (App (t, quote d e vA v), fB v)
        | _ -> failwith "quote_neu: ill-typed application subject"
        end

On a well-typed value, quoting always succeeds. This justifies two choices made in this implementation of quote and quote_neu.

1. In ill-typed cases, we simply use failwith.
2. In the last case of quote, where we call quote_neu, the expected type of VN n (i.e. the input vA of quote) will necessarily equal its actual type (the type synthesized by quote_neu). We simply throw out the synthesized type.

Normalization

A roundtrip through eval and quote brings a term into normal form.

let norm (t : tm) (tA : tp) : tm =
    quote 0 [] (eval [] tA) (eval [] t)

To enable normalization of open terms – we will encounter open terms during typechecking – it suffices to slightly generalize norm. I’ll delay the discussion of that generalization until we get to implementing the typechecker, in Part 2 of the series.

Next steps

In this post, we built a normalization procedure for a small, dependently-typed calculus by following the Normalization by Evaluation strategy. That strategy consists broadly in the following steps.

1. Define a semantics into which the language’s syntax is to be interpreted. In doing so, separate the introduction forms (normal terms) from the elimination forms (neutral terms), giving rise to a syntactic characterization of beta-normal forms.
2. Implement an evaluation procedure eval that interprets the syntax into the semantics. This eliminates all beta-redexes from a given term.
3. Implement the evaluation procedure’s inverse, called quote. This procedure converts a semantic term back into a syntactic term. This is a type-directed procedure. Typing information is crucially used to enable eta-expansion. This allows our system to identify as definitionally equal the terms f and λx. f x.
4. A roundtrip through eval and quote brings a term into normal form. This ultimately gives us a way to compare types which may contain terms requiring evaluation, as is common in a dependently-typed system.

Equipped with the normalization procedure implemented in this post, we’re ready to code up the typechecker for this small language. The central challenge in implementing the typechecker, as usual, will be the handling of variables and substitutions. I’ll approach that challenge in a few different ways in the next post.

At long last I combine my two passions to form an unlikely duo: mathematical rigour and bodybuilding.

For me, the biggest challenge in bodybuilding has always been to eat enough. The apps out there for tracking macros and meal planning always left me dissastisfied, so I even made my own called macro-traco. But in macro-traco, I made the same mistake that major apps like MyFitnessPal make: my model of food considered as distinct a number of things that are, at their core, not different at all. Those are foods, recipes, meals, meal plans/journals. Addressing that mistake is what led me to implement Nutcalc in a very natural way.

Getting to the bottom of food

What is a meal plan? It’s a schedule of meals to eat in a day. Really it’s the same as a meal journal, but with a forward-thinking perspective. So a meal plan (or journal) is composed of meals.

What is a meal? (I promise this isn’t a silly question.) Take the bodybuilder classic: chicken, broccoli and rice. In this case, the meal is composed of three foods. Take a less simple meal: pasta with bolognese sauce. Whereas pasta is just a food, bolognese sauce is a recipe you make yourself, unless of course you buy it in a jar.

Then, a recipe is something that doesn’t have a nutrition facts label, that you make yourself by combining various foods via a mysterious process known as “cooking.”

Finally, what is a food? A food, practically speaking, has a “nutrition facts” label on it, listing the nutrients of the food.

At last we reach the bottom of this model: nutrients sit indivisible, at the bottom of this hierarchy.

To recap, this exploration reveals the following levels, from largest to smallest: a meal plan/journal is composed of meals; a meal is composed of recipes; a recipe is composed of foods; and foods are composed of nutrients.

Modelling the hierarchy

A rigid model of this setup, as used in apps like MyFitnessPal or macro-traco could define each layer separately to consist of a list of items from the layer below. Such a model is inflexible – adding new layers requires substantial work – and even its initial setup is painful. (Five database tables?) And moreover, at a basic level, a straightforward such model would require that meals consist only of recipes, making it annoying to model meals like “chicken broccoli and rice” that are simply composed naturally of foods, living two (instead of one) layers away in the hierarchy.

Instead of choosing a particular number of layers, let’s define all infinitely many layers by induction. Here is an inductive definition of Food. (I’ll use capital-F Food to refer to items generated by this inductive process at any layer as opposed to lowercase-F food for those items described above as having nutrition facts labels, living at layer 1 of the model.)

• Base case. Each Nutrient is a Food.
• Step case. If \(F_1,\ldots,F_n\) are each a Food, then the Compound \(\langle F_1, \ldots, F_n \rangle\) is a Food.

Mathematically speaking, a Food is an \(n\)-ary tree whose leaves are labelled according to what nutrient is represented there.

In this inductive model, nutrition facts is defined by a totally straightfoward recursion on Food, which one normally learns to implement in kindergarten.

• If the Food is a Nutrient, then it is its nutrition facts.
• If the Food is a Compound \(\langle F_1, \ldots, F_n \rangle\), then its nutrition facts is simply the sum of the nutrition facts of each \(F_i\).

How much chicken?

I’ve thus far omitted a crucial aspect of the model: what are the quantities of the Foods that go into forming a Compound? Recall that Compounds are things like meals, which we would naturally express as “two cups of cooked white rice, one large chicken breast, and 200 grams of broccoli.”

In that natural language description, the foods are “cooked white rice”, “chicken breast”, and “broccoli”, but their respective quantities are “two cups”, “one large”, and “200 grams”. Quantities are composed of a number together with a unit. The units here are “cups”, “large,” and “grams”.

In light of this, I define a Quantified Food as a Quantity together with a Food. A Quantity is just a number and a unit. The unit is specific to the associated Food, so a “cup” of “cooked white rice” is distinct from a “cup” of “dry steel-cut oats.”

Now let’s revise the step case of the definition of Food.

• If \(Q_1, \ldots, Q_n\) are each a Quantified Food, then \(\langle Q_1, \ldots, Q_n \rangle\) is a Food.

Mathematically speaking, we still have an \(n\)-ary tree, but the edges in that tree are now labelled with Quantities.

The recursive definition of nutrition facts changes slightly: it suffices to multiply the nutrition facts of each Food in a Compound by its Quantity before summing.

In practice

Nutcalc implements this inductive model of food. It is a domain-specific programming language for defining Foods and performing computations on them.

A Nutcalc program consists of a series of Food definitions.

1 cup 'cooked white rice' weighs 158 g:
- 0.4 g fat + 4.3 g protein + 45 g carbs
- 1.9 mg iron

1 large 'chicken breast' weighs 120 g:
- 4.3 g fat + 37 g protein
- 89 mg sodium + 102 mg cholesterol + 1.2 mg iron + 307 mg potassium

100 g broccoli:
- 0.4 g fat + 7 g carbs + 2.8 g protein
- 33 mg sodium + 316 mg potassium

The effect of executing these statements in the Nutcalc interpreter is to define the Foods ‘cooked white rice’ and ‘chicken breast’ together with the respective units ‘cup’ and ‘large’ whose equivalent weights are given. For broccoli, since the definition is already for a particular weight, there’s no need for a ‘weighs’ clause.

Since Nutcalc uses the inductive model of food, we use the same syntax to define meals. For meals, the most natural unit is usually ‘serving’, but that’s long to type so how about x?

1 x 'chicken broccoli rice':
- 1 large 'chicken breast' + 2 cup 'cooked white rice' + 200 g broccoli

When a ‘weighs’ clause is omitted but the definition defines a new unit (here x), the weight of the new unit is inferred as the sum of the weights of the constituent Foods. For meals, not only is that assumption about weight usually correct, but we often don’t care anyway about the weights of Foods at the higher levels in the model.

Again, we use the same syntax to define a meal plan:

1 x Monday:
- 1 x 'oatmeal breakfast'
- 1 x 'eggs sausage bacon toast lunch'
- 1 x 'chicken broccoli rice'
- 1 x 'protein shake'

Of course, the same syntax is used to define a food journal:

1 x '2025-02-20':
- 0.5 cup 'dry steel-cut oats' + 50 g walnuts + 2 cup '3.25% milk'
- 4 x 'breakfast sausage' + 4 large egg + 2 slice toast + 2 tsp butter
- 1 medium 'chicken break' + 1 medium 'chicken leg' + 2.5 cup 'cooked white rice' + 200 g broccoli

After loading a file with our definitions, we can compute aggregate nutrition facts easily:

$ nutcalc -i journal.nut
nutcalc> facts 1 x 'oatmeal breakfast'
energy: 1122.83 kcal
protein: 44.30 g
fat: 46.57 g
carbs: 131.62 g
water: 150.60 g
calcium: 602.41 mg
iron: 4.35 mg
potassium: 930.49 mg
sodium: 84.34 mg
zinc: 2.71 mg
cholesterol: 30.12 mg
nutcalc>

In fact, facts accepts an expression on the right, which can be a lone Quantified Food or a sum thereof, e.g. 2 cup 'cooked white rice' + 200 g broccoli.

A programming language?

Some might laugh at the idea of calling Nutcalc a “programming language.” It has no loops, no conditions, no mutable variables, no functions, no recursion. It is certainly not Turing-complete.

In my view, Nutcalc is a programming language because it has a syntax and a semantics. Both are very simple, but I see that as a strength of the language, not a weakness.

Overall, whether Nutcalc is or isn’t a programming language is incidental. It’s useful. Anyway, the main contribution is the inductive model of food, giving a unified view of foods, meals, and so on.

As I prepare for technical interviews, I’ve been grinding Leetcode to practice my data structures and algorithms. (And if you’re looking to hire in 2025, hit me up.) I want to share detailed explanations of three different, fairly short solutions to Problem 10, “regular expression matching.” I especially find it interesting when a problem admits multiple, workable solutions. Part of the beauty of computer science and software engineering is exactly that there are so many different, interesting ways to solve problems.

The problem: not really regex!

We’re asked to define a function match_pattern(s: str, p: str) -> bool that determines whether the given string s matches the given regular expression p, subject to some constraints.

• the lengths of both s and p are between 1 and 20 (inclusive)
• s contains only lowercase English letters
• p contains only lowercase English letters, plus . and *.

The . pattern matches any character, and the * is a modifier meaning to match zero or more times the preceding character. There is an additional restriction that the input p be “well-formed,” meaning that there isn’t a * as the very first character, nor that there are two *s in a row.

Moreover, for our match_pattern to output True, the pattern must cover the entire input string.

This might sound different from the regex you know and love:

• There are no parentheses, so repetitions with the star are only for a single character.
• There are no alternations, so no patterns like a|b.

This vastly simplifies the problem down from full regex matching!

In the absense of any *, this is just a simple string-matching problem. We could just traverse s and p simultaneously, checking along the way that each character of p matches the corresponding character of s, letting any . in p match any character in s.

In other words, when e.g. s = "abc" and p = "a.c", we can walk both strings together and compare character by character. When comparing 'b' with '.', we say yes. We would arrive at the end of both inputs at the same time, and along the way said yes everywhere, so we decide that this string matches this pattern.

To illustrate the challenge involved in handling the *, let’s walk through an example in detail. Say s = "aab" and p = "a*b". Set up two pointers: let i be an index into s and j be an index into p, both starting at zero.

• We see a star at index j+1, so we have to decide:
  • do we skip the star, setting j = j+2 but leaving i the same? Or,
  • do we use the star, setting i = i+1, but leaving j the same?

That decision is the crux of the challenge: we can’t know for sure which choice to make.

Using brute force: backtracking search

The idea is this: we don’t know which choice to make, so try both ways!

A backtracking search will, in the worst case, explore the entire decision tree of the problem, which we can visualize like this. Rather than show the indices, I’ll show s and p changing as we move through them. States in red are those where we return false, and the green state is where we return True as we successfully matched the string with the pattern.

• The leftmost state does not have a “use star” branch: to use the star, the head of s (i.e. 'b') must match the head of p (i.e. 'a').
• When the current state’s pattern’s head doesn’t have a star, the tree doesn’t branch.

We can code this backtracking search as a recursive algorithm.

def match_pattern(s, p):
    if s == '' and p == '': return True
    if p == '': return False # pattern empty, but more string left
    if len(p) >= 2 and p[1] == '*':
        # right branch: skip star
        if match_pattern(s, p[2:]): return True
        # but if that fails, then we have to use the star
        return (
            len(s) # for that, we need at least one char in s
            and match_char(s[0], p[0]) # it must match the head of p
            and match_pattern(s[1:], p) # and p has to match the rest of s
        )
    else: # we're not handling a star, so we need:
        return (
            len(s) # there to be at least one character in s
            and match_char(s[0], p[0]) # that it match the head of p
            and match_pattern(s[1:], p[1:])
            # that the rest of the string match the rest of the pattern
        )

The match_char function is simply an equality check on characters that accounts for the pattern character '.' being equal to any string character.

To assess the time complexity of this solution, let’s consider an input with lots of stars: match_pattern("aaaaab", "a*a*a*a*c")

Of course, the output should be False due to the last characters being mismatched, but to discover this, the backtracking search has to explore the entire tree. So how big does the tree become for this input?

When there’s a star, we make two recursive calls, each on an input that’s smaller only by one. Overall, we get an exponential time complexity.

However, you might notice that a sequence of a*a*... is equivalent to just a single a*. Therefore, we can ‘optimize’ the regular expression before interpreting it. By doing this and prioritizing the ‘use star’ branch of the backtracking search we get a greedy formulation of the search that avoids exponential blowup in most cases. This formulation is performant enough to pass the time requirements on Leetcode.

Avoid repeating work: dynamic programming

The subproblems to solve in the recursive algorithm above are overlapping.

To see why, let’s draw out some of the recursive calls for match_pattern("aaab", "a*a*a*c")

See? There are two different paths that lead us to solving the subproblem match_pattern("aab", "a*a*c"):

1. Left (‘use star’) then right (‘skip star’)
2. Right (‘skip star’) then left (‘use star’)

The essence of dynamic programming is to solve each unique subproblem exactly once, and there are two broad approaches to accomplishing that.

Strategy 1: top-down with memoization

This technique is quite simple: take the recursive algorithm and attach a cache of solutions to it. This cache will be a hashtable, initially empty. Each key will be a tuple of inputs to match_pattern. At the very beginning of match_pattern, we consult the cache to see if we’ve already solved the problem, returning the cached result in that case. Else, we actually do the computation, and store its result in the cache before returning.

def match_pattern_dp(s, p, cache):
    if s == '' and p == '': return True
    if p == '': return False

    # answer already known:
    if (s, p) in cache: return cache[(s, p)]

    if len(p) >= 2 and p[1] == '*':
        outcome = (
            # skip star:
            match_pattern_dp(s, p[2:], cache)
            # or use star:
            or len(s)
            and match_char(s[0], p[0])
            and match_pattern_dp(s[1:], p, cache)
        )

    else:
        outcome = (
            len(s)
            and match_char(s[0], p[0])
            and match_pattern_dp(s[1:], p[1:], cache)
        )

    cache[(s, p)] = outcome
    return outcome

With the cache in place, whenever we run into a subproblem we’ve seen before, we can just look up its answer.

Just by looking this program, however, it’s not so obvious what the time complexity is. Instead, it’s more illuminating to look at the tree again:

With the addition of the cache, it doesn’t matter anymore by what path we arrive at a particular subproblem. We can visualize this by fusing the duplicated middle nodes.

Imagining that we continue fusing duplicated middle nodes all the way, the structure that this top-down memoizing algorithm ends up exploring takes the shape of a rectangle.

To determine the time complexity of our algorithm, it suffices to count the nodes in the rectangle. One side’s length is bounded by the length of s – let’s call that \(n\) – and the other side’s length is bounded by the length of p – let’s call that \(m\). We get an \(O(n\times m)\) complexity overall then.

Equipped with this better understanding of this problem’s underlying structure, we can exploit this structure to design an even better implementation.

Strategy 2: bottom-up table construction

Function calls are slow. Rather than express the exploration of the problem’s rectangular state space as a recursive function + a hashtable, we can instead directly build the rectangle as a matrix, using nested loops. CPUs hate function calls, but love loops. Of course, this won’t improve the asymptotic complexity of the algorithm – it will still be \(O(n \times m)\) – but the constant factors, which matter in Real Life, will be better.

The game plan is to write a function that will construct a 2D array T, such that T[i][j] holds the answer to the problem “does the pattern formed by taking the last i characters of p match the string obtained by taking the last j characters of s.”

Sheesh. That’s a complicated definition. I picked it because it really is a bottom-up version of what we did in the previous section.

• The base case of the recursive algorithm ends up in T[0][0] = True – the empty pattern matches the empty string.
• But wait, there’s more base case: an empty pattern can never match a nonempty string; that’s the rest of the first row of T. We’ll set T[0][j] = False for all j>0 up to and including len(s).
• Then, we build T row by row until we’ve filled in T[len(p)][len(s)], which holds the answer to the full problem.
  • T[i][0] (for i>0) asks “does the pattern formed from the last i characters of p match the empty string?”
    • This is possible only when that pattern’s head is starred, and when T[i-1][0] is True
  • the value of T[i][j] (for i>0 and j>0) depends crucially on p[-i]
    • If it’s unstarred, then take the AND of T[i-1][j-1] – does the rest of the pattern match the rest of the string? – and whether p[-i] matches s[-j].
    • If it’s starred, then take the OR of the skip star and use star calculations.
    • Skip star is easy: it’s precisely T[i-1][j]
    • Use star on the other hand, has us check that p[-i] match s[-j] and AND that with T[i][j-1].

To simplify the implementation slightly, we’ll preprocess the input pattern p to ‘tokenize’ it. This process will turn a pattern like ba*a*bb*c into ['b', 'a*', 'a*', 'b', 'b*', 'c']. This avoids an uncomfortable situation where we take e.g. the last two characters of that pattern, giving the illegal pattern *c.

def preprocess(p):
    out = []
    i = 0
    while i < len(p):
        if i+1 < len(p) and p[i+1] == '*':
            out.append(p[i:i+2])
            i += 2
        else:
            out.append(p[i])
            i += 1
    return out

def match_pattern_tabular(s, p):
    p = preprocess(p)
    T = [[True] + [False] * len(s)]
    for i in range(1, len(p) + 1):
        p_c = p[-i]
        T[i][0] = p_c.endswith('*') and T[i-1][0]
        for j in range(1, len(s) + 1):
            c = s[-j]
            T[i][j] = (
                match_char(p_c, c) and T[i-1][j-1]
                if not p_c.endswith('*') else (
                    # use star
                    match_char(p_c[0], c) and T[i][j-1]
                    # skip star
                    or T[i-1][j]
                )
            )
    return T[len(p)][len(s)]

The nested loops in this approach make obvious the \(O(n \times m)\) time complexity.

We could further improve the memory usage of this approach by observing that only the previous row is needed to compute the next row. I leave that as an exercise to the reader.

The general solution: interpret an NFA

I was first intrigued by this Leetcode problem because it had me remember my theory of computation class. We spent quite some time in that course talking about different kinds of automata, which are (idealized) models of computation. I learned in that course that regular expressions correspond to nondeterministic finite automata (NFAs), and that a regex can be converted into an NFA according to an algorithm called Thompson’s construction. Yes, that’s Ken Thompson, co-creator of Unix!

But what’s an NFA?

Let’s start with an example. Here’s an NFA that corresponds to the regex a*b*c:

This diagram is a representation of an abstract machine into which we input a string, that either accepts or rejects the string.

Fundamentally, an NFA is a state machine. The circles in the diagram are the different states the machine can be in. The letters inside the states are just some arbitrary names for those states. The arrows are the transitions of the machine, and the label of an arrow indicates what letter we need to see at the head of the string to advance to the corresponding state on the rest of the string. However, the label \(\epsilon\) is special: this is a transition that we can make “for free” without needing to consume the head letter from the string.

The machine accepts a given string if after processing the whole string, the machine is in an accepting state, indicated by the double circle. The machine rejects the string under two situations.

1. The machine’s current state cannot accept the head letter of the string.
2. The string becomes empty while the machine is not in an accepting state.

An NFA is nondeterministic. What that means is that the transition to take to move to the next state isn’t (always) uniquely determined. For example, in the first state on the left (or the second one), we could follow the \(a\) (or \(b\)) transition or the \(\epsilon\) transition whenever the head of the string is an a (or b, respectively). More generally, NFAs might have multiple transitions with the same label.

This somewhat complicates the idea of “accepting” a string. How does the NFA “know” which choice to make, when more than one transition is possible? It doesn’t know! This is a mathematical object, after all. In reality, the acceptance criterion for an NFA is more precisely that there exist a sequence of transitions that consumes the entire input string and ends on an accepting state.

But let’s not get ahead of ourselves. Before we see how to discover whether such a sequence of transitions exists, we need to decide how to represent an NFA in code, and how to translate a regex into an NFA.

Constructing an NFA from a regex

Fortunately, the regex in this problem are not really regex! Therefore, we don’t need to implement the entire Thompson construction. Just a part of it will do.

First, let’s choose a way of representing an NFA. We’ll identify the states of the NFA with numbers, and use an array to represent the NFA itself. Indexing into that array at \(i\) will give us a description of the transitions that are possible in state \(i\). We’ll describe a set of transitions as a list of key-value pairs, with the key being the letter we need to see and the value being the state index to move to in that case. To represent \(\epsilon\)-moves, we’ll use the empty string.

We can represent the NFA from before, for the regex a*b*c, as the following.

abc = [
    [('a', 0), ('', 1)],
    [('b', 1), ('', 2)],
    [('c', 3)],
    [],
]

Due to the simplified nature of the regex in this problem, we’ll always have a unique accepting state, and it will always be the last state in the list.

def regex_to_nfa(p):
    nfa = []
    i = 0
    while i < len(p):
        new_state_index = len(nfa)
        if i + 1 < len(p) and p[i+1] == '*':
            new_state = [ ('', new_state_index + 1), (p[i], new_state_index) ]
            i += 2
        else:
            new_state = [ (p[i], new_state_index + 1) ]
            i += 1
        nfa.append(new_state)
    nfa.append([]) # the accepting state
    return nfa

Equipped with a Python model of NFAs and a procedure for constructing an NFA from a regex, we’re ready to tackle the nondeterminism of the NFA to discover how to actually run one of these machines.

Running an NFA

The computers that we use are deterministic, so to run an NFA, we have to simulate its nondeterminism. One way to do this is to keep track of a set of states that the NFA is in simultaneously. To see how this works, let’s walk through running the NFA of the regex a*b*c on the string aac.

• The initial set of states is \(\{A, B, C\}\), i.e. the state marked by “start” plus all the states reachable from there by using only \(\epsilon\)-moves.
• For each state that the NFA is in, we check whether that state can consume the head letter a, and form the set of all the states that result from following that transition + all the states then reachable by using only \(\epsilon\)-moves. The states \(B\) and \(C\) can’t accept a, so they “die”. The state \(A\), however, can consume a, leading back to state \(A\) + the states \(B\) and \(C\) which are reachable from \(A\) by using \(\epsilon\)-moves. Therefore we get back the same set of states, \(\{A, B, C\}\), only now the string that’s left to process is ac.
• The same thing happens again, leading to the next set of states being \(\{A, B, C\}\) and the remaining string to process being just c.
• States \(A\) and \(B\) can’t consume c, so they die, but state \(C\) can consume c, leading to state \(D\).
• The input string runs out, and the final set of states is just \(\{D\}\).
• Among that set of states is an accepting state, so overall the string is accepted.

You also might have noticed from the description above that “all the states reachable from there by using only \(\epsilon\)-moves” shows up a few times. This concept has a fancy math-name: it’s called the \(\epsilon\)-closure of a state. It’s easily expressed as a depth-first search through the NFA, following only \(\epsilon\)-moves. The NFAs in our setting never have \(\epsilon\)-moves that jump “backwards”, so we don’t need to worry about marking nodes as visited.

def epsilon_closure(nfa, i):
    closure = set([i]) # the state itself is part of the closure
    for (letter, j) in nfa[i]:
        if letter == '':
            closure.update(epsilon_closure(nfa, j))
    return closure

Finally, we’re ready to simulate an NFA.

def run_nfa(nfa, s):
    active_states = epsilon_closure(nfa, 0)

    for c in s:
        # each iteration of the outermost loop
        # calculates a new set of active states.
        next_active_states = set()
        for i in active_states:
            for (letter, j) in nfa[i]:
                if match_char(c, letter):
                    next_active_states.update(epsilon_closure(nfa, j))
        active_states = next_active_states

    # accept or reject the string according to whether
    # the accepting state is among the active states, now that
    # the entire string has been traversed.
    return len(nfa)-1 in active_states

Now let’s evaluate the time complexity of run_nfa. The two outermost loops require that it be at least \(O(n \times m)\), but it might be worse than that. We need to account for the call to epsilon_closure that happens in the innermost loop!

Recall that calculating an \(\epsilon\)-closure requires making a depth-first search through the NFA. In a general graph, a depth-first search has a time complexity of \(O(V + E)\); in our setting the vertices are the states of the NFA and the edges are the transitions. The number of transitions per state is bounded by a constant, namely \(2\) – just look at regex_to_nfa. Therefore, to calculate the \(\epsilon\)-closure of one state takes \(O(m)\) time.

It suffices at this point to determine how many \(\epsilon\)-closures we need to calculate. Notice that the call to epsilon_closure is guarded by the condition if match_char(c, letter). This condition can be true at most once per entire loop! That’s because there are at most \(2\) transitions per state in our NFAs, and at most \(1\) that’s not an \(\epsilon\)-move.

Therefore, thanks again to the restricted nature of the regex in this problem, the simulation of the NFAs that result from these regex still takes \(O(n \times m)\) time.

In a general NFA, where the number of transitions per state labelled with a given letter is bounded by the total number of states in the NFA (instead of a constant), the time complexity would be worse.

Conclusion

Leetcode problem 10 is a fascinating one because it admits three different, workable solutions. The backtracking, exponential-time algorithm is performant enough to pass the automated checker so long as we make a simple optimization of collapsing a*a*... down to a*. The two dynamic programming solutions admit a much better time complexity of \(O(n\times m)\).

1. The top-down implementation uses a hashtable to store previously seen solutions of subproblems.
2. The bottom-up implementation directly exploits the resulting rectangular shape of the problem’s state space.

Finally, the most theoretically rigorous and (in principle) general solution to the problem comes from the study of the theory of computation: we translate the regular expression into a nondeterministic finite automaton and interpret it by tracking a set of active states to simulate the nondeterminism. Of course, thanks to the simplified nature of the regex in the problem, we were able to cut some corners to simplify the implementation as well.

Some languages define a variant of return called yield. When a function returns normally, it’s finished, but when a function yields, the execution context of the function is saved, enabling us to re-enter the function to resume its execution. Such resumable functions are especially convenient for defining lazily generated sequences. Such functions that generate sequences are called generators.

This article will first describe native generators in Python before translating the idea into OCaml using mutable variables and continuation-passing style (CPS). Then, we will see how to eliminate the mutable variable and how to eliminate the higher-order functions that arise from CPS, leading to an implementation that we translate into C.

Native generators in Python

For example, consider this recursive algorithm in Python that enumerates all truth assignments on n variables. Each truth assignment is represented as a list of booleans of length n. The sequence that arises from this enumeration has length 2^n since there are two possible choices for the value of each of the n variables.

def enumerate_assignments(n):
    a = [True] * n

    def go(i):
        if i == 0:
            yield a
        else:
            a[n - i] = True
            yield from go(i - 1)

            a[n - i] = False
            yield from go(i - 1)

    return go(n)

In the base case when i == 0, the variable a contains a truth assignment with some Trues and some Falses put in it by the stack of recursive calls that leads to the base case. We yield that truth assignment, suspending execution of the function. The function, upon being re-entered, proceeds from that yield, in this case returning back to its caller. The caller might be the recursive call yield from go(i - 1), after which a[n - i] is set to False and another recursive call is made.

Therefore, calling this function does not generate the whole sequence at once. In fact, nothing happens just yet if we call enumerate_assignments(5) except to ‘prime’ the generator to run. The call returns a so-called iterator, which stores the state of the execution. Then, calling the function next() on the iterator will resume execution of the function up to the next yield. Also, next() returns to us the yielded value. Repeatedly calling next() until the generator exits is exactly what a Python for-loop does, so we can print all the truth assignments on five variables like this:

for a in enumerate_assignments(5):
    print(a)

What if the language we’re working in doesn’t have yield though? How can we implement something like this in, say, OCaml? In doing so, we will explore what happens under the hood of generators.

Evaluation contexts and continuations

Before diving into the implementation of generators in OCaml using continuations, some background on continuations is required. I’ve already written an article on continuations in general here, but I’ll give a somewhat different take on the core idea in this section.

To motivate the discussion of continuations in this section, let’s look ahead at what we will want to accomplish in the next section: we will implement a generator as a stateful function next. Each time we call this function – next () – it returns the next item in the sequence. The key idea in implementing this function is that just before returning an item in the sequence, we store the current evaluation context in a reference. Then, when the user of our generator calls next () again, we restore the saved evaluation context. Evaluation then proceeds from that saved point up to the next item in the sequence.

I said ‘evaluation context’ several times in the last paragraph (and in the title of this section) so it’s about time I define what that is and how it’s related to continuations.

When a program is “in a call to a function”, for instance, there’s somewhere that call will return to. That location corresponds to how the return value of the function is used. Or, said differently, that location is the evaluation context.

For example, let’s imagine we have a function f : unit -> int. Maybe somewhere in our program there is let x = f () in E. The evaluation context of that call to f – the way the result of f is used – is that it is associated to the variable x in evaluating E. We can write this evaluation context as let x = _ in E as this shows where the return value of the call will be used. Notice that let x = _ in E is not a program! Rather, it’s a program with a hole in it.

Or as another example, perhaps the program contains f () + 5. The evaluation context inside the call to f is that the result of the call will be added with five. We can write that again with this ‘hole’ syntax: _ + 5.

Of course, evaluation can reach quite deep into a subexpression. Consider this evaluation context let x = if _ > 17 then E1 else E2 in E3.

The relevant insight about evaluation contexts, which are a concept outside our programming language, is that we can represent them in our programming language: we represent an evaluation context as a function. Let’s see how this applies to the examples seen so far.

• let x = _ in E is represented by fun r -> let x = r in E
• The second example _ + 5 is represented by fun r -> r + 5
• The last example let x = if _ > 5 then E1 else E2 in E3 is represented by fun r -> let x = if r > 5 then E1 else E2 in E3

This functional representation of an evaluation context is called a continuation.

Now that we have a way of representing evaluation contexts as continuations, we can write OCaml programs that manipulate continuations. Since continuations are functions, we cannot inspect them: we can only construct them and call them. Calling a continuation k with an argument a represents filling the hole in the evaluation context represented by k with the value a and proceeding to evaluate the resulting expression.

For example, if k is the continuation fun r -> let x = if r > 5 then E1 else E2 in E3 and we apply this to 3, then evaluation continues from the point where the value of the hole _ was required in the represented evaluation context, namely in computing _ > 5. The value of 3 > 5 is needed in the context if _ then E1 else E2, and the value of that if-then-else expression is needed in the context let x = _ in E3. From this apparent nestedness of evaluation contexts, we can observe that the contexts form a stack. This observation will be expanded on considerably in the last section of this article when we translate these ideas into C.

    (fun r -> let x = if r > 5 then E1 else E2 in E3) 3
==> let x = if 3 > 5 then E1 else E2 in E3    -- substitute 3 for r
==> let x = if true then E1 else E2 in E3     -- compute '>'
==> let x = E1 in E3                          -- select then-branch
==> ...

Great, evaluation contexts can be represented as functions called continuations and our programs can use continuations by calling them. Calling a continuation corresponds to filling the hole in the represented evaluation context and continuing the evaluation from there.

The big idea of continuation-passing style (CPS) is to equip the functions we write with an extra parameter. You guessed it, that extra parameter is the continuation of the function. So instead of having a function f : A -> int that we use like f a + 5, we instead write f a (fun r -> r + 5). The upshot is that in the implementation of f, we will have access to (a representation of) the evaluation context in which the call to f takes place!

In traditional CPS, functions no longer ‘return normally’. Instead, they return by calling the continuation with the value they want to return. In our current setting of implementing generators, this isn’t quite what we want. We would like our functions to return normally – this is how the generator will emit a value – but we nonetheless want access to the evaluation context so that we can store it in a reference just before returning. This will enable us to resume the function from the point where it emitted a value. In the next section, we’ll see how we can transform the Python code from the beginning of the article into OCaml using this form of CPS.

Implementing generators using state and CPS

In this section, we mimic Python’s approach to generators: generators in Python are stateful objects, so our implementation in this section will use a reference to store an evaluation context.

To come up with our OCaml implementation, let’s first translate the Python program from the first section into pseudo-OCaml with yield & yield from. Rather than use a mutable array to construct the truth assignment, we’ll build it up one value at a time in a parameter.

let enumerate_assignments n =
    let rec go n a =
        if n = 0 then
            yield a
        else begin
            yield from go (n-1) (true :: a);
            yield from go (n-1) (false :: a)
        end
    in
    go n []

Let’s concentrate specifically on the inner function go. We will need to find a way to represent yield and yield from in genuine OCaml. First, let’s deal with yield.

Operationally, yield saves the current evaluation context and emits the given value. We choose to represent ‘emitting a value’ as simply returning. To gain access to the current evaluation context so that we can save it, we rewrite go in CPS. We also introduce a ref to store the evaluation context.

let state = ref (* what to put here initially ? *) in
let rec go n a next =
    if n = 0 then
        (state := next; a)
    else begin
        yield from go (n-1) (true :: a);
        yield from go (n-1) (false :: a)
    end
in ...

Next, let’s address yield from. This keyword causes the current generator to invoke a ‘sub-generator’ and to yield all of its values.

Since we choose to represent yielding a value as simply returning it, we implement yield from go ... as simply calling go. In doing so however, we must pass go a continuation. We can determine what continuation to pass by looking at the evaluation context of yield from go (n-1) (true :: a):

_; yield from go (n-1) (false :: a)

We represent this evaluation context as a function: fun x -> x; yield from go (n-1) (false :: a). Since yield from does not compute anything – it performs an effect – we observe that x : unit. A value of type unit does not convey any information, so we simplify to fun () -> yield from go (n-1) (false :: a).

Next we must translate the inner yield from go ... which appears inside the continuation. Again, we translate this simply to a call to go, but in doing so we must decide what continuation to pass in this call. We ask ourselves what evaluation context this call takes place in: what happens next, after yield from go (n-1) (false :: a) finishes generating its sequence of values? The answer is that go returns to its caller. In other words, the evaluation context of yield from go (n-1) (false :: a) is the evaluation context of go itself, which was passed to go in the continuation next. Therefore, the continuation we pass to this second call to go is simply next.

Let’s take stock of the translation so far:

let state = ref (* ? *) in
let rec go n a next =
    if n = 0 then
        (state := next; a)
    else
        go (n-1) (true :: a) (fun () -> go (n-1) (false :: a) next)
in ...

Let’s zoom out, look at enumerate_assignments again, and see how we can fit it together with our adjusted go.

let enumerate_assignments n =
    let state = ref (* ? *) in
    let rec go n a next =
        if n = 0 then
            (state := next; a)
        else
            go (n-1) (true :: a) (fun () -> go (n-1) (false :: a) next)
    in
    go n [] (* ? *)

Calling go right away at the end of enumerate_assignments can’t be correct anymore. This will return the first item of the sequence, and then we’ll have no way to call the stored continuation to get the next one!

Instead, we need to adjust the return type of enumerate_assignments. Let’s look back to the pseudocode using yield and yield from to inform how.

let enumerate_assignments n =
    let rec go n a =
        if n = 0 then
            yield a
        else begin
            yield from go (n-1) (true :: a);
            yield from go (n-1) (false :: a)
        end
    in
    go n []

First of all, what is this enumerate_assignments supposed to return? What is the return type of go in this pseudocode? If we think back to the original Python implementation, it was returning an iterator. This is some kind of object that holds the suspended execution state. We call the function next() on this iterator to resume execution up to the next yield.

We can imagine in our OCaml pseudocode that go returns something of type bool list gen and therefore that enumerate_assignments n : bool list gen. We can further imagine that there’s a function next : 'a gen -> 'a option which pumps the generator for one more item. This hypothetical next returns an option because the sequence can end and we need a way to signal this.

Now we need to turn the fantasy of the type 'a gen and its associated function next : 'a gen -> 'a option into reality. Here there are several approaches available to us, but one particularly clean one using higher-order functions is to represent 'a gen as a function unit -> 'a option. Then, next becomes trivial to implement.

let next f = f ()

Hence, we implement enumerate_assignments to return a function unit -> bool list option, such that each time this function is called, it emits the next item in the sequence.

let enumerate_assignments n =
    let state = ref (* ? *) in
    let rec go n a next =
        if n = 0 then
            (state := next; a)
        else
            go (n-1) (true :: a) (fun () -> go (n-1) (false :: a) next)
    in
    (* previously: go n [] *)
    fun () -> ???

We’re almost finished now. What’s left is to make it so that the first time the returned function is called, it emits the first item in the sequence.

The missing insight has to do with the initial value of state: it should be a function that makes the initial call go n []. This makes it so that the ‘driver function’ returned by enumerate_assignments can be implemented as fun () -> !state (). Recall that the state stores a function that generates the next item in the sequence, so initially we store the function that generates the first item of the sequence fun () -> go n []. Moreover, we implemented go to save the current continuation back into the state just before it returns. Therefore, the next time the driver function is called, it emits the next item in the sequence!

But there’s a small wrinkle in fun () -> go n []: there’s an argument missing! What continuation do we pass in this initial call to go?

The continuation passed here ends up saved by go in state after the whole sequence of truth assignments is generated. Therefore, we need to arrange that when this continuation ends up stored, the driver function returns None. The driver function has the form fun () -> !state (), so this initial continuation passed to go ought to be fun () -> None.

But previously, the continuation had the type unit -> 'a, so we need to adjust go slightly to accommodate this change. This leads to the finalized generator using CPS and state.

let enumerate_assignments n =
    let rec go n a next =
        if n = 0 then
            (* we wrap the next item in Some *)
            (state := next; Some a)
        else
            go (n-1) (true :: a) (fun () -> go (n-1) (false :: a) next)
    and state = ref (fun () -> go n [] (fun () -> None))
    (* and arrange that the last continuation to be stored in
       `state` just returns `None`. *)
    in
    fun () -> !state ()

Now the state reference, which initially contains a call to go, must be mutually recursive with go, which refers back to the state reference.

Let’s witness the fruits of our handiwork. Here’s an OCaml REPL demonstrating the generator.

> let next = enumerate_assignments 5 ;;
val next : unit -> bool list option = <fun>

> next () ;;
- : bool list option = Some [true; true; true; true; true]

> next () ;;
- : bool list option = Some [false; true; true; true; true]

> next () ;;
- : bool list option = Some [true; false; true; true; true]

> next () ;;
- : bool list option = Some [false; false; true; true; true]

In the next section, we explore how to eliminate the mutable variable from this implementation. This will give rise to an implementation suitable to purely functional languages, which lack (genuine) mutable variables.

Eliminating state

The implementation we have arrived at can be made stateless. Notice that there are two ways that the driver function fun () -> !state () can return: it can return Some a after storing the continuation, or it can return None which happens when the sequence ends and the initial continuation fun () -> None has been stored in the variable state.

Rather than store the continuation in some hidden stateful variable, we can simply return both the assignment and the continuation in the Some case. That gives rise to the following intuitive implementation.

let enumerate_assignments n =
  let rec go n a next =
    if n = 0 then
      Some (a, next) (* return the value _and_ the continuation *)
    else
      go (n-1) (true :: a) (fun () -> go (n-1) (false :: a) next)
  in
  go n [] (fun () -> None)

Slight problem with this implementation: it doesn’t typecheck! And the error is not a simple “expect this type, got this other type,” but rather

Error: The expression `go (n-1) (false :: a) next`
       has type (bool list * (unit -> 'a)) option
       but an expression was expected of type 'a
       The type variable 'a occurs inside (bool list * (unit -> 'a)) option

According to this error, the return type of go, which is so far inferred as (bool list * (unit -> 'a)) option (due to the expression Some (a, next)) has to equal the return type of next which is so far inferred as 'a. This circularity is forbidden, so OCaml rejects the program.

We can resolve this by introducing a recursive type, let’s say L, such that L = (bool list * (unit -> L)) option. The base case of this recursive type arises from the None constructor of the option type. Now we can fix the circular variable type variable 'a to be L and eliminate the forbidden circular instantiation.

type l = Next of (bool list * (unit -> l)) option

We can slightly refactor this type by introducing a second constructor and eliminating the option.

type l =
  | Done
  | More of bool list * (unit -> l)

And we can generalize the type by replacing bool list with a type variable.

type 'a l =
  | Done
  | More of 'a * (unit -> 'a l)

And would you look at that! This is simply a list, but whose tail is computed by a function unit -> 'a l rather than being already materialized.

Now we’re equipped to rewrite enumerate_assignments but having the type int -> unit -> bool list l.

let enumerate_assignments n =
  let rec go n a next =
    if n = 0 then
      More (a, next) (* return the value _and_ the continuation *)
    else
      go (n-1) (true :: a) (fun () -> go (n-1) (false :: a) next)
  in
  fun () -> go n [] (fun () -> Done)

Have we achieved our goal of making enumerate_assignments stateless? Yes and no.

Indeed we have eliminated the mutable variable, so on the one hand we can say “mission accomplished.” But on the other hand, the state of the walk through the space of truth assignments is still very much present in our program. That state is captured in the continuation, which is returned explicitly via the More constructor of our lazy list type l. The state of our generator implementation is no longer mutable and hidden, but rather immutable and explicitly passed around. We can therefore view lazy lists as purely functional generators.

Another consideration is that in order to make this approach work in a strongly and statically typed setting as in OCaml, we did need to introduce the recursive type l. In the setting of a different type system, this might not be necessary. For instance, in a dynamically-typed setting, e.g. in Python, it is unnecessary to introduce an extra type: we can simply return None when the sequence ends and (a, next) when the sequence continues.

In the next section, we revisit our implementation using hidden, mutable state, and eliminate from it the higher-order functions. This will give rise to a first-order implementation suitable for translation into a language such as C, which is moreover well-equipped to handle mutable state.

Defunctionalizing the continuation of a generator

Recall from the first section the enumerate_assignments implementation using mutable state.

let enumerate_assignments n =
  let rec go n a next =
    if n = 0 then
      (state := next; Some a)
    else
      go (n-1) (true :: a) (fun () -> go (n-1) (false :: a) next)

  and state = ref (fun () -> go n [] (fun () -> None))
  in
  fun () -> !state ()

We can apply defunctionalization to eliminate the higher-order functions present in this program (see here). In short, we replace each function type \(T = T_1 \to T_2\) that occurs in the program with a new datatype \(D(T)\). Then, for each function \(\Gamma \vdash \text{fun}\, x \to e_i : T_1 \to T_2\), define a constructor \(C_i : P(\Gamma) \to D(T)\) where \(P(x_1 : S_1, \ldots, x_n : S_n) = (S_1, \ldots, S_n)\). Next, define the function \(\text{apply}\, : D(T) \to T1 \to T2\) as follows (in pseudo-OCaml)

let apply (f : D(T)) (x : T1) : T2 = match f with
  | C_i (x1, ..., xN) -> e_i

In other words, the function apply takes the representation \(f : D(T)\) of the original function and recovers the original function: notice that \(\text{apply}\, f : T_1 \to T_2\) has the original function’s type!

Concretely for enumerate_assignments, we introduce a type stack to represent the function type unit -> bool list option. There are three functions of this type passed as arguments to other functions.

• fun () -> go n [] (fun () -> None): this is the initial continuation and it contains a free variable n : int. We generate from this a constructor Start : int -> stack.
• fun () -> None: this is the final continuation and it contains no free variables. We generate a constructor Finished : stack
• fun () -> go (n-1) (false :: a) next: this is the continuation passed when making a recursive call to go. It has the free variables n : int, a : bool list, and next : unit -> bool list option. Notice that the function we’re translating refers to a function of the type we’re defunctionalizing. This will make our type stack into a recursive type. Moreover, since the new continuation defined by this anonymous function refers to exactly one other continuation, we are finally justified in calling our type “stack”: these continuations were implicitly forming a linked list that now is explicitly represented. From this analysis we generate the constructor Continue : int * bool list * stack -> stack.

We will make a small refactoring: we will separate the “stack frames” from the stack by introducing a type frame and changing the constructor Continue to instead have the type Continue : frame * stack -> stack.

type frame = { n : int; a : bool list }
type state =
  | Start
  | Continue of frame * state
  | Finished

Equipped with this representation of the functions of type unit -> bool list option occurring in the program, we can translate enumerate_assignments to use these constructors instead of using anonymous functions. Anywhere we call an anonymous function, we instead call the function apply that we implement to recover the behaviour of the anonymous function. The parameter next : unit -> bool list option becomes s : stack.

let enumerate_assignments n =
  let state = ref (Start n) in
  (* ^ previously: fun () -> go n [] (fun () -> None) *)
  let rec go n a s =
    if n = 0 then
      (state := s; Some a)
    else
     go (n-1) (true :: a) (Continue ({n; a}, s))
     (* ^ previously: fun () -> go (n-1) (false :: a) next *)
  (* Apply pops a frame from the stack and runs until the next item is produced, if any.
     When `go` runs, it will manipulate the stack, saving it into `state` in particular right
     before returning `Some a`. *)
  and apply s = match s with
    | Start n -> go n [] Finished
    | Continue ({n; a}, s) -> go (n-1) (false :: a) s
    | Finished -> None
  in
  (* And now we have implemented a function `unit -> bool list option` without using any
     higher-order functions internally! *)
  fun () -> apply !state

This program is now completely first-order with the exception of the “higher-order interface” provided at the very end in the form of the function fun () -> apply !state.

This first-order nature will make it possible for us to (somewhat) straightforwardly translate this into C. Of course, we can translate it “as is”, which would mean using a linked list structure for the type stack and for the type bool list, but since we’re going lower-level, we may as well choose more efficient representations for these types too.

Translating to C

The resulting C program is around 100 lines of code whereas the defunctionalized OCaml program is around 20 lines of code. (Both counts ignore blank lines and comments.) We need the following includes in this development.

#include <stdint.h>
#include <stdlib.h>
#include <stdio.h>

Our first choice for efficient data representation will be to represent bool list as simply uint64_t. This limits the number of variables to 64, but it also has the nice property that on a 64-bit machine such as most, a truth assignment fits into a register. We will need to define some operations for setting and clearing specific bits.

typedef uint64_t truth_assignment;
typedef uint8_t var_index;

truth_assignment set_true(truth_assignment a, var_index i) {
    return a | (1UL << i);
}

truth_assignment set_false(truth_assignment a, var_index i) {
    return a & ~(1UL << i);
}

truth_assignment const EMPTY_TRUTH_ASSIGNMENT = 0;

Next, we will translate the type frame from the OCaml implementation into a simple C struct.

struct frame {
    var_index i;
    truth_assignment a;
};

struct frame make_frame(var_index i, truth_assignment a) {
    return (struct frame) {
        .i = i,
        .a = a,
    };
}

Next, recall the stack type from the OCaml implementation.

type stack =
    | Start of int
    | Continue of frame * stack
    | Finished

We will represent the linked list structure as a simple array of struct frames. Notice that the depth of recursion is bounded by the parameter n given to enumerate_assignments. The maximum value of the parameter n in the development here is 64. The maximum recursion depth informs the maximum stack size.

var_index const MAX_VARS = 64;
var_index const STACK_LIMIT = MAX_VARS;

To manage this array of frames, we will need to track a frame pointer: this is the index of the next unused frame in the stack. Seen differently, the frame pointer is the count of frames currently in the stack.

The frame pointer, being at least 8 bits wide, can accommodate values greater than the count of frames we will ever store. This means we can use the upper bits of the frame pointer to help identify what state the generator is in. When the generator is in the start state, the stack is empty, so we can use the value 1 << (WIDTH-1) for the frame pointer to signify that the generator is in the start state.

In the Start state, we need to know the count n of variables we’re enumerating truth assignments for, but afterwards we can forget this n and just keep the array of frames. This suggests using a union to reuse space here.

struct generator {
    uint8_t frame_pointer;
    union {
        var_index num_variables;
        struct frame stack[STACK_LIMIT];
    } data;
};

enum state {
    START = 1,
    CONTINUE = 0,
};

enum state generator_state(struct generator *gen) {
    // extract the uppermost bit of the frame pointer
    return gen->frame_pointer >> 7;
}

// constructs a generator's initial state
struct generator enumerate_assignments(var_index num_variables) {
    return (struct generator) {
        .data.num_variables = num_variables,
        .frame_pointer = 1 << 7,
    };
}

Next, we need operations to push to and pop from the stack held in a generator.

/**
 * Returns -1 if pushing fails: wrong generator state or stack is full.
 * Otherwise copies the given frame into the top of the stack, increments the frame pointer,
 * and returns 1. */
int gen_stack_push(struct generator * gen, struct frame const * frame) {
    if (gen->frame_pointer >= STACK_LIMIT) {
        return -1;
    }
    gen->data.stack[gen->frame_pointer++] = *frame;
    return 1;
}

/**
 * Returns -1 if popping is forbidden: wrong generator state.
 * Returns 0 if the stack is empty.
 * Otherwise decrements the frame pointer, and copies the top frame into `out`,
 * and returns 1 */
int gen_stack_pop(struct generator * gen, struct frame * out) {
    if (START == generator_state(gen)) return -1;
    if (0 == gen->frame_pointer) return 0;

    *out = gen->data.stack[-- gen->frame_pointer];
    return 1;
}

Now that we’ve translated all the type definitions, we can translate the programs go and apply from OCaml into C. Recall the OCaml implementation:

let enumerate_assignments n =
    let state = ref (Start n) in
    let rec go n a s =
        if n = 0 then
            (state := s; Some a)
        else
            go (n-1) (true :: a) (Continue ({n; a}, s))
    and apply s = match s with
        | Start -> go n [] Finished (* Finished comes from fun () -> None *)
        | Continue ({n; a}, s) -> go (n-1) (false :: a) s
        | Finished -> None
    in
    fun () -> apply !state

Notice that go refers to the variable state that is not a parameter of go. In other words, the definition of go constructs a closure. Sadly, C does not have closures, so we will implement this by passing our translation of go a pointer to the generator. This way when go makes a recursive call, it can simply use gen_stack_push to implement the expression Continue ({n; a}, s) at the same time as state := s. In other words, rather than building up the stack in a parameter to save it just before returning, our translation of go will be mutating the stack along the way.

Observe also that go is tail-recursive: the recursive call is the last thing the function does. The OCaml compiler transforms this into a while-loop during compilation. This transformation is called tail-call optimization. We will also perform this transformation in our translation, to avoid using the C call stack.

int go(struct generator * gen, var_index i, truth_assignment * ta) {
    struct frame frame;

    // due to the --> 'operator',
    // `i` will have its value decremented by one inside the loop
    while (i --> 0) {
        *ta = set_true(*ta, n);

        // in the OCaml program, the frame that gets pushed by the recursive call
        // `go (n-1) (true :: a) (Continue ({n; a}, s))`
        // stores the value `n`, but here we are storing n-1 as a consequence
        // of the decrement that happens in the while loop condition.
        frame = = make_frame(n, *ta);
        if(-1 == gen_stack_push(gen, &frame)) return -1;
    }
    return 1;
}

Notice that go returns a status code here, whereas the original OCaml program returned the truth assignment. The C program ‘returns’ the truth assignment via the pointer parameter ta, and rather than construct a new truth assignment, it simply modifies the given one.

Finally, we can translate apply. We will call it next in the C program, since it will dispatch on the current generator state to compute the next item in the sequence, updating the generator state.

int next(struct generator * gen, truth_assignment * ta) {
    int status;
    struct frame frame;

    switch (generator_state(gen)) {
    case CONTINUE:
        status = gen_stack_pop(gen, &frame);

        // handle generator exit
        if (0 == status || -1 == status) return status;

        *ta = set_false(frame.a, frame.i);
        if (-1 == go(gen, frame.i, ta)) return -1;
        return 1;

    case START:
        *ta = EMPTY_TRUTH_ASSIGNMENT;
        gen->frame_pointer = 0;
        if(-1 == go(gen, gen->data.num_variables, ta)) return -1;
        return 1;
    }
}

And what good is all this code if we don’t try it out. Here’s a main function to run the generator until it exits, printing out the truth assignments along the way.

int main() {
    struct generator gen = enumerate_assignments(5);
    int status = 0;
    truth_assignment a;

    for (; status = next(&gen, &a), 1 != status;) {
        printf("truth assignment: %d\n", a);
    }

    if (-1 == status) {
        printf("Generator encountered an error, sorry.\n");
        return EXIT_FAILURE;
    }

    return EXIT_SUCCESS;
}

Collecting all this C code into a file enumerate.c, we can witness the fruits of our labour:

$ gcc enumerate.c && ./a.out
truth assignment: 31
truth assignment: 30
truth assignment: 29
truth assignment: 28
truth assignment: 27
truth assignment: 26
truth assignment: 25
truth assignment: 24
truth assignment: 23
truth assignment: 22
truth assignment: 21
truth assignment: 20
truth assignment: 19
truth assignment: 18
truth assignment: 17
truth assignment: 16
truth assignment: 15
truth assignment: 14
truth assignment: 13
truth assignment: 12
truth assignment: 11
truth assignment: 10
truth assignment: 9
truth assignment: 8
truth assignment: 7
truth assignment: 6
truth assignment: 5
truth assignment: 4
truth assignment: 3
truth assignment: 2
truth assignment: 1
truth assignment: 0

And there you have it: the most complicated program you’ve ever seen to count down from 2^n -1.

Conclusion

This article covered a lot of ground.

At first, we were motivated by Python’s elegant generator syntax and curious about how these ideas are implemented. We implemented the idea of stateful generators using continuation-passing style to give us access to a representation of the evaluation context, so we could store that context in a mutable variable. That gave us the following implementation.

let enumerate_assignments n =
    let rec go n a next =
        if n = 0 then
            (* we wrap the next item in Some *)
            (state := next; Some a)
        else
            go (n-1) (true :: a) (fun () -> go (n-1) (false :: a) next)
    and state = ref (fun () -> go n [] (fun () -> None))
    (* and arrange that the last continuation to be stored in
       `state` just returns `None`. *)
    in
    fun () -> !state ()

In the following section, we explored a purely functional take on this idea, motivated by the simple idea of returning the continuation together with the generated value. To make this typecheck, we needed to introduce the following recursive type.

type 'a l =
  | Done
  | More of 'a * (unit -> 'a l)

This recursive type (or some variant thereof) is often presented in programming languages courses simply as “a lazy list”. The development in this article, on the other hand, derived this representation by eliminating the mutable variable from the program in the previous section. This demonstrates that lazy lists are purely functional generators, where the state of the generator is captured in the continuation of type unit -> 'a l that is explicitly returned.

In the following section, we applied defunctionalization to the stateful CPS generator to convert it into a state machine using an explicit stack.

let enumerate_assignments n =
  let state = ref (Start n) in
  let rec go n a s =
    if n = 0 then
      (state := s; Some a)
    else
     go (n-1) (true :: a) (Continue ({n; a}, s))
  and apply s = match s with
    | Start n -> go n [] Finished
    | Continue ({n; a}, s) -> go (n-1) (false :: a) s
    | Finished -> None
  in
  fun () -> apply !state

In the final section, we translated the defunctionalized program into C, using efficient data representations along the way where possible.

I hope that this sheds some light on the connection between CPS, generators, and state machines.

Once upon a time, some seven years ago, I wrote a post about a Python trick for accessing the first element of any kind of sequence.

The gist of that article is that just using indexing isn’t general enough: if the sequence is lazily computed by a generator, then indexing won’t work. What does work, however, is to iterate over the sequence using a for loop. But since we just want the first element, let’s immediately break out of the loop.

for x in seq:
    print(x) # or do whatever
    break

But there’s something I missed in that post! What if the sequence is empty?

An obvious solution is to use a boolean:

entered = False
for x in seq:
    entered = True
    print(x)
    break
if not entered:
    pass # do whatever

But that’s pretty ugly if not outright disgusting. Instead, we can take advantage of a little-known feature of Python’s for-loops: in Python, a for-loop can have an else block! The semantics are that if the loop terminates normally, then the else-block is run; else, if the loop exits early, e.g. by a break, then the else-block is skipped.

These semantics are frankly weird to me, but they end up working out marvellously for our current use case. It almost makes me think that this is why else was allowed in the first place for a for-loop. Consider this.

for x in seq:
    print(x)
    break
else:
    print("nothing in the sequence")

If the sequence is nonempty, then we will print the first item in the sequence and abort the loop, skipping the else-block. If the sequence is empty, then the loop will exit normally, so the else-block will run!

Originally developed as a technique for implementing compilers for (higher-order) functional languages, defunctionalization (d17n) is a program transformation that eliminates higher-order functions. It is based on the following observation: although there might be infinitely many possible functions one might pass to a higher-order function, there are only finitely many functions that are actually passed in any given program. Therefore, we can define a (finite) data type to represent our choice of function and we can define an interpreter for this data type that recovers the behaviour of the original function.

One situation we might want to employ d17n as programmers is in networked applications – higher-order functions give us lots of expressive power, but we unfortunately can’t send functions over the network! In fact, you might have already done this without realizing it. Rather than send a function, which as more or less impossible, we send some representation of a function that the remote side interprets to execute the function we wanted. The canonical example in web development is the use of a return_to URL parameter: when a user attemps to perform an action while not logged in, we redirect them to a login page with a return_to URL parameter. Once the user logs in, they are redirected to the URL stored in that parameter. That parameter is precisely a defunctionalized continuation.

This article will demonstrate d17n first for the standard higher-order function filter, followed by a discussion of defunctionalized continuations.

Defunctionalizing filter

Let’s see d17n in action with an example. Suppose our program filters lists.

let rec filter p = function
  | [] -> []
  | x :: xs -> if p x then x :: filter p xs else filter p xs

let _ = filter (fun x -> x mod 2 = 0) [1;2;3;4;5;6;7;8;9]

Now let’s see the defunctionalized form of this program.

(* This is the data structure that represents our choice of function. *)
type predicate = IsEven

(* This is the interpreter for this data structure,
   which recovers the behaviour of the function. *)
let apply p_rep x = match p_rep with
  | IsEven -> x mod 2 = 0

(* This is the defunctionalized filter.
   Instead of receiving a function as input, it receives a
   value of type `predicate`, which is a _representation_ of a function. *)
let rec filter_df p_rep = function
  | [] -> []
  | x :: xs ->
    (* We pass the predicate to our interpreter. *)
    if apply p_rep x
    then x :: filter_df p_rep xs
    else filter_df p_rep xs

let _ = filter_df IsEven [1;2;3;4;5;6;7;8;9]

Look, no more higher-order functions!

That was a quite simple example though. Let’s see how to accommodate more complicated functions one by one.

First, what happens when the function we pass to filter is a closure, i.e. it contains variables that are not its parameters?

let divisible_by k l = filter (fun x -> x mod k = 0) l

Notice that the function fun x -> x mod k = 0 refers to k, which is not a parameter of the function. To accommodate this, we will add a constructor to our type predicate called IsDivisible, and that new constructor will crucially have one field to store the value of this k. Then, when our interpreter apply matches on the predicate it can recover the value of k and use it to recover the behaviour of the original function.

type predicate = IsEven | IsDivisible of int

let apply p_rep x = match p_rep with
  | IsEven -> x mod 2 = 0
  | IsDivisible k -> x mod k = 0

(* The implementation of `filter` itself is unchanged. *)

let divisible_by k l = filter_df (IsDivisible k) l

Next, what happens when we combine multiple predicates into one? For example, we might want to filter a list to select all elements that satisfy two properties. Actually, we can define this as a separate combinator, which takes two functions as input and produces a function as output.

let both f g = fun x -> f x && g x

Notice that the type of f and of g is also the type of both f g, namely 'a -> bool. That’s precisely the higher-order type that we’re eliminating via d17n. This will end up making our representation type predicate into a recursive type.

type predicate =
  | IsEven
  | IsDivisible of int
  | Both of predicate * predicate

(* And correspondingly, our implementation of `apply` will be recursive too. *)
let rec apply p_rep x = match p_rep with
  | IsEven -> x mod 2 = 0
  | IsDivisible k -> x mod k = 0
  | Both (p1, p2) -> apply p1 x && apply p2 x

Now the astute reader might have noticed that our type predicate is actually quite limiting: it only works for one type and in this case that’s int. The crux of the issue is that the predicate IsDivisible and IsEven only work for int, whereas the predicate Both (p1, p2) should work for any x : 'a provided that both p1 and p2 are predicates that work on 'a. More to the point, what if we had a predicate IsPalidrome that should work on on string? How would we define a well-typed apply? It would need to be able to accept either a string or an int, requiring that the given predicate be “for a string” or “for an int”.

It is possible to further generalize the type predicate by making it into a generalized algebraic datatype (GADT). With a GADT, it becomes possible to implement a well-typed apply, but a demonstration of this will need to wait for a future article.

Defunctionalizing continuations

One amazing use of higher-order functions is a technique called continuation-passing style (CPS). In this style of programming, rather than write a function that returns its result normally, we instead write a function that returns its result by passing it to another function. The upshot is that every call in a CPS program is a tail-call. In the presence of tail-call optimization, such programs do not use the call stack as function calls are implemented (more or less) as simple jump instructions. For the functional programmer, this means we have goto in OCaml!

A few months ago, I wrote an article explaining CPS in some detail and showing a mindblowing use of it to implement a backtracking search through a tree. I recommend reading that article before the rest of this one.

That article presents a simple definition of boolean formulas and an evaluator for such formulas:

type formula =
  | Conj of formula * formula
  | Disj of formula * formula
  | Neg of formula
  | Var of string

type name = string
type env = (name * bool) list

let rec eval (r : env) = function
  | Var x -> List.assoc x r
  | Neg e -> not (eval r e)
  | Conj (e1, e2) -> eval r e1 && eval r e2
  | Disj (e1, e2) -> eval r e1 || eval r e2

The challenge is this: devise a way to find for a given phi : formula a satisfying assignment to its variables, i.e. a value r : env such that eval r phi = true, in “one pass”. Of course, it can’t actually be in one pass because this is an NP-complete problem; that’s why I put quotes around “one pass”. What I mean by “one pass” is that we want to solve this formula without separately enumerating all the possible truth assignments of the variables and evaluating them.

I found a way to do this using CPS. The core idea of the below algorithm is that we want to save the current execution state whenever we encounter a variable we haven’t seen before. These occur in the leaves of the expression tree. If we find at the end of evaluating the tree that the result is false, then we successively jump back to those saved states to try a different value for the corresponding variable.

let rec solve (r : env) (fail : unit -> 'r) (phi : formula)
    (assign : env -> bool -> (unit -> 'r) -> 'r) : 'r =
  match phi with
  | Var x -> begin match List.assoc_opt x r with
    | Some b -> (* we already have a value for `x` *)
      assign r b fail
    | None -> (* we don't have a value for `x` *)
      assign ((x, true) :: r) true @@ fun () ->
      assign ((x, false) :: r) false fail
    end
  | Neg e ->
    solve r fail e @@ fun r b fail -> assign r (not b) fail
  | Conj (e1, e2) ->
    solve r fail e1 @@ fun r b1 fail ->
    if b1 then (* short-circuiting *)
      solve r fail e2 @@ fun r b2 fail -> assign r (b1 && b2) fail
    else
      assign r false fail
  | Disj (e1, e2) ->
    solve r fail e1 @@ fun r b1 fail ->
    if b1 then
      assign r true fail
    else
      solve r fail e2 @@ fun r b2 fail ->
      assign r (b1 || b2) fail

(* From the CPS solver, we recover a function `formula -> env option` that decides
   whether the given formula is satisfiable, giving the satisfying assignment
   in that case.

   To do so we provide the initial continuations for `fail` and `assign`:
   - The initial `fail` continuation is reached if we have exhausted every
     assignment to the variables. The formula is therefore unsatisfiable.
   - The initial `assign` continuation is reached when we finish evaluating
     the whole expression, having worked out an assignment `r : env` under
     which the expression's value is `res : bool`.
*)
let solve_enter (phi : formula) : env option =
  solve [] (fun () -> None) phi (fun r res fail -> if res then Some r else fail ())

What’s challenging and mysterious about this implementation is that the assign continuation, which represents a normal return from solve, takes in addition to an env and a bool a function of type unit -> 'r. This function represents an abnormal return, like throwing an exception. Therefore, we can undo a return from solve by having the continuation we pass as assign call its given failure continuation.

The most interesting and important part of the algorithm is in the None subcase of the Var case. In that case, we have encountered a variable for which we don’t have an assigned value in the environment r : env. So we call assign with an extended environment and an augmented failure continuation. It is exactly here that we express the idea “try to return true here but if that fails then return false.”

At first it seems terrifying to try to defunctionalize this, but fortunately, d17n is a completely systematic, mechanical process.

1. For every type of function that is passed, we define a new datatype. We define one interpreter for each of these types, although at this point all we can do is work out the signature of the interpreter, as we haven’t decided what the constructors of these new types will be.
2. Find every position where we pass an anonymous function. Each of those functions becomes a constructor of the datatype corresponding to the function’s type. Then, we have to identify for each anonymous function all the variables it contains that aren’t its parameters – the fancy math-name for those is free variables. The types of the free variables become the fields of the constructor corresponding to the function.
3. Find every place where we call a function received through a parameter; these will become calls to the apply interpreter we will write. Replace every anonymous function with its corresponding constructor generated by the above process.
4. Implement apply to recover the original functions

We’ll follow these steps one by one in modifying the implementation of solve.

Step 1: type definitions

There are two higher-order types we seek to eliminate in the original program: we define a type failure to represent the function unit -> 'r and a type assign to represent the function env -> bool -> (unit -> 'r) -> 'r.

There is sadly a small wrinkle arising from the fact that without using a GADT, we can’t express polymorphism in our defunctionalized program. Therefore, we will have to decide what 'r is going to be. We observe that the initial continuations return env option, so that is what we fix 'r to be. The interpreters we define will be apply_failure : failure -> env option and apply_assign : assign -> env -> bool -> failure -> env option.

type failure = ...
type assign = ...

let apply_failure (sf : failure) : env option = failwith "todo"
let apply_assign (sa : assign) (r : env) (b : bool) (sf : failure) : env option =
  failwith "todo"

Step 2: constructor definitions

Here’s the code for the CPS solver, this time annotated to identify all the anonymous functions that are passed as arguments. Annotations F-n indicate an anonymous function passed as a fail continuation whereas A-n indicate a function passed as an assign continuation.

let rec solve (r : env) (fail : unit -> 'r) (phi : formula)
    (assign : env -> bool -> (unit -> 'r) -> 'r) : 'r =
  match phi with
  | Var x -> begin match List.assoc_opt x r with
    | Some b -> assign r b fail
    | None ->
      assign ((x, true) :: r) true @@ fun () ->
      assign ((x, false) :: r) false fail (* F-1 *)
    end
  | Neg e ->
    solve r fail e @@ fun r b fail -> (* A-1 *)
    assign r (not b) fail
  | Conj (e1, e2) ->
    solve r fail e1 @@ fun r b1 fail -> (* A-2 *)
    if b1 then
      solve r fail e2 @@ fun r b2 fail -> (* A-3 *)
      assign r (b1 && b2) fail
    else
      assign r false fail
  | Disj (e1, e2) ->
    solve r fail e1 @@ fun r b1 fail -> (* A-4 *)
    if b1 then
      assign r true fail
    else
      solve r fail e2 @@ fun r b2 fail -> (* A-5 *)
      assign r (b1 || b2) fail

let solve_enter (phi : formula) : env option =
  solve []
    (fun () -> None) (* F-2 *)
    phi
    (fun r res fail -> if res then Some r else fail ()) (* A-6 *)

Since only two functions are passed as the fail continuation, let’s start there. The initial continuation has no free variables, so its constructor will just be Failed with no fields. On the other hand, the function F-1 has the free variables assign, x : name, r : env, and fail. This leads to the following definition.

type assign = ...
type failure =
  | Failed
  | Retry of assign * name * env * failure

And we can refactor this by observing that this pair of constructors give rise to a list structure:

type failure_frame = assign * name * env
type failure = failure_frame list

Next, we apply the same reasoning to the functions passed as assign. We can immediately observe that a list structure will again arise as each function passed as assign refers to the outer assign as a free variable, except for the initial assign continuation which has no free variables. I will annotate each constructor field with the name of the free variable that corresponds to that field.

type failure = failure_frame list
and failure_frame = assign * name * env

and assign = assign_frame list
  (* and the [] case of the list corresponds to A-6 *)
and assign_frame =
  | Neg1 (* A-1 *)
  | Conj1 (* A-2 *) of formula (* e2 *)
  | Conj2 (* A-3 *) of bool (* b1 *)
  | Disj1 (* A-4 *) of formula (* e2 *)
  | Disj2 (* A-5 *) of bool (* b1 *)

Due to the list structure, we should keep in mind that when our interpreter examines a Disj1, for example, there will be a sublist of assign_frames as well. This sublist corresponds to the assign free variable present in the original function. Calling assign from within an augmented assign continuation in the original program will be translated into a call to apply_assign on the sublist of assign_frames.

Step 3: Replace unknown functions with apply and anonymous functions with constructors

In this step, we change the implementation of solve and solve_enter. By “unknown function”, I mean a function received through a parameter.

let rec solve (r : env) (fail : failure) (phi : formula) (assign : assign) : 'r =
  match phi with
  | Var x -> begin match List.assoc_opt x r with
    | Some b ->
      apply_assign assign r b fail
    | None ->
      (* Here we used to both call assign and pass it an anonymous function. *)
      apply_assign assign ((x, true) :: r) true @@ (assign, x, r) :: fail
      (* previously: fun () -> assign ((x, false) :: r) false fail *)
    end
  | Neg e ->
    solve r fail e (Neg1 :: assign)
    (* previously: fun r b fail -> assign r (not b) fail *)
  | Conj (e1, e2) ->
    solve r fail e1 (Conj1 e2 :: assign)
    (* previously: fun r b1 fail ->
    if b1 then (* short-circuiting *)
      solve r fail e2 @@ fun r b2 fail -> assign r (b1 && b2) fail
    else
      assign r false fail *)
  | Disj (e1, e2) ->
    solve r fail e1 (Disj1 e2 :: assign)
    (* previously: fun r b1 fail ->
    if b1 then
      assign r true fail
    else
      solve r fail e2 @@ fun r b2 fail ->
      assign r (b1 || b2) fail *)

let solve_enter (phi : formula) : env option =
  solve [] (* the initial environment *)
    [] (* previously: fun () -> None *)
    phi
    [] (* previously: fun r res fail -> if res then Some r else fail () *)

Step 4: Implement apply

This step is straightforward. We write recursive functions apply_failure and apply_assign to process the lists of type failure and assign. Notice that in the ‘previously’ comments from the above code, some of the continuations we replaced with constructors will need to call solve. This means that apply_failure, apply_assign and solve will all need to be mutually recursive.

To implement each case of apply, we just have to take the code from those ‘previously’ comments: we continue translating each anonymous function into a constructor and each call to an unknown function into a call to the corresponding apply.

let rec solve r fail phi assign = ... (* as above *)

and apply_failure (sf : failure) : env option = match sf with
  | [] -> None
  | (sa, x, r) :: sf -> apply_assign sa ((x, false) :: r) false sf

and apply_assign (assign : assign) (r : env) (b : bool) (fail : failure) : env option =
  match assign with
  | [] -> if b then Some r else apply_failure fail
  | a :: assign -> match a with
    | Neg -> apply_assign assign r (not b) fail
    | Conj1 e2 ->
      if b then solve r fail e2 (Conj2 b) else apply_assign assign r false fail
    | Conj2 b1 ->
      apply_assign assign r (b1 && b) fail
    | Disj1 e2 ->
      if b1 then apply_assign assign r true fail else solve r fail e2 (Disj2 b1)
    | Disj2 b1 ->
      apply_assign assign r (b1 && b) fail

Conclusion

This is certainly a strange way of programming. Whereas in the original implementation with higher-order continuations all the code was in pretty much one tight function, now we have the code spread across three functions! Let’s take a step back and think about what we’ve done.

At first, our CPS program expressed the logic of what to do after the recursive call by placing that logic within a continuation. Since each new continuation refers to the previous one, these form a linked list, although this structure is not immediately obvious. At positions where the solver wants to return a value, it instead invokes the continuation, e.g. assign r (not b) fail. The continuation assign contains all the pending operations to do ‘on the way back’ of the recursion. In its defunctionalized form, the solver expresses the logic of what to do next by pushing onto a literal stack some kind of token e.g. Conj1 e2. When it wants to return, it invokes apply_assign passing it the stack of tokens representing the remaining work to do. Then, apply_assign dispatches on the stack to perform the logic that we used to express in the anonymous function.

Ultimately, what we have done is rewrite the higher-order CPS code into a first-order state machine. Calls to functions such as solve, apply_failure, and apply_success represent a state transition and the collection of arguments given to such functions is the state. These functions examine the current state to figure out what to do: apply_assign, for instance, examines the call stack (represented as the type assign) to decide whether we finished evaluating the formula.

Transforming our code in this way can serve as a guide for implementation in a lower-level language. We can fairly straightforwardly translate this code into C: since every call in the defunctionalized program is to a statically-known first-order function, we won’t even need to use function pointers in implementing the C code. The resulting C program won’t even need to be recursive: we can code it as one big while loop by effectively performing tail-call optimization ourselves. Seeing how to translate the final code we arrived at in this article might be the topic of a future post.

Naturally, defunctionalization has an inverse called refunctionalization (r17n). The idea of r17n really is the opposite: we identify each position where we branch on a first-order data structure and replace the data structure with a function that we instead call. Remarkably, this concept is actually a well-known refactoring in the OOP world; we call it “replace conditional with polymorphism” in that context since (object) polymorphism is effectively how one gets higher-order functions in an OOP language. I have also written an article previously (and unknowingly!) about r17n in JavaScript. Specifically, the very first example discussed in that article shows how to replace a first-order implementation of a ‘thunk’ with a higher-order implementation that has some garbage-collection related benefits.

Finally, the takeaway is that d17n and r17n are fascinating refactorings that apply across a wide range of programming languages. Using d17n, we can prototype something in a higher-order way and systematically lower it to a first-order implementation. The benefits are multiple: we can reimplement in a lower-level language for performance, or we can transfer defunctionalized continuations across nodes in a networked application. Using r17n we can convert natural first-order code into higher-order code. The benefits in that case are less obvious, but I have yet to see so far a refunctionalized program shorter than its corresponding first-order implementation.

(This article was originally drafted in February 2022. The topic is very much related to the previous article’s, Implementing environment-based evaluation of recursive functions in OCaml, but they can very much be read independently.)

Whew, that’s a title that takes some unpacking! Asynchronous recursion is a concept in JavaScript, and presumably in other languages with some form of async-await. Maybe another way to call it would be “indirect recursion”. A picture is worth a thousand words, so let me paint a picture with some code. Let’s count down from a given number until we reach zero, pausing for a second at each recursive call.

function countdown(n) {
    if (n <= 0) {
      console.log(
          "No more bottles of beer on the wall, no more bottles of beer!",
      );
      return;
    }
    console.log(`${n} bottles of beer on the wall, ${n} bottles of beer! ...`);
    setTimeout(() => countdown(n - 1), 1000);
}

This is what I mean by “asynchronous recursion” or “indirect recursion”. Rather than making a recursive call as a statement of the main body of countdown, the recursive call is made in a callback function to an asynchronous operation – in this case, a timeout.

This pattern of recursion can be converted to a kind of continuation-passing style (CPS). A JavaScript programmer is probably already intimately familiar with this style of programming. For example in NodeJS, most of the standard library works this way. Here’s an example where we delete a file:

const fs = require('fs');

fs.unlink('/tmp/hello', (err) => {
  if (err) { console.log('yikes'); return; }
  console.log('deleted /tmp/hello');
});

What’s special about this is that the code to execute after the delete has taken place is represented as a function passed as the second parameter of unlink. The call to unlink returns immediately so other work can be done concurrently. When the deletion finishes, the NodeJS runtime invokes the callback we passed, and we see either “yikes” or “deleted /tmp/hello”.

From the internal implementation of unlink’s point of view, it has been passed a function that it can consider a “return function”. When unlink’s IO operation is finished, it would like to return back into our code, but because all of this happened asynchronously, there isn’t a stack frame in our code that we can simply return to! Instead, it calls the “return function” it was passed.

This concept of “return function” is an instance of taking a baked-in notion of control flow – in this case, returning from a function call – and reflecting it into our code as an explicit function we can call.

To refactor the example of asynchronous recursion I showed earlier, we can apply this same idea. Let’s take the control-flow idea of “making a recursive call” and reflect it into our code as an explicit function, by rewriting countdown to take an extra parameter which I’ll unimaginatively call recurse:

function countdown(recurse, n) {
    if (n <= 0) {
      console.log(
          "No more bottles of beer on the wall, no more bottles of beer!",
      );
      return;
    }
    console.log(`${n} bottles of beer on the wall, ${n} bottles of beer! ...`);

    recurse(n - 1);
}

Notice that I also got rid of the setTimeout. The big idea here is that countdown doesn’t care what kind of recursion the caller wants. The caller can choose what to pass as recurse and get either synchronous or asynchronous recursion as desired! The accomplishment here is also a separation of concerns: we decoupled the behaviour for a single iteration of a recursive loop from the “loop behaviour”.

The change we made to countdown comes at a cost, however. What exactly are we supposed to put as a value for recurse when we call countdown??

To address this, we will model each kind of recursive behaviour as a separate, higher-order function to which we can pass countdown. The result of applying such a recursion combinator to countdown should be a function that takes all the “real” parameters of countdown, i.e. n. Let’s figure out how to implement a “standard” recursion first, then we’ll move on to asynchronous recursion.

const recursively = (f) => (...args) => f(recursively(f), ...args);

To see how this works, let’s evaluate recursively(countdown). There’s only one step to do: substitute countdown for f and we arrive at (...args) => countdown(recursively(countdown), ...args). The result is that when we call this function with a number such as 100, we in fact end up calling countdown passing recursively(countdown) itself as the argument for the recurse parameter. The process then continues recursively, we might say.

Finally, to make this asynchronous, we’ll need to construct a new recurse function that’s more complicated than just f. But only a bit more complicated. It will simply need to call itself within a call to setTimeout.

const delayedRecursively = (f, delay) =>
    (...args) => f(
        (...args) => setTimeout(
            () => f(delayedRecursively(f, delay), ...args),
            delay,
        ),
        ...args
    );

Let’s convince ourselves that this works by evaluating delayedRecursively(countdown, 100). Substituting, we get

(...args) => countdown((...args) =>
    setTimeout(() =>
        countdown(delayedRecursively(countdown, 100), ...args), 100), ...args)

If we imagine for a moment that JavaScript allows partial application, then we can substitute further and get the following:

(n) => {
    /* ... the countdown implementation ... */
    setTimeout(() => countdown(delayedRecursively(countdown, 100), n-1), 100)
}

And if we ran this as delayedRecursively(countdown, 500)(10) we would see the messages printed out slowly.

This approach works! We were able to get different recursive behaviours out of the same implementation of countdown, provided it recurse through an auxiliary function rather than directly.

There is however a glaring issue with this approach: there is no way for the recursive call to meaningfully return a value to the caller! Sure, recursively simply returns whatever f returns, so one could simply write const result = recurse(...); but what about when we use delayedRecursively? We would then get whatever setTimeout returns! To address this, we will need a uniform way to return a value, that works whether the recursion is synchronous or asynchronous.

Returning via yet another function

The trick, as always, is to introduce yet another layer of indirection.

As a motivating example, let’s consider a recursive algorithm that sums the integer values contained in a binary tree.

const sumTree = (t) => {
    if (t === null) return 0;
    else return sumTree(t.left) + t.value + sumTree(t.right);
}

First and as before, we make the recursion indirect via a recurse function.

const sumTree = (recurse, t) => {
    if (t === null) return 0;
    else return recurse(t.left) + t.value + recurse(t.right);
}

Next, we observe that we have a problem if we use delayedRecursively: the return value of recurse(t.left), for example, won’t be the sum of the left subtree’s elements! Let’s introduce another function parameter called resolve this time. Rather than returning via the return statement, our function will instead return by calling resolve.

const sumTree = (recurse, t, resolve) => {
    if (t === null) resolve(0);
    else
        recurse(t.left, (n1) =>
            recurse(t.right, (n2) =>
                resolve(n1 + t.value + n2)));
}

What’s nice about this approach is that it works without us needing to modify any of our recursion combinators from the previous section. We can evaluate delayedRecursively(sumTree, 100)(sampleTree, (n) => console.log(n)) and it (slowly) calculates the sum of the tree and prints out the sum.

Conclusion

In this article, we saw how to separate two concerns that at first glance seem inextricably tried: the pattern of recursion was isolated from the recursive algorithm itself. We introduced some combinators, recursively and delayedRecursively, to each represent a different pattern of recursion. Then, we rewrote our recursive algorithm to recurse via a function which we named recurse. That refactored version of the algorithm expresses the base and step cases of the recursive algorithm without explicitly performing the recursion, leaving it up to the recursion combinator to decide how exactly that will be done. Finally, to account for a desire to return values from our functions, we introduced one more layer of indirection by returning via a function which we named resolve. No changes to the recursion combinators were necessary to accommodate this.

The choice of resolve for the name of this function is no accident. The code recurse(t.left, (n1) => recurse(t.right, (n2) => resolve(n1 + t.value + n2))) is honestly gross. It’s callback hell. There is certainly a way of writing some slightly different recursion combinators that take advantage of JavaScript’s promises and async/await syntax. I leave it to the interested reader to work this out.