Memoization in OCaml?

Question

It is possible to improve "raw" Fibonacci recursive procedure

Fib[n_] := If[n < 2, n, Fib[n - 1] + Fib[n - 2]]

with

Fib[n_] := Fib[n] = If[n < 2, n, Fib[n - 1] + Fib[n - 2]]

in Wolfram Mathematica.

First version will suffer from exponential explosion while second one will not since Mathematica will see repeating function calls in expression and memoize (reuse) them.

Is it possible to do the same in OCaml?

How to improve

let rec fib n = if n<2 then n else fib (n-1) + fib (n-2);;

in the same manner?

Answer 1

The solution provided by rgrinberg can be generalized so that we can memoize any function. I am going to use associative lists instead of hashtables. But it does not really matter, you can easily convert all my examples to use hashtables.

First, here is a function memo which takes another function and returns its memoized version. It is what nlucaroni suggested in one of the comments:

let memo f =
  let m = ref [] in
    fun x ->
      try
        List.assoc x !m
      with
      Not_found ->
        let y = f x in
          m := (x, y) :: !m ;
          y

The function memo f keeps a list m of results computed so far. When asked to compute f x it first checks m to see if f x has been computed already. If yes, it returns the result, otherwise it actually computes f x, stores the result in m, and returns it.

There is a problem with the above memo in case f is recursive. Once memo calls f to compute f x, any recursive calls made by f will not be intercepted by memo. To solve this problem we need to do two things:

1. In the definition of such a recursive f we need to substitute recursive calls with calls to a function "to be provided later" (this will be the memoized version of f).

2. In memo f we need to provide f with the promised "function which you should call when you want to make a recursive call".

This leads to the following solution:

let memo_rec f =
  let m = ref [] in
  let rec g x =
    try
      List.assoc x !m
    with
    Not_found ->
      let y = f g x in
        m := (x, y) :: !m ;
        y
  in
    g

To demonstrate how this works, let us memoize the naive Fibonacci function. We need to write it so that it accepts an extra argument, which I will call self. This argument is what the function should use instead of recursively calling itself:

let fib self = function
    0 -> 1
  | 1 -> 1
  | n -> self (n - 1) + self (n - 2)

Now to get the memoized fib, we compute

let fib_memoized = memo_rec fib

You are welcome to try it out to see that fib_memoized 50 returns instantly. (This is not so for memo f where f is the usual naive recursive definition.)