How does this short memoization function in the GHC test suite work?

Question

Here is the complete runnable code for the following memoization function:

memo f = g
  where
    fz = f Z
    fs = memo (f . S) 
    g  Z    = fz
    g (S n) = fs n
    -- It is a BAD BUG to inline 'fs' inside g
    -- and that happened in 6.4.1, resulting in exponential behaviour

-- memo f = g (f Z) (memo (f . S))
--        = g (f Z) (g (f (S Z)) (memo (f . S . S)))
--        = g (f Z) (g (f (S Z)) (g (f (S (S Z))) (memo (f . S . S . S))))

fib' :: Nat -> Integer
fib'             =  memo fib
  where
  fib Z          =  0
  fib (S Z)      =  1
  fib (S (S n)) = fib' (S n) + fib' n

I tried to figure it out with expansion of the terms by hand, but that expansion looks just like the slow, unmemoized function. How does it work? And how is the commented out code derived?

Answer 1

This is quite tricky to explain. I'll start with a simpler example.

One has to keep in mind the difference between

\x -> let fz = f 0 in if x==0 then fz else f x
let fz = f 0 in \x -> if x==0 then fz else f x

Both compute the same function. However, the former will always (re)compute f 0 when called with argument 0. Instead, the latter will compute f 0 only the first time it is called with argument 0 -- when that happens, fz is evaluated, and the result stored there forever, so that it can be reused again the next time fz is needed.

This is not too different from

f 0 + f 0
let fz = f 0 in fz + fz

where the latter will call f 0 only once, since the second time fz will be already evaluated.

So, we could achieve a light memoization of f, storing only f 0 as follows:

g = let fz = f 0 in \x -> if x==0 then fz else f x

Equivalently:

g = \x -> if x==0 then fz else f x
   where
   fz = f 0       

Note that here we can not bring \x -> on the left of =, or we lose memoization!

Equivalently:

g = g' 
   where
   fz = f 0       
   g' = \x -> if x==0 then fz else f x

Now we can bring \x -> on the left without any problem.

Equivalently:

g = g' 
   where
   fz = f 0       
   g' x = if x==0 then fz else f x

Equivalently:

g = g' 
   where
   fz = f 0       
   g' 0 = fz
   g' x = f x

Now, this only memoizes f 0 instead of every f n. Indeed, computing g 4 twice will cause f 4 to be computed twice.

To avoid this, we can start making g to work on any function f instead of a fixed one:

g f = g'    -- f is now a parameter
   where
   fz = f 0       
   g' 0 = fz
   g' x = f x

Now, we exploit that:

-- for any f, x
g f x = f x
-- hence, in particular
g (f . succ) (pred x) = (f . succ) (pred x) = f (succ (pred x)) = f x

So, g (f . succ) (pred x) is a complicated way of writing f x. As usual, g memoizes the function at zero. However this is (f . succ) 0 = f 1. In this way, we obtained memoization at 1, instead!

Hence, we can recurse

g f = g'    -- f is now a parameter
   where
   fz = f 0       
   g' 0 = fz
   g' x = g (f . succ) (pred x)

If called with 0, this uses fz to store f 0, memoizing it.

If called with 1, this will invoke g (f . succ) which will allocate another fz for the 1 case. This looks good, but that fz won't last for long, since it will be reallocated every time g' x is called, negating the memoization.

To fix this issue, we use another variable, so that g (f . succ) will be computed only once, at most.

g f = g'    -- f is now a parameter
   where
   fz = f 0       
   fs = g (f . succ)
   g' 0 = fz
   g' x = fs (pred x)

Here, fs is evaluated at most once, and will cause the allocation of another fz for the 1 case. This fz now won't disappear.

Recursively, one can be convinced that now all the values f n are memoized.