Rust has always felt like a strange beast, culturally speaking. The community is made of a mix of people with very different perspectives, including anything from hardcore low-level kernel hackers to category-theorist and functional programming gurus. This is also what makes this community so fertile: whether you’re coming from C, Haskell or TypeScript, you’re likely to learn a lot from other perspectives.

I’d like to add my modest contribution by introducing a pattern coming from the functional programming world, recursion schemes¹. Recursion schemes are a design pattern for representing and traversing recursive data structures (typically trees) which help factor the common part of recursive traversals, making transformations nicer to write, to read and to compose.

Even in the functional programming world, recursion schemes are not so well-known. Like monads, they are usually presented in Haskell with frightening words like zygohistomorphic prepromorphisms. It’s a pity because recursion schemes can be both simple, useful and practical. I’d even argue that in Rust, the most interesting part is perhaps the representation technique, more than the traversal, despite the latter being the original and the usual motivation for using recursion schemes.

In this post, we’ll work through a concrete example to introduce recursion schemes and what they can do. We’ll point to a more real life example of how we use them in the implementation of the Nickel configuration language, and we’ll discuss the pros and cons of using recursion schemes in the particular context of Rust.

(In)flexible representations

Let’s say you’re writing a JSON parser library. You’ll need to expose a type representing JSON values. For the sake of argument, let’s assume that you support an extension of the JSON language with pairs, so you can write {"foo": ("hello","world")}. Here’s a natural representation:

pub enum JsonValue {
  String(String),
  Number(f64),
  Pair(Box<JsonValue>, Box<JsonValue>),
  Array(Vec<JsonValue>),
  Object(HashMap<String, JsonValue>),
}

This data structure is recursive: JSON values can contain other JSON values. We thus have to use Box (or any other indirection) around recursive occurrences of JsonValue. Otherwise, this enum would have an infinite size (excepted for Array and Object since Vec and HashMap add their own indirection, but it’s somehow luck).

Now, user requestor asks that your parser adds location information to the output, because they validate some user-provided configuration and would like to point to specific items on error. This is a reasonable request which is sadly very hard to satisfy in the serde ecosystem. Anyway, our parser isn’t interfacing with serde, so we can add span information:

pub type Span = std::ops::Range<usize>;

pub struct Spanned<T> {
  pos: Span,
  data: T,
}

pub type SpannedValue = Spanned<JsonValue>;

pub enum JsonValue {
  String(String),
  Number(f64),
  Pair(Box<SpannedValue>, Box<SpannedValue>),
  Array(Vec<SpannedValue>),
  Object(HashMap<String, SpannedValue>),
}

You can go different ways about this. We could have added a second argument to each constructor of the enum, such as in String(String, Span), to avoid the additional Spanned layer, but that would be a lot of repetition. We could also have moved Box to data: Box<T>. Still, the general idea is that we now have two layers:

• a struct layer gathering the JSON data and the span together;
• the original enum layer, the core of JSON, which is almost unchanged.

So far, so good. But user conservator is now complaining that you’ve spoiled their performance. They’re using JSON as a machine exchange format and don’t care about position information. Could you restore the old representation and a way to produce it, ignoring spans?

Unfortunately, we had to change JsonValue. Copy-pasting the original JsonValue enum under a different name is possible, but it’s unsatisfying, as we now have multiple copies to maintain. It also doesn’t scale. Beside adding position information, you might want to have a value representation that uses Rc instead of Box, because you’re going to need to keep reference to arbitrary nodes during some complex transformation.

The functorial representation

The recursion schemes pattern has two components: a representation technique and a transformation technique. I believe the representation part is particularly interesting for Rust, so let’s start with that.

We’ll try to make our JSON representation more generic to accommodate for the different variations that we mentioned in the previous section. The fundamental idea is to replace the recursive occurrences of JsonValue within itself, Box<JsonValue> (or JsonValue for Array and Object), by a generic parameter T. Doing so, we’re defining just one layer of a JSON tree where recursive children can be anything, not necessarily JSON values (we use the F suffix for that generic version because it’s technically a functor, but that doesn’t really matter).

pub enum JsonValueF<T> {
  String(String),
  Number(f64),
  Pair(T, T),
  Array(Vec<T>),
  Object(HashMap<String, T>),
}

Let’s play a with a few examples to get familiar with this representation.

• If we set T = (), we get a type that is isomorphic (modulo some ()) to:

  JsonValueF<()> ~ enum {
    String(String),
    Number(f64),
    Pair,
    Array,
    Object,
  }

  This is precisely a single node of a JSON tree, that is either a leaf or a marker of a node with children but without actually including them.

