How to memoize recursive functions?

Question

Consider a recursive function, say the Euclid algorithm defined by:

let rec gcd a b =
  let (q, r) = (a / b, a mod b) in
  if r = 0 then b else gcd b r

(This is a simplified, very brittle definition.) How to memoize such a function? The classical approach of defining a high-order function memoize : ('a -> 'b) -> ('a -> 'b) adding memoization to the function is here useless, because it will only save time on the first call.

I have found details on how to memoize such function in Lisp or Haskell:

• How do I memoize a recursive function in Lisp?
• Memoization with recursion

These suggestions rely on the ability found in Lisp to overwrite the symbol definition of a function or on the “call-by-need” strategy used by Haskell, and are therefore useless in OCaml.

Answer 1

The winning strategy is to define the recursive function to be memoized in a continuation passing style:

let gcd_cont k (a,b) =
  let (q, r) = (a / b, a mod b) in
  if r = 0 then b else k (b,r)

Instead of defining recursively the gcd_cont function, we add an argument, the “continuation” to be called in lieu of recursing. Now we define two higher-order functions, call and memo which operate on functions having a continuation argument. The first function, call is defined as:

let call f =
    let rec g x =
      f g x
    in
    g

It builds a function g which does nothing special but calls f. The second function memo builds a function g implementing memoization:

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

These functions have the following signatures.

val call : (('a -> 'b) -> 'a -> 'b) -> 'a -> 'b = <fun>
val memo : (('a -> 'b) -> 'a -> 'b) -> 'a -> 'b = <fun>

Now we define two versions of the gcd function, the first one without memoization and the second one with memoization:

let gcd_call a b =
  call gcd_cont (a,b)

let gcd_memo a b =
  memo gcd_cont (a,b)

Answer 2

# let memoize f =
    let table = Hashtbl.Poly.create () in
    (fun x ->
      match Hashtbl.find table x with
      | Some y -> y
      | None ->
        let y = f x in
        Hashtbl.add_exn table ~key:x ~data:y;
        y
    );;
val memoize : ('a -> 'b) -> 'a -> 'b = <fun>


# let memo_rec f_norec x =
    let fref = ref (fun _ -> assert false) in
    let f = memoize (fun x -> f_norec !fref x) in
    fref := f;
    f x
  ;;
val memo_rec : (('a -> 'b) -> 'a -> 'b) -> 'a -> 'b = <fun>

You should read the section here: https://realworldocaml.org/v1/en/html/imperative-programming-1.html#memoization-and-dynamic-programming in the book Real World OCaml.

It will help you truly understand how memo is working.