Using the Free Monad to Avoid Stack Overflows in Scala

One of the joys of Scala is actually being able to code in a much more functional style than you can in Java. But inevitably as you start down that route you start using the machinery of advanced functional programming: monads, monad transformers and their ilk.

So it’s a bit of a shock when what looks like some perfectly innocent code stack overflows in your face.

A simple example

(This is similar to the situation where I first observed the problem: imagine you want to manage access to a connection with Reader, so you map over a list of “things to do”, giving you a list of Readers, which you sequence and then feed in the connection. In this case the “thing to do” is trivial: just return what you got. And then I add them up, just so you can add this to your list of “stupid ways to get the length of a list”.)

Okay, that’s pretty bad, what’s going on here? The problem is that we’re building up a big Reader object (and a Reader is just a function, remember), which we only actually execute at the end. And the way sequence works means we end up building a huge stack of functions (proportional to the length of the list!), and so when we run them we stack overflow.

Uh oh

This is bad. We really don’t want to have to break out of monadic style arbitrarily where we’re worried that we might end up processing a lot of data. And we just generally don’t want our programs to die depending on the length of the data involved!

(Yes, we’d be okay if the Scala compiler had better tail-call optimization, but at the moment it only optimizes self-recursive calls, which isn’t quite what’s going on here.)

A solution?

I actually hit this problem at work, and I popped into the #scalaz IRC room to seek help. Naturally, what I got was a somewhat cryptic comment from Tony Morris that “the Free monad” was what I needed, but this was enough to put me on the trail. And yes, it does solve our problem, but as ever, it’s no fun if we don’t delve into the theory for a bit. So humour me for a while (or you can skip ahead to the bit where I tell you how to fix the problem).

Category theory time

Yep, as ever, the terminology leads us smack bang into category theory territory. I’m going to assume a bit more knowledge this time, if you’re not familiar with the basics up to functors (and how they relate to Scala), then you might want to brush up now. I explain the basics in a previous post as well.

Free functors, loosely, are functors from one category to another that add “as little structure as possible”. For example, the free monoid functor, $F_{Mon}: \mathbf{Set} \rightarrow \mathbf{Mon}$, is the functor that takes a set of elements to the “list monoid” on that set. We all know that list is a monoid:

• Singleton sets are the basic elements.
• The monoid operation is appending sets.
• The zero of the monoid operation is the empty set.

But in a sense list is “the most basic” monoid. You can describe stuff that happens in any monoid in terms of a list of the elements involved and the “real” monoid operation: you just fold the operation over the list! A list just, well, lists the elements that were put together but “doesn’t decide” which operation to use to put them together. And that’s the sense in which lists are the free monoids over sets.

Freedom and Forgetfulness

We can pin down what a free functor has to do a bit more precisely, however. Free functors have a relationship with another kind of functor: forgetful functors. The forgetful functor is really boring: it just takes an object and “forgets” some of its structure. So the forgetful functor $U_{Mon}: \mathbf{Mon} \rightarrow \mathbf{Set}$ just takes a monoid and gives you the set of its elements. It “forgets” that there was an operation on them.

The final concept we need is that of two functors being “adjoint”. This is a bit of a complicated one, but given $F: \mathcal{C} \rightarrow \mathcal{D}$ and $G: \mathcal{D} \rightarrow \mathcal{C}$, we say that $F$ is left-adjoint to $G$ (or vice versa with “right adjoint”) if there is a bijection $Hom(X, FY) \cong Hom(GX, Y)$ for any X and Y.1