• If we set T = Box<JsonValueF<T>>, we get back the original JsonValue. But wait, you can’t define the generic parameter T to be something which depends on T itself! In fact we can, but we need to introduce an extra indirection:

  pub struct JsonValue {data: JsonValueF<Box<JsonValue>>}

  The price to pay is an additional struct layer, so you need to match on value.data, and wrap new values as JsonValue { data: JsonValueF::Number(0) }. Note that this layer doesn’t have any cost at run-time.

  Another difference is that we now box the values in Array and Object, which isn’t needed. For now I’ll just ignore that, but you could take a second generic parameter U to represent the occurrences of T that don’t need an indirection if this really matters to you.

• If we extend our intermediate layer a bit, we can get SpannedValue!

  pub struct SpannedJsonValue {
    data: JsonValueF<Box<SpannedJsonValue>>,
    span: Span,
  }

  You can create any extension of JsonValue with additional metadata lying at each node of the tree, which is pretty neat.

• We are also able to change the ownership model of JSON values. It’s simple to write a reference-counted variant:

  pub struct SharedJsonValue {data: JsonValueF<Rc<SharedJsonValue>>}

  Or a borrowed version, that you could allocate in an arena:

  pub struct ArenaJsonValue<'a> {data: JsonValueF<&'a ArenaJsonValue>}

This idea of putting a self-referential type within JsonValueF is referred to as tying the knot. The power of this approach is that you can keep the core JsonValueF type unchanged. This applies to any tree-like recursive structure.

Some methods can be implemented only once on JsonValueF for any T, say is_string or is_number. With additional trait constraints on T, we can write more involved functions, still operating on the generic functor representation.

Let’s now see how to traverse our JSON values.

Traversals

The strong point of recursion schemes is to provide an interface for traversing recursive structures that let you focus on what the function actually does, which is otherwise mixed with how the recursion is done. The idea is to use generic combinators which factor out the plumbing of recursive traversals.

Let’s count the number of String nodes in a JSON value, the naive way.

fn count_strings(value: &JsonValue) -> u32 {
    match &value.data {
        JsonValueF::String(_) => 1,
        JsonValueF::Number(_) => 0,
        JsonValueF::Pair(fst, snd) => count_strings(fst) + count_strings(snd),
        JsonValueF::Array(array) => array.iter().map(|elt| count_strings(elt)).sum(),
        JsonValueF::Object(object) => object.values().map(|elt| count_strings(elt)).sum(),
    }
}

We’ll see how to write this function in the style of recursion schemes. First, we need to define one core combinator: map.

map takes a JsonValueF<T>, a function f from T to U and returns a JsonValue<U>. That is, map takes a JSON layer where all the direct children (the recursive occurrences in our full type) are of some type T and applies f to transform them to something of type U. This is the secret sauce for defining traversals.

impl<T> JsonValueF<T> {
    fn map<U>(self, f: impl FnMut(T) -> U) -> JsonValueF<U> {
        match self {
            JsonValueF::String(s) => JsonValueF::String(s),
            JsonValueF::Number(n) => JsonValueF::Number(n),
            JsonValueF::Pair(fst, snd) => JsonValueF::Pair(f(fst), f(snd)),
            JsonValueF::Array(array) => {
                JsonValueF::Array(array.into_iter().map(|elt| f(elt)).collect())
            }
            JsonValueF::Object(object) => {
                JsonValueF::Object(object.into_iter().map(|(k, v)| (k, f(v))).collect())
            }
        }
    }
}

map isn’t specific to JsonValueF. It can be defined mechanically for any functor representation (e.g. through a macro) of a data structure.

Note that there’s no recursion in sight: there can’t be, because T and U are entirely generic and could very well be (), but we saw that JsonValueF<()> is a single node. map only operates at the current layer.

The trick is that f can use map itself. Let’s see how to use it for count_strings:

fn count_strings(value: JsonValue) -> u32 {
    match value.data.map(|child| count_strings(*child)) {
        JsonValueF::String(_) => 1,
        JsonValueF::Number(_) => 0,
        JsonValueF::Pair(fst, snd) => fst + snd,
        JsonValueF::Array(array) => array.iter().sum(),
        JsonValueF::Object(object) => object.values().sum(),
    }
}

If you look closely, there’s no more recursion in the body of the pattern matching. It’s factored out in the map call. Let’s break down this example:

1. map, given a function from T to U, promises you that it can transform the direct children of type T in JsonValueF<T> to U, providing JsonValueF<U>. We use it immediately with a recursive call to count_strings, which can indeed transform the direct children from a Box<JsonValue> to a u32. If the children have children itself, count_strings will do that recursively as its first action, down to the leaves.
2. Once we’ve reduced potential children of deeper layers to u32s, we get a JsonValueF<u32>. We sum its content at the current layer.

