Polymorphic Recursion

Re-reading for the n-th time Chris Okasaki's From Fast Exponentiation to Square Matrices: An Adventure in Types, I tried my hand at seeing how does the latest OCaml version cope with non-regular types and the consequent polymorphic recursion. I first scoped the terrain and Googled what others had to say about this topic. I found a ten-year old OCaml Mailing List post where Pierre Weiss mentions work by Pessaux and Leroy that ended in the inclusion of support for recursive types. So I fired OCaml with the not-always-useful option -rectypes.

That post gives uses as a running example the type of non-regular lists, or functional vectors, which Okasaki calls "random-access lists" (Okasaki, Chris. Purely Functional Data Structures, p. 143 and ss.). The type in question encodes in constructors the length of the vector as a binary integer, with a recursive type argument:

type 'a vec =
  Nil
| Zero of ('a * 'a) vec
| One  of 'a * ('a * 'a) vec

(I depart from the names given in the above reference). This permits indexing in logarithmic, as opposed to linear, time. For instance, the vector [1; 2; 3; 4; 5] is represented as:

One (1, Zero (One (((2, 3), (4, 5)), Nil)))

The problem is with the function cons:

(* Doesn't work *)
let rec cons x = function
  Nil         -> One (x, Nil)
| Zero    ps  -> One (x, ps)
| One (y, ps) -> Zero (cons (x, y) ps)

Consing an element to a One (x, xs) list must "carry over" through a recursive call to cons whose argument is not α but α×α. The typechecker unifies both types in the recursive type α×α as α, whose only value is:

let rec omega = (omega, omega)

Such that cons omega Nil works, but the more useful cons 1 Nil doesn't. The correct type declaration for cons looks like:

val cons : 'a -> 'a vec -> 'a vec

and such a declaration would be required in Haskell, since polymorphic recursion is undecidable. It is, however, not quite right: the parameter is not universally quantified, since that would mean that the recursion must have the same type as the argument. The type variable should be, instead, existentially quantified.

Existential types are also undecidable and so need explicit declarations. OCaml has two ways to introduce existential types: abstract types in modules, and records. The second alternative is more lightweight:

type cons = { cons : 'a . 'a -> 'a vec -> 'a vec }

(I abuse the different namespaces in OCaml punning types and values). The parameter α is existential since it is "to the right" of a type arrow. Now I can write cons through such a record:

let cons =
  let rec f = { cons = fun x l -> match l with
    Nil         -> One (x, Nil)
  | Zero    ps  -> One (x, ps)
  | One (y, ps) -> Zero (f.cons (x, y) ps)
  } in f.cons

Voilà! Now I can write:

# cons 1 (cons 2 (cons 3 (cons 4 (cons 5 Nil)))) ;;
- : int vec = One (1, Zero (One (((2, 3), (4, 5)), Nil)))

As another example, the length of such a vector can be found in logarithmic time:

type length = { length : 'a . int -> 'a vec -> int }

let length v =
  let rec f = { length = fun n l -> match l with
    Nil         -> n
  | Zero    ps  -> f.length (    2 * n) ps
  | One (_, ps) -> f.length (1 + 2 * n) ps
  } in f.length 0 v

For a more challenging, and useful, example, indexing on a vector must take into account the position and the length of the current vector. I follow directly Okasaki's implementation (op. cit. p. 145):

type index = { index : 'a . int -> 'a vec -> 'a }

let index n v =
  let rec f = { index = fun n l -> match l with
    Nil         -> failwith "index"
  | One (x, ps) -> if n = 0 then x else f.index (n - 1) (Zero ps)
  | Zero    ps  ->
    let (l, r) = f.index (n / 2) ps in
    if n mod 2 = 0 then l else r
  } in f.index n v

Update: Thank you to SpiceGuid for pointing out the error that prevented index from typechecking. Also, be sure to read bluestorm's comment for alternatives to achieve the same effect.