# Transparent memoisation in Haskell with MemoTries

Last month I participated in Advent of Code, along with a bunch of friends and colleagues. Typically I prefer to use Advent of Code as a mechanism to learn a new programming language, but after not finishing the last couple of years, I wanted to make life easier for myself, so I decided to use Haskell and give myself a little break.

Haskell is a language designed to effectively express computations. As such, functions are “pure”: they may have no side effects. Even I/O is handled by returning a sequence of instructions which are then actually enacted by the VM. This is incredibly useful when solving logic puzzles such as Advent of Code; you never have to worry about accidentally modifying something, because you *can’t*.

However, it makes memoisation very difficult.

This article is based on the paper, “Memo functions, polytypically!”, by Spenser Jannsen, and the Haskell project MemoTrie. Almost none of the ideas are mine; if you’re interested in going deeper, I recommend starting with those.

### A brief primer on memoisation

*Memoisation* (or “memoization”, if you’re American) is the act of invisibly recording inputs to functions and caching the outputs, so that a second call to the same function responds instantly with the cached value. This needs to be done without exposing the details of the memoisation outside; it should be a drop-in replacement. This is really helpful in certain algorithms, because it allows us to avoid explicit caching for performance.

Typically, we *memoise* a function by wrapping it in another function that caches the results. It might look like this in Python:

```
def fib(n):
if n == 0:
return 0
elif n == 1:
return 1
else:
return memoisedFib(n - 2) + memoisedFib(n - 1)
def memoise(f):
computed = {}
def memoisedF(x):
if x in computed:
return computed[x]
else:
computed[x] = f(x)
return computed[x]
return memoisedF
memoisedFib = memoise(fib)
```

Here we have a naive, slow, Python function that computes a Fibonacci number recursively. Running `fib(40)`

without memoisation takes a few seconds on my computer, because it recomputes the same value many times. However, with memoisation, it’s instant, as it remembers the result for each input and only computes it once.

We can use this technique to cheaply speed up a function we know we will be calling many times with the same input. Often, there is another solution that doesn’t require this, but sometimes it can be harder to read or require a rewrite; memoisation can be a lot cheaper.

### And why this doesn’t work

However, you can’t do that in Haskell. The code above mutates a dictionary, and mutation is a side effect. In order to perform a mutation, we would need to represent that in the type. So while the function itself would look like this:

```
fib 0 = 0
fib 1 = 1
fib n = fib (n - 2) + fib (n - 1)
```

The memoised version might look like this:

```
memoize :: (Eq a, Ord a) => (a -> State (Map a b) b) -> a -> State (Map a b) b
memoize f x = do
computed <- get
case Map.lookup x computed of
Just result ->
return result
Nothing -> do
result <- f x
put $ Map.insert x result computed
return result
fib :: Integer -> State (Map Integer Integer) Integer
fib 0 = return 0
fib 1 = return 1
fib n = do
a <- memoisedFib (n - 2)
b <- memoisedFib (n - 1)
return $ a + b
memoisedFib = memoise fib
```

Note the change in the type signature. This function needs to carry around state, and so that’s reflected in the signature. This means we can’t simply swap out our old `fib`

for this faster version. If we were to use this, the rest of the code would have to change accordingly to accommodate the new type. And, because it’s recursive and therefore needs to call `memoisedFib`

, even the original implementation needs to be modified to carry around the state.

And we need to call it in a strange way, to pass an initial state in, and then discard it later:

```
λ evalState (fib 40) Map.empty
102334155
```

However, there’s a trick we can use to get the same benefits of memoisation, without having to change the type signature of the function. This relies on laziness.

If you’re not familiar with laziness in Haskell, you’re in luck: I wrote an article just for you. Please go read it. And when you’re done, come back.

### Tries, or general-purpose lazy data structures

A *trie* (confusingly pronounced “tree”, and sometimes referred to as a “prefix tree” or “digital tree”), is a way of representing data in a more compact space.

For example, let’s say we have a set of strings:

```
["apple", "avocado", "banana", "cabbage", "cherry", "chive", "lemon", "lettuce"]
```

We can represent this in a little less space by creating a tree of common prefixes:

```
{
"a" -> {"pple" -> {}, "vocado" -> {}}
"banana" -> {}
"c" -> {
"abbage" -> {}
"h" -> {"erry" -> {}, "ive" -> {}}
}
"le" -> {"mon" -> {}, "ttuce" -> {}}
}
```

This won’t always be more space-efficient, but sometimes it can really help. However, it also has another upside: it allows us to be lazy about each node.

In a Haskell list, the head and tail are individually evaluated on demand; if you never want to see the tail, it’s never computed. This is true for any data structure. For example, we can create a data structure representing an `if`

/`else`

construction, and look them up based on a boolean value.

```
data If a = If a a
select True (If ifTrue _) = ifTrue
select False (If _ ifFalse) = ifFalse
if' condition ifTrue ifFalse = select condition (If ifTrue ifFalse)
-- example:
λ let things = If 3 5
λ select True things
3
λ select False things
5
λ if' True 7 2
7
```

