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?

Liberec

Answer 1

You pretty much do what the mathematica version does but manually:

let rec fib = 
  let cache = Hashtbl.create 10 in
  begin fun n ->
    try Hashtbl.find cache n
    with Not_found -> begin
      if n < 2 then n
      else 
        let f = fib (n-1) + fib (n-2) in
        Hashtbl.add cache n f; f
    end
  end

Here I choose a hashtable to store already computed results instead of recomputing them. Note that you should still beware of integer overflow since we are using a normal and not a big int.

Bolen

You can abstract the memoization procedure into a higher-order function that is very convenient, let memo_f f = let tbl = Hashtbl.create 100 in (fun x -> if Hashtbl.mem tbl x then Hashtbl.find tbl x else begin let n = f x in Hashtbl.add tbl x n; n end) – Cowboy

@nlucaroni: this won't catch the recursive calls that f makes, so it is pretty much useless. It does not decrease the complexity of naive Fibonacci. – Rimbaud

Answer 2

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.)

Rimbaud

There are two typos at the beginning of your answer: “I am going to use associtive lists insetad (…)” (associtive → associative, insetad → instead). I can’t suggest an edit of less than six chars. – Psychosis

Thanks. Can't you edit my answer? (I seem to be able to edit other people's answers.) – Rimbaud

Oh then just quote some Shakespeare to satisfy the administrative requirement. – Rimbaud

This is great. But it's only for single-argument functions, right? Not for multi-arg curried functions? Yet it seems to work with more arguments. I don't understand why. The alist uses only the first argument. Ah, got it--it's only the function returned by passing in the first argument that's being memoized, not the value of the entire function call with all arguments. – Browband

What? If you memoize f : α → β → γ that will memoize on the first argument, since this example is just f : α → δ where δ = (β → γ). You seem to be stating wrongly that it is memoizing the result of f, i.e., a function of type β → γ, but that's not the case. – Rimbaud

"This argument is what the function should use instead of recursively calling itself" -> fib does not need to be recursive anymore, you can remove 'rec' ;) – Assyria

It does? I dont' see it. – Rimbaud