I am tackling the Advent of Code challenges in Haskell. In particular, Advent of Code #10 was fun (spoilers clicking on that link, but no spoilers in this blog post itself). Part two required using memoization to be solved, and I had already used memoization in other programming languages but not in Haskell. So I learned the hows and the whys of memoization in Haskell – thus the reason why this article exists.

In order to learn about memoization in Haskell, I did a quick Google and it led me the Memoization article on the Haskell Wiki. That article is good, but it only explains the how, not the why. I will cover both how and why in this blog post.

Motivation

The usual example to start with is the following:

slow_fib :: Int -> Integer
slow_fib 0 = 0
slow_fib 1 = 1
slow_fib n = slow_fib (n-2) + slow_fib (n-1)

Why is this a slow implementation? It’s easy to see if you try to calculate slow_fib 5 by hand:

slow_fib 5 = slow_fib 3 + slow_fib 4
slow_fib 4 = slow_fib 2 + slow_fib 3
slow_fib 3 = slow_fib 1 + slow_fib 2
slow_fib 2 = slow_fib 0 + slow_fib 1
slow_fib 1 = 1
slow_fib 0 = 0

Note how slow_fib 5 calls slow_fib 3 and slow_fib 4 but both of those call slow_fib 2 – so there are some computations that are happening multiple times which slows down the whole process.

When trying to rewrite a recursive function like that to its memoized version in Haskell, the first step is to make the function a bit more general. That is:

slow_fib' :: (Int -> Integer) -> Int -> Integer
slow_fib' f 0 = 0
slow_fib' f 1 = 1
slow_fib' f n = f (n-2) + f (n-1)

Why is this more general? Because we can pass something else other than slow_fib' as f and it will compute something different. We can consume this function as follows:

Main> let f = (\x -> slow_fib' f x) in f 10
55

There is a common abstraction for let f (\x -> slow_fib' f x) – a neater way to make the same call is using fix, where fix f = f (fix f):

Main> :t fix
fix :: (a -> a) -> a
Main> let f = fix slow_fib' in f 10
55

At this point, we have a more generalized slow_fib but it’s still as slow as the previous implementation – try running both on the input 100. This is the part where the magic happens.

Evaluation model

Before we start implementing the memoized version of slow_fib, we will spend some time explaining how memoization works in Haskell by inspecting the evaluation model.

Sharing

Consider the following piece of code:

f x + f y + f x

In this example, there will be one evaluation for f x, one for f y, and another evaluation for f x. So, f x was evaluated twice. However, there’s a trick we can use. Consider the next code:

let x = map (+1) [0..] in (x !! 0, x !! 0)

Executing it means we are mapping the function (+1) on the infinite list, and this happens lazily – it will only get evaluated when we need it. Now, when we actually retrieve the first element x !! 0, it is like computing 0 + 1. Note that we are doing that twice in the tuple. However, Haskell will only evaluate 0 + 1 once, even though it sees “two” 0 + 1.

The reason for this is that in general, when Haskell sees let f = ... f ..., i.e. f referencing itself in the body, it will share the value of f instead of recalculating it. The rationale is: If you name it, you can share it. This allows Haskell to skip the recomputation of values – avoid duplicate computation.

Inlining

You can test this; If you create a file named test.hs with the following contents, you will be able to see when the evaluation happens:

import Debug.Trace

test = let x = trace "forced" (map (+1) [0..]) in (x !! 0, x !! 0)

It will produce:

Main> test
(forced
1,1)

However, if you run this expression directly in the REPL you will get different results:

Main> let x = trace "forced" (map (+1) [0..]) in (x !! 0, x !! 0)
(forced
1,forced
1)

The reason for that is that GHC optimized the expression by inlining it – it replaced (x !! 0, x !! 0) with ("forced" `trace` 1, "forced" `trace` 1). The primary optimization mechanism in GHC is inlining, inlining and inlining.

However, note that this problem occurs just for this specific example, and memoization is not affected by it. The reason for that is GHC will not inline recursive things.

Monomorphism and Polymorphism

For some reason, GHC thought that it should inline this expression. We can disable that behavior as shown in the following example:

Main> :set -XMonomorphismRestriction
Main> let x = trace "forced" (map (+1) [0..]) in (x !! 0, x !! 0)
(forced
1,1)
Main>

You might be wondering, why on earth did this happen? We start to dig a little into the internals of Haskell but I will keep it short.

A function with a polymorphic type is t a. A function with a monomorphic type is [Int] (or List Int) – so monomorphic is the opposite of polymorphic, it’s like a particular instance of a polymorphic.

If we consider the expression x = 4 with a type of Num a => a then it requires re-computation every time it will be used. The reason for that is Haskell can’t be certain of what x really is, because x :: Int is different from x :: Double.

This will not happen with monomorphic types since their values are static and can be cached, so that’s why -XMonomorphismRestriction fixes the problem. I don’t know about GHC internals much, but maybe it could use a map of (Type, Value) to cache values of polymorphic types? Shrug.

Memoization

From the Haskell docs, the following memoize function is shown:

memoize :: (Int -> a) -> (Int -> a)
memoize f = (map f [0 ..] !!) -- not using eta-reduction may affect performance

Note how this function shares values, so it will avoid recomputation. Now, finally, we have:

Main> fix (memoize . slow_fib') 10
55
Main> fix (memoize . slow_fib') 100
573147844013817084101

Works pretty fast on the two example inputs.

Putting it all together

We learned a lot about the Haskell evaluation model and how it allows us to cache/memoize computation. We also learned how complex the compiler is and how some of the assumptions it makes defeats our intentions (e.g. inlining is intended for optimization but it produced wrong results w.r.t. our expectations).

I would like to thank dminuoso and monochrom from #haskell@freenode for explaining the evaluation model to me. Thanks to u/rifasaurous for pointing at a typo and for hinting how complex the Haskell compiler is.