There is a catch though: our count_strings function takes an owned argument, which consumes the original JSON value. I’ll come back to that later.

While I find the second version of count_strings a little cleaner, the difference between the two isn’t really astonishing.

As a more compelling example, let’s define a generic bottom-up traversal function on JsonValue. This traversal is able to map — that is to rewrite — nodes (more exactly entire subtrees). map_bottom_up takes a generic transformation f and applies this function to every subtree starting from the leaves. You could use such a function to apply program transformations or optimizations on an abstract syntax tree.

impl JsonValue {
    pub fn map_bottom_up(self: JsonValue, f: impl FnMut(JsonValue) -> JsonValue) -> JsonValue {
        let data = self.data.map(|v| Box::new(v.map_bottom_up(f)));
        f(JsonValue { data })
    }
}

This example is quite remarkable: it’s almost a one-liner and there is no pattern matching at all! Once again, the structural recursion is entirely factored out in the map function. We implemented map_bottom_up on JsonValue directly, but with some trait constraints on T, we can write a more generic version JsonValueF that works on both the Boxed and Rced version (the arena one is more tricky as it requires an explicit allocator). This example is only scratching the surface.

Mapping is just one example: another common traversals are folds (known as catamorphisms in the recursion schemes jargon), which generalize the well-known Iterator::fold from sequences to trees. In fact, count_strings would make more sense as a fold, but we’ll leave that for another time.

Are recursion schemes useful in Rust?

Haskell has a number of features that make recursion schemes particularly nice to use and to compose, not the least of which is garbage collection. You don’t have to think about ownership; it’s references all the way down. Recursive data structures are easy to express.

On the other side, there is Rust, which culturally doesn’t like recursive functions that much, for good and bad reasons². Though sometimes recursion is hard to avoid, especially on tree-like data structures.

An important issue is that our count_strings consumes its argument, which is unacceptable in practice. It is possible to write a version of map that takes a value by reference, and thus similarly for count_strings, but it’s not entirely straightforward nor free. You can find a by-reference version and more explanations in our associated repository. At any rate, you can always write specific traversals manually without resorting to the recursion schemes way if needed. It’s not an all or nothing approach.

In fact, even if you don’t use map at all, the functor representation alone is quite useful.

How we use recursion schemes in Nickel

In the implementation of the Nickel configuration language, we use the functor representation for the abstract syntax tree of a static type. Here are the stages we went through:

1. In the parser and most of the Nickel pipeline, we used to have a simple Box-based, owned representation, akin to JsonValue.

2. However, during type inference, the Nickel typechecker needs to handle new type constructions, in particular unification variables. Those are as-of-yet unknown types, similar to unknowns in an algebraic equation. Extending the base representation is readily done as for SpannedJsonValue:

   pub enum UnifType {
     Concrete(Box<TypeF<UnifType>>),
     /// A unification variable.
     UnifVar(VarId),
     //.. rigid type variables, etc.
   }

3. More recently, we’ve split the historical, all-powerful unique representation of expressions (including Nickel types) into two intermediate ones. The new initial representation is arena-allocated, which makes it natural to use bare references as the recursive indirection instead of allocating in the heap through e.g. Box. This is easy with recursion schemes: that is precisely the ArenaJsonValue example. For a smooth transition, we need to temporarily keep the old Box-ed Type representation in parts of the codebase, but having different representations co-exist is a basic feature of recursion schemes.

We use map-based traversal typically to substitute type variables (that is, a Nickel generic type, as our T in Rust) for a concrete type and similar rewriting operations. We have variants of the core map function that can also thread mutable state, raise errors, or both. Traversal by reference are implemented manually, with a plain recursive function.

On the downside, type and core function definitions can be a bit verbose and tricky to get right. For example, Nickel’s TypeF has sub-components that themselves contain types leading to 4 generic parameters. There are multiple possibilities for Box placement in particular, only some of them are correct and they are subtly different. Though once you’ve defined a new variant, this complexity is mostly hidden from the consumers of your API. It can still manifest as terrible Rust type errors sometimes if, God forbid, you’ve put a Box at the wrong place.

Conclusion

We’ve introduced recursion schemes, a design pattern for representing and traversing recursive data structures. While the traversal part isn’t as good a fit as in purer functional languages like Haskell, it can still be useful in Rust. The representation part is particularly relevant, making it easy to define variations on a recursive data structure with different ownership models or metadata. We’ve shown how we use recursion schemes in Nickel, and while there are performance and complexity trade-offs to consider, they can bring value for moderately complex tree types that need to be extended and transformed in various ways.