We can prove this is lazy by causing a failure if one of the branches is reached, but not the other:

```
λ let broken = If (error "Whoops.") "Hi!"
λ select False broken
"Hi!"
λ select True broken
"*** Exception: Whoops.
CallStack (from HasCallStack):
error, called at <interactive>:1:1 in interactive:Ghci4
```

The data structure was constructed with an exception, but we only triggered the exception when we evaluated that path of the data structure; if we hadn’t gone there, it would never have been evaluated.

And just like the list of Fibonacci numbers from that laziness article, we can put values in these data structures that are expensive to compute, and only ask for what we need.

### Memoising booleans

If you can express a computation operationally, you can also express it as a data structure. We already saw this for Fibonacci numbers in the last post, and `if`

/`else`

above. In general, a function, `input -> output`

, can be modelled as a data structure that maps inputs to outputs. In the case of the Fibonacci sequence, the most obvious mapping was a list (which could be thought of as a mapping from integers to integers); with `if`

/`else`

, it was a data structure with two values, `a`

and `b`

.

We can do this more generically.

```
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
class HasTrie input where
-- Yes, the name of the data type here is ":->:"; just roll with it
data input :->: output
instance HasTrie Bool where
data Bool :->: output = BoolTrie output output
select :: Bool -> (Bool :->: output) -> output
select True (BoolTrie _ ifTrue) = ifTrue
select False (BoolTrie ifFalse _) = ifFalse
if' :: Bool -> output -> output -> output
if' condition ifTrue ifFalse = select condition (BoolTrie ifTrue ifFalse)
```

Once we have that structure, we can add conversion functions back and forth between functions and tries.

```
class HasTrie input where
infixr 2 :->: -- sets the precedence
data input :->: output
trie :: (input -> output) -> input :->: output
unTrie :: (input :->: output) -> input -> output
instance HasTrie Bool where
data Bool :->: output = BoolTrie output output
trie f = BoolTrie (f False) (f True)
unTrie (BoolTrie ifFalse _) False = ifFalse
unTrie (BoolTrie _ ifTrue) True = ifTrue
```

We can see here that `select`

is just `unTrie`

with the arguments flipped. And so we can define `if'`

as such:

```
if' :: Bool -> output -> output -> output
if' condition ifTrue ifFalse = unTrie (trie f) condition
where
f True = ifTrue
f False = ifFalse
```

This is exactly the same as the standard `if`

structure in Haskell… except it constructs a data structure in the middle. Right now it’s not very useful, but if we flip the arguments, suddenly everything changes.

There is a function in the Haskell standard library, `bool`

, which is effectively a reversed form of `if'`

:

```
bool :: a -> a -> Bool -> a
bool x y p = if p then y else x
```

If we redefine this using our `trie`

and `unTrie`

functions above, a really wonderful property falls out:

```
bool' :: a -> a -> Bool -> a
bool' ifFalse ifTrue = unTrie (trie f)
where
f False = ifFalse
f True = ifTrue
```

It’s very important that `bool'`

does not receive the `condition`

parameter, because we don’t want to evaluate `unTrie (trie f)`

once for each `condition`

argument; we want to evaluate it once for all conditions.

We can use this to construct a function of type `Bool -> output`

:

```
someFibs :: Bool -> Integer
someFibs = bool' (fib 30) (fib 40)
```

And once we’ve done that, we can call it:

```
λ :set +s
λ someFibs False
832040
(2.52 secs, 921,589,872 bytes)
λ someFibs False
832040
(0.01 secs, 70,472 bytes)
```

You see what happened there? Calling `fib 30`

the first time took around 2.5s and a gigabyte of memory. The second time, it was instant, because the result was cached in the `BoolTrie`

.

Our `trie`

and `unTrie`

functions effectively memoised the function, `f`

. This is totally generalisable:

```
memo :: (a -> b) -> a -> b
memo f = unTrie (trie f)
-- or, point-free:
memo :: (a -> b) -> a -> b
memo = unTrie . trie
```

### Everything is a trie

This works for booleans. Can we do this for other values too?

It turns out, yes. Haskell data types are modelled as *sums of products*, and it turns out there’s a mechanical transformation from a data type to the corresponding trie type.

Booleans are defined as a sum type: `data Bool = False | True`

, which becomes a trie with two values. This works for all sum types. For example, `Maybe a`

becomes a trie as follows:

```
instance HasTrie a => HasTrie (Maybe a) where
data Maybe a :->: output = MaybeTrie output (a :->: output)
trie f = MaybeTrie (f Nothing) (trie (f . Just))
unTrie (MaybeTrie ifNothing ifJust) = maybe ifNothing (unTrie ifJust)
```

`MaybeTrie`

is a data structure with a value that corresponds to `Nothing`

, and a trie that corresponds to `Just a`

.

We can transform `Either a b`

to a trie in a similar fashion.