That’s a bit of a lot to take in, so let’s do an example. Free functors are left adjoint to forgetful functors! And we actually demonstrated (most of) that bijection earlier: given a set X, and a monoid Y, if there’s a function from X to the set of elements of Y (i.e. $g: X \rightarrow U_{Mon}Y$, then we can make a function $f: Free_{Mon}X \rightarrow Y$ by doing (in pseudocode) $\lambda x \rightarrow x.map(g).fold(mzero_{Y})(mappend_{Y})$.

That’s just like what we described before - going from lists to any monoid by folding, it’s just that this time we don’t assume that the list monoid and the other monoid have to have the same basic elements: just so long as we’ve got a mapping between them, we’re fine.2

Damn straight.

This isn’t actually scary at all, it’s just like saying that we have a type-constructor that takes a type parameter that’s itself a type-constructor. We actually do this all the time, but the way we write it in Scala is to write our type parameter as F[_], indicating that it itself has a “hole” for a type parameter. In this case, we’re going to have two type parameters as we still want to parametrize over our “base” type. So in Scala we’d have something like:

In the case of our free monad functor, we’re going to claim that, given a functor F, it’s going to give us a monad in the remaining parameter A.

Okay, so what should a free monad look like? It’s actually not too difficult to puzzle out. Suppose we have a functor F and we want to make a monad just from that functor. Well, let’s start by just using our functor. As a first bash:

We’ve got a functor, might as well use it.

Flatmap is a bit harder. In fact, let’s just think about flatten/join (remember, you just need to define one or the other to have a monad). How on earth are we going to turn a Free[S, Free[S, A]] into a Free[S, A]? We’ve just got nothing to work with.

Okay, so we’re going to need to define our Free functor a bit differently. We can do whatever we like, though, so what would make it work? Well, how about if a Free[S, Free[S, A]] just was a Free[S, A]?

Let’s run with this for a bit. We’re going to define our own Free monad3, imaginatively named Freee (to avoid name collisions with Scalaz’s Free), as the union of several cases, so we do that with a sealed abstract trait and some case classes that extend it. Also, in order to do some stuff with partially applying type constructors, I’m going to start using the super-ugly “type-lambda” trick.

So we’ve got a case for just putting a thing in there straight (Return), and then we can have things that are examples of our functor wrapping something else. Which might itself be just a wrapped value (a Return), or it might be another thing wrapped in a functor (a Roll). This is a bit weird, since it’s kind of recursive, but it solves our flattening problem: we don’t flatten, we just declare that we’ve got an instance of Roll and we’re done!

Also, this isn’t actually that odd. Consider the recursive definition of a list (this isn’t exactly how it’s done in scala, but near enough):

This is actually very similar to our definition of Freee! Which isn’t super surprising, given that List is a free functor itself.4

Okay, let’s go ahead and actually define the Monad instance for Freee:

Hopefully that all makes some sense: we’re really just propagating the maps and joins down our list of functors until we hit a Return and we can do the obvious thing.

I’ve thrown in a couple of other useful functions in there as well: liftFree is like taking x to [x], and foldFree is, well, like the folding operation I described. Since the free monad doesn’t actually do anything, just makes a big list of all the functors you’ve stacked up, foldFree takes a way to to “get out” of your functor and then just works it’s way down the list. (In the scalaz version, this function is called “go”).

Cool stuff

Free monads are pretty cool, because of the properties that come from them being an adjunction, particuarly the “folding” property. That means that you can define a functor that represents stuff that you want to do, wrap it in a free monad, and end up with functions that you can use in monadic style and which produce a kind of “program”. That is, it just lists all the things that you’ve put together, but you can then “interpret” the program by folding over it with an interpreter that takes your functor objects and does something with them.

Again, this is all pretty reasonable when we think about our definition and the similarity to lists: the free monad doesn’t actually do anything, it just builds up a “list” of all the functors you’ve stuck on, and then you have to come along and provide it with a way to actually flatten all those functors – maybe doing something along the way. And that’s what foldFree does for us.

This is really powerful, and it’s particularly cool because you can easily stick on different interpreters, say adding on the possibility of executing your “program” in a concurrent fashion.5

This is all well and good, but how does it help us with our stack problems? Well, the answer is that it doesn’t directly, but in scalaz the implementation of the free monad has two useful properties:

1. It is a proper implementation of the free monad, so we can do all the stuff we’ve been talking about with it.
2. It implements flatMap by returning continuation objects, thus effectively transferring the execution stack instead to the heap, and cleverly ensuring that the evaluation remains tail-recursive.

These are really two distinct things, but they go nicely together precisely because the free monad provides so little structure. This makes it easy to bolt on when you really want property 2!

The idea for this comes from Runar’s excellent paper Stackless Scala with Free Monads. It’s well worth a read: he generalises the idea of a “trampoline” and shows that you end up with… the free monad! A trampoline is essentially just a free monad where the “suspension functor” is a Function0, i.e. a no-arguments continuation. But if you generalise this to allow other suspension functors, you get the free monad. Plus, with a little messing around you can do the whole thing without making any non-tail-recusive calls. Awesome!

(I’m not going to go into the details of how this works, as Runar’s paper does a better job than I would (and this post is already far too long!). There are a few clever tricks, but if you’ve made it through this post, you should have all the conceptual machinery you need.)

The scalaz implementation, then, is a bit different from the one we derived above. In particular, there are a few naming differences:

• Roll becomes Suspend
• foldFree becomes go

The rest is mostly machinery around making property 2 work, but most of it gets explained in Runar’s paper.

So, let’s refactor our original problem case to use the scalaz’s Free.

There are a few things going on here.

1. We’re stacking monads on top of each other, so we’re going to need to start using monad transformers. If you’re not familiar with them, they’re not too scary: they’re really just monads wrapping other monads with all the plumbing handled for you.6
2. We’re not interested in using the interpreter-like properties of the free monad, we’re pretty much just using it for property 2 above, so we’re just going to us Id as our “suspension” functor.
3. We have to do a little bit of a dance to get in and out of Free: we have to point the output of our Reader functions into FreeId and then (trivially) fold the whole thing down at the end using go.

And that’s it! You can make the list as long as you want now, and you won’t hit stack overflow errors (well, unless it’s so long you OOM, but that’s not really fair). This does mean you’ll take a performance/memory hit, as you’re trading stack space for heap space, which both means that it’s not as optimised for what we’re doing, and, you know, you use up heap space. From a code point of view, the main cost you bear is having to switch over to using monad transformers, even if you only had one in play to begin with. This can be alleviated with some sensible type aliases, but it does make things conceptually more complicated. There is also the cost of having to get in and out of Free, but this should be fairly localised and can be hidden with some sensible methods.

Boom. No more stack overflows.

1. Technically, these bijections have to be “natural”, which really means “behave sensibly”, so don’t worry about that.

2. Also, the adjunction requires a bijection, and I’ve only shown one direction. The reverse is left as an exercise for the reader.

3. By “our own”, I mean blatantly stolen from the Haskell implementation, with the main source for my plagiarism being this excellent SO answer.

4. In fact, we can write List using the Free monad as follows:

type List2[A] = Freee[({type f[+x] = (A, x)})#f, Unit]

This is actually a bit different from the usual list, as join will act as concatenation, instead of flattening.

5. If you’re interested in this kind of thing, this is an excellent discussion of free monads in the context of interpreters.

6. One subtlety is that they often appear to be “inside out” – that is, a FooT[BarT[Baz]] is sometimes really a Foo wrapped in a Bar wrapped in a Baz, which is the opposite to how it appears.