```
instance (HasTrie a, HasTrie b) => HasTrie (Either a b) where
data Either a b :->: output = EitherTrie (a :->: output) (b :->: output)
trie f = EitherTrie (trie (f . Left)) (trie (f . Right))
unTrie (EitherTrie ifLeft ifRight) = either (unTrie ifLeft) (unTrie ifRight)
```

So, we can see that sum types of the form `a | b | c`

become tries of the form `Trie a b c`

.

There’s also a mechanical transformation for products. Let’s take the tuple type:

```
instance (HasTrie a, HasTrie b) => HasTrie (a, b) where
newtype (a, b) :->: output = Tuple2Trie (a :->: b :->: output)
trie f = Tuple2Trie $ trie $ \a -> trie $ \b -> f (a, b)
unTrie (Tuple2Trie f) (a, b) = unTrie (unTrie f a) b
```

Tuples become nested tries, which kind of makes sense: the prefix of `(a, b)`

is `a`

. This holds for all product types, too.

Now that we can handle sum and product types, we can turn a list into a trie. Lists are defined in Haskell something like this:

```
data [a] = [] | a : [a]
```

There’s a lot of operators in there, so think of it like this if it’s easier:

```
data List a = Nil | Cons a (List a)
```

That is, it’s the sum type of `[]`

(or `Nil`

) and the product `a : [a]`

(or `Cons a (List a)`

). So we can transform this like a combination of `Maybe`

and tuples.

```
instance HasTrie a => HasTrie [a] where
data [a] :->: output = ListTrie output (a :->: [a] :->: output)
trie f = ListTrie (f []) (trie $ \x -> trie $ \xs -> f (x : xs))
unTrie (ListTrie ifNil _) [] = ifNil
unTrie (ListTrie _ ifCons) (x : xs) = unTrie (unTrie ifCons x) xs
```

We can check if this works, for example, by memoising the Fibonacci number function *after* we have converted a list of booleans representing little-endian bits to a number. Let’s define some conversion functions:

```
-- these won't work for negative numbers
bits :: Integral a => a -> [Bool]
bits 0 = []
bits n = (n `mod` 2 == 1) : bits (n `div` 2)
unBits :: Integral a => [Bool] -> a
unBits [] = 0
unBits (False : bits) = 2 * unBits bits
unBits (True : bits) = 2 * unBits bits + 1
λ bits 10
[False,True,False,True]
λ bits 99
[True,True,False,False,False,True,True]
λ unBits (bits 99)
99
```

Now we’ve done that, we can create a slow Fibonacci function that works on bits:

```
fibBits :: [Bool] -> [Bool]
fibBits = bits . fib . unBits
```

But the real power here is that `[Bool] -> output`

can be converted to a trie, and therefore we can memoise it. Let’s do that:

```
memoFib :: Integer -> Integer
memoFib = unBits . memo (bits . memoFib' . unBits) . bits
where
memoFib' 0 = 0
memoFib' 1 = 1
memoFib' n = memoFib (n - 2) + memoFib (n - 1)
λ memoFib 30
832040
(0.01 secs, 811,520 bytes)
```

We’ve gone from `fib 30`

taking 2.5s to almost no time at all, just like the Python version, because we recurse over the memoised function. And unlike our first attempt, there is no change to the type signature; we get the caching “for free”. (Of course, if your trie becomes huge, you will pay the price, but let’s pretend that won’t happen.)

This is a little clunky. Fortunately, we have all the building blocks to define a trie over integers; after all, aren’t they just lists of bits? So let’s define it as such:

```
instance HasTrie Integer where
newtype Integer :->: output = IntegerTrie ([Bool] :->: output)
trie f = IntegerTrie $ trie (f . unBits)
unTrie (IntegerTrie bitTrie) = unTrie bitTrie . bits
```

And now we can simplify the definition of `memoFib`

:

```
memoFib :: Integer -> Integer
memoFib = memo memoFib'
where
memoFib' 0 = 0
memoFib' 1 = 1
memoFib' n = memoFib (n - 2) + memoFib (n - 1)
```

### Odds and ends

There’s a library for all that, of course. You don’t have to write it yourself. It’s called MemoTrie, and you can find it on Hackage. That said, it’s pretty easy to implement this if you want to, or if you want slightly different semantics.

I ended up copying the library in parts, just so I could actually understand what was going on. In doing so, I found a couple of cool tricks. The main one I wanted to share with you here is that if you’re going to memoise a function that takes multiple parameters, of the form:

```
f :: a -> b -> c -> d
```

I recommend uncurrying it first, and writing the appropriate tuple trie if it doesn’t already exist:

```
f' :: (a, b, c) -> d
f' = uncurry3 f
```

If you use the `memo3`

function from the package (which recursively applies `memo`

), you may end up discarding tries which could be handy later, because they only exist in the scope defined by the nested function. By converting to a tuple first, you ensure that you keep all data around for as long as you need it.

Aside from this small quirk, it’s a wonderful library, and I hope it serves you well.