Fun with (Type) Class

I'd like to motivate OCaml's object-oriented features (the "O" in "OCaml") by showing an encoding of (first-order) Haskell type classes. This will require making use of almost everything in OCaml's object buffet: class types, virtual classes, inheritance of class types and of classes, double dispatch and coercion.

Haskell could be said to have the best type system among all languages in wide use at the moment, striking a delicate balance between power and practicality without compromising purity (neither conceptual nor functional). In the words of Oleg Kiselyov, it sports "attractive types". A great feature is the presence of type-indexed values via the type-class mechanism. In a nutshell, type-indexed values are functions from types to values. A classic example is overloaded operators, where for instance the same symbol + denotes addition over the integers, floating-point and rational numbers, depending on the type of its parameters. Haskell provides "type magic" to automatically infer the intended function (the value) based on the types of the arguments to which it is applied, achieving a form of type-safe "ad-hoc" polymorphism.

The ML family of languages doesn't have the machinery to attain such expressive power. There are, however, a number of ways to encode type-indexed families with an explicit parameter encoding the selector type. The oldest, and most complicated technique is that of Danvy (see also this complete example) which is not really equivalent to Haskell's type classes but a sophisticated application of phantom types. Another technique is using an explicit dictionary (Yaron Minsky has a worked example) from which to select the appropriate operation. Normally these dictionaries are records; the labels provide type-safe instantiation of the generic value family. The problem with records is that they are not extensible, so it is not immediate how to encode the equivalent of inheritance of type classes. OCaml's object-oriented sublanguage is especially well-suited for the task, allowing for a natural encoding that shows that the "class" in Haskell's "type classes" is more object-oriented than it seems at first sight.

I'll start with a bit of Haskell's standard prelude, as shown in the Report in a nice graph. I'll focus on the two most common nontrivial classes, Eq and Ord. The first is defined as:

class  Eq a  where
    (==), (/=)  ::  a -> a -> Bool

    x /= y  = not (x == y)
    x == y  = not (x /= y)

I encode a Haskell type class in OCaml with… (surprise!) a class type and associated functions:

class type ['a] eq = object
  method equal     : 'a -> 'a -> bool
  method not_equal : 'a -> 'a -> bool
end

Since the type class has a type parameter a, my class type also has a type parameter 'a. Each member of the type class corresponds to a (named) method and to a regular function taking the class type as a dictionary parameter:

let (===) x y (eq : 'a #eq) = eq#equal x y
and (<=>) x y (eq : 'a #eq) = eq#not_equal x y

(I've opted to not shadow the regular operators in Pervasives). The weird order of the parameters (you'd expect the dictionary to come first) is such that I can write:

    if x === y $ eq_int then …

where ($) is the application operator. Haskell provides default implementations for both members; as the Report notes:

This declaration gives default method declarations for both /= and ==, each being defined in terms of the other. If an instance declaration for Eq defines neither == nor /=, then both will loop.

I find that unsettling, so I'll just define one via a virtual class (not class type):

class virtual ['a] eq_base = object (self)
  method virtual equal : 'a -> 'a -> bool
  method not_equal x y = not (self#equal x y)
end

In this case, equal is left virtual. It is important that the class be left virtual even if equal is defined, so that it cannot be instantiated. This means that it cannot be coerced to the class type eq, since it is illegal to shadow virtual members that were declared non-virtual. Now come the instances of the type class. In Haskell, a very simple instance is that for unit which is built-in (via deriving) but could be declared as:

instance  Eq ()  where
    () == () = true

To encode instances, I use an immediate object bound to an identifier:

let eq_unit : unit eq = object
  inherit [unit] eq_base
  method equal     () () = true
  method not_equal () () = false
end

Note the coercion to the corresponding monomorphic instance of the class type. I inherit from the default implementation in the class base, even if in this case I explicitly implement all methods for efficiency. The same encoding works for booleans too:

let eq_bool : bool eq = object
  inherit [bool] eq_base
  method equal     = (==)
  method not_equal = (!=)
end

OCaml has built-in ad-hoc polymorphic comparisons that I could use directly:

class ['a] eq_poly : ['a] eq = object
  inherit ['a] eq_base
  method equal     = Pervasives.(=)
  method not_equal = Pervasives.(<>)
end

With that I can have type classes for the primitive types:

let eq_int    = (new eq_poly :> int    eq)
and eq_float  = (new eq_poly :> float  eq)
and eq_char   = (new eq_poly :> char   eq)
and eq_string = (new eq_poly :> string eq)

Here I instantiate objects of the eq_poly class coerced to the correct monomorphic instance of eq (recall that in OCaml coercion is always explicit). The type class instance for pairs requires type classes for each of its components. In Haskell, we'd have:

instance  (Eq a, Eq b) => Eq ((,) a b)  where …

(not really, but let me gloss over that). Encoding that in OCaml also requires a witness for each type parameter as explicit arguments to the constructor:

let eq_pair (l : 'a eq) (r : 'b eq) : ('a * 'b) eq = object (self)
  inherit ['a * 'b] eq_base
  method equal (p, q) (p', q') = l#equal p p' && r#equal q q'
end

An inefficiency of this encoding is that it creates a new object with every invocation of the witness function. For instance, in general eq_pair eq_unit eq_bool != eq_pair eq_unit eq_bool and so on for every witness. The same is needed to encode an equality witness for option types:

let eq_option (o : 'a eq) : 'a option eq = object
  inherit ['a option] eq_base
  method equal x y = match x, y with
    None  , None   -> true
  | Some x, Some y -> o#equal x y
  | _              -> false
end

And for lists:

let eq_list (o : 'a eq) : 'a list eq = object
  inherit ['a list] eq_base
  method equal =
    let rec equal x y = match x, y with
    | [], [] -> true
    | _ , []
    | [], _  -> false
    | x :: xs, y :: ys when o#equal x y -> equal xs ys
    | _      -> false
    in equal
end

(I use explicit closed recursion for efficiency). In all these cases, I encoded a derived instance via an explicit witness passed as a parameter. Derived type classes require genuine inheritance of class types, as seen in the encoding of Haskell's Ord type class:

class  (Eq a) => Ord a  where
    compare              :: a -> a -> Ordering
    (<), (<=), (>), (>=) :: a -> a -> Bool
    max, min             :: a -> a -> a

(I omit showing the default implementations). I use the SML ordering type for fidelity:

type ordering = LT | EQ | GT

let of_comparison c = if c < 0 then LT else if c > 0 then GT else EQ
and to_comparison   = function LT -> -1 | EQ -> 0 | GT -> -1

The class type encoding the Ord type class is:

class type ['a] ord = object
  inherit ['a] eq
  method compare       : 'a -> 'a -> ordering
  method less_equal    : 'a -> 'a -> bool
  method less_than     : 'a -> 'a -> bool
  method greater_equal : 'a -> 'a -> bool
  method greater_than  : 'a -> 'a -> bool
  method max           : 'a -> 'a -> 'a
  method min           : 'a -> 'a -> 'a
end

Again, operators take an explicit witness as a dictionary:

let (<==) x y (ord : 'a #ord) = ord#less_equal    x y
and (<<<) x y (ord : 'a #ord) = ord#less_than     x y
and (>==) x y (ord : 'a #ord) = ord#greater_equal x y
and (>>>) x y (ord : 'a #ord) = ord#greater_than  x y

And I provide a virtual class with default implementations:

class virtual ['a] ord_base = object (self)
  inherit ['a] eq_base
  method virtual compare : 'a -> 'a -> ordering
  method equal         x y = match self#compare x y with EQ -> true  | _ -> false
  method less_equal    x y = match self#compare x y with GT -> false | _ -> true
  method less_than     x y = match self#compare x y with LT -> true  | _ -> false
  method greater_equal x y = match self#compare x y with LT -> false | _ -> true
  method greater_than  x y = match self#compare x y with GT -> true  | _ -> false
  method max           x y = match self#compare x y with LT -> y | _ -> x
  method min           x y = match self#compare x y with GT -> y | _ -> x
end

Note how the class type inheritance ord :> eq translates into class inheritance ord_base :> eq_base. The instances for unit and bool use inheritance and selective overriding for efficiency:

let ord_unit : unit ord = object
  inherit [unit] ord_base
  method equal     = eq_unit#equal
  method not_equal = eq_unit#not_equal
  method compare () () = EQ
  method max     () () = ()
  method min     () () = ()
end

let ord_bool : bool ord = object
  inherit [bool] ord_base
  method equal     = eq_bool#equal
  method not_equal = eq_bool#not_equal
  method compare b b' = match b, b' with
  | false, true -> LT
  | true, false -> GT
  | _           -> EQ
  method max = (||)
  method min = (&&)
end

In order to profit from the efficient implementation of equal and not_equal I explicitly delegate on the methods in the corresponding eq instances. Now as before I can use the ad-hoc polymorphic operators in Pervasives to speed writing the type classes for atomic types:

class ['a] ord_poly : ['a] ord = object
  inherit ['a] ord_base
  method compare x y   = of_comparison (Pervasives.compare x y)
  method less_equal    = Pervasives.(<=)
  method less_than     = Pervasives.(< )
  method greater_equal = Pervasives.(>=)
  method greater_than  = Pervasives.(> )
  method max           = Pervasives.max
  method min           = Pervasives.min
end

let ord_int    = (new ord_poly :> int    ord)
and ord_float  = (new ord_poly :> float  ord)
and ord_char   = (new ord_poly :> char   ord)
and ord_string = (new ord_poly :> string ord)

(Note the use of the injection into ordering). The case for pairs requires appropriate witnesses as before:

let ord_pair (l : 'a ord) (r : 'b ord) : ('a * 'b) ord = object
  inherit ['a * 'b] ord_base
  method equal = (eq_pair (l :> 'a eq) (r :> 'b eq))#equal
  method compare (p, q) (p', q') =
    match l#compare p p' with EQ -> r#compare q q' | c -> c
end

In this case, delegating equal requires passing the witnesses to eq_pair explicitly coerced to the eq class type. Finally, the classes for options and lists are straightforward applications of the encoding:

let ord_option (o : 'a ord) : 'a option ord = object
  inherit ['a option] ord_base
  method equal = (eq_option (o :> 'a eq))#equal
  method compare x y = match x with
    None   -> (match y with None -> EQ | Some _ -> LT)
  | Some x -> (match y with None -> GT | Some y -> o#compare x y)
end

let ord_list (o : 'a ord) : 'a list ord = object
  inherit ['a list] ord_base
  method equal = (eq_list (o :> 'a eq))#equal
  method compare =
    let rec compare x y = match x, y with
    | [], [] -> EQ
    | _ , [] -> GT
    | [], _  -> LT
    | x :: xs, y :: ys -> match o#compare x y with
    | EQ -> compare xs ys
    | c  -> c
    in compare
end

As it is clear, the encoding shows a direct correspondence between Haskell type classes and OCaml's classes and objects:

• type classes into class types with top-level functions dispatching on a witness object of that class type
• derived type classes into inheriting class types
• default implementations into virtual classes
• instances into objects inheriting from the default virtual class
• derived instances into objects having witnesses of the appropriate class type

In every case, the witness objects are used purely as closed vtables without state. It would be interesting to investigate which use cases would open using row polymorphism or data members. Also, note that there is no need to abstract over the object types, as each witness carries with it the intended implementation in a type-safe manner.

This technique is analogous to the one using functors to encode type classes, but then why have both in the same language?, in the words of Jacques Garrigue. One reason is that functors are not first-class values in OCaml but objects are, which makes using them as type witnesses much more lightweight. On the other hand, one limitation that functors don't have is that this encoding is limited to first-order type classes. In other words, it is not clear what would be required to implement Haskell type classes like:

instance  (Monad m) => Functor m  where
    fmap f x = bind x (return . f)

where m is a type constructor with sort * → * using the encoding presented here, whereas it is straightforward to concoct a (ML) functor:

module MonadFunctor (M : MONAD) : FUNCTOR = struct
  type 'a t = 'a M.t
  let fmap f x = M.bind x (fun v -> M.return (f v))
end

Another limitation is the need to pass a witness object to the type-indexed functions. One would ideally want to write something like:

let (===) (x : eq) = x#equal

where the value is its own witness, but obviously this would require that every value be wrapped into an object, requiring, for instance, code like the following:

if (new eq_int 5) === (new eq_int x) then …

Moreover, in order to actually operate on the wrapped values, the classes would need an explicit accessor. OCaml supports double dispatch making possible forcing the types of x and y to be equal:

class type ['a] eq = object ('self)
  method value     : 'a
  method equal     : 'self -> bool
  method not_equal : 'self -> bool
end

class eq_int x : [int] eq = object (self : 'self)
  method value                    = x
  method equal     (that : 'self) = self#value == that#value
  method not_equal (that : 'self) = self#value != that#value
end

Canonicalizing the witnesses for finite types like bool is very tedious:

let eq_bool : bool -> bool eq =
  let _t : bool eq = object
    method value          = true
    method equal     that =      that#value
    method not_equal that = not (that#value)
  end and _f : bool eq = object
    method value          = false
    method equal     that = not (that#value)
    method not_equal that =      that#value
  end in fun b -> if b then _t else _f

(essentially, a Church-encoded boolean). More importantly, this encoding is much more dynamic, making it rather remote with respect to the original, fully static Haskell semantics (this, incidentally, is a definite advantage of the functorial encoding).

In closing, it would be interesting to see how far this technique can be carried, especially to implement a numeric tower analogous to Haskell's, or maybe a better type class hierarchy to encode would be that of the Numeric Prelude. Another necessary avenue of investigation is measuring the abstraction penalty incurred by this encoding. In any case, I believe this is a neat showcase of the often-forgotten object-oriented features in OCaml.

(Dis)Functional Bowling

It is saddening to see people invested so heavily in cargo-cult programming practices that they cannot see natural solutions to simple problems after having their thoughts twisted by object-orientation as an ends and TDD as a fighting technique. Uncle Bob is struggling (and failing) to approach functional programming in the natural way by what seems to be his biases and commercial interests. His Bowling Game Kata (warning, PPT; right-click to download) is to find the score of a 10-frame bowling game:

The game consists of 10 frames as shown above. In each frame the player has two opportunities to knock down 10 pins. The score for the frame is the total number of pins knocked down, plus bonuses for strikes and spares.

A spare is when the player knocks down all 10 pins in two tries. The bonus for that frame is the number of pins knocked down by the next roll. So in frame 3 above, the score is 10 (the total number knocked down) plus a bonus of 5 (the number of pins knocked down on the next roll.)

A strike is when the player knocks down all 10 pins on his first try. The bonus for that frame is the value of the next two balls rolled.

In the tenth frame a player who rolls a spare or strike is allowed to roll the extra balls to complete the frame. However no more than three balls can be rolled in tenth frame.

In order to sidestep the problem of validating games for length, I'll consider a bowling game as an infinite number of balls, a finite prefix of which is non-zero:

let repeat x = let rec l = x :: l in l

let game strikes = strikes @ repeat 0

Scoring a frame according to the rules is a matter of pattern matching on the first few balls. According to the result, a complete frame is removed and the rest of the game returned:

let score_frame = function
| 10 :: (n :: m :: _ as g) -> 10 + n + m, g
| n :: m :: (r :: _ as g) when n + m = 10 -> 10 + r, g
| n :: m :: g -> n + m, g

To score a game, I keep track of how many frames I'm into the game, and stop at the tenth:

let rec score frame g =
 if frame = 10 then 0 else
 let n, g = score_frame g in
 n + score (succ frame) g

The game begins at the zeroth-frame, obviously:

# score 0 (game [1; 4; 4; 5; 6; 4; 5; 5; 10; 0; 1; 7; 3; 6; 4; 10; 2; 8; 6]) ;;
- : int = 133

10 minutes of thinking, and some looking around to realize that I had mistyped the scorecard. What TDD do you need, Uncle Bob?

Monadic Golf

Reading Bonsai Code's solutions to Programming Praxis's weekly puzzles (which I can't recommend highly enough) makes me feel acutely aware of how verbose OCaml is, and how inadequate its standard library when compared to Haskell's. However, I've found that the latest puzzles yield concisely to a monadic style over the lists.

Breaking with my usual literate top-down presentation I'll concentrate on the code required to solve the puzzles and leave the obvious scaffolding to the reader. I'll still putt over the par, especially if scored against Remco's brutally terse solutions, but I hope that what is missing is straightforward to fill in. I've written about point-free style and its limitations in the face of OCaml's value restriction. In this case I'll use monadic style for the solutions as a way to show that procedural expression too has its place in functional programming.

This week's problems are drawn from the International Mathematical Olympiads and are very much in the spirit of Project Euler's problems, yielding nicely to brute-force search. The first:

Determine all three-digit numbers N having the property that N is divisible by 11, and N/11 is equal to the sum of the squares of the digits of N.

can be solved simply by:

let imo_1960_01 =
  range 1 99 >>= fun i ->
  let n = 11 * i in
  guard (sum % fmap square % digits $ n = i) >>
  return n

In this solution and the ones that follow I express the problem in monadic terms via fmap, return, bind (and the two synonyms >>= and >>) and guard. Here % is composition, and range, sum, square and digits are obvious. Equally pythy is the solution to the second problem:

Find the smallest natural number n which has the following properties:

1. Its decimal representation has 6 as the last digit.
2. If the last digit 6 is erased and placed in front of the remaining digits, the resulting number is four times as large as the original number n.

Since n = 10*k + 6, the condition is equivalent to asking that 4*(10*k + 6) = 6*10^b + k, where b = ν(k) the number of decimal digits in k. Simplifying, the problem is equivalent to finding the smallest integer k = 2/13*(10^b - 4) with exactly b digits. In code:

let imo_1962_01 =
  range 1 9 >>= fun b ->
  let e = pow 10 b in
  guard (e mod 13 = 4) >>
  let k = (e - 4) / 13 * 2 in
  guard (List.length % digits $ k = b) >>
  return (10 * k + 6)

The number is not too large, and a 31-bit version of pow is sufficient. The third problem will require more scaffolding:

Five students, A, B, C, D, E, took part in a contest. One prediction was that the contestants would finish in the order ABCDE. This prediction was very poor. In fact no contestant finished in the position predicted, and no two contestants predicted to finish consecutively actually did so. A second prediction had the contestants finishing in the order DAECB. This prediction was better. Exactly two of the contestants finished in the places predicted, and two disjoint pairs of students predicted to finish consecutively actually did so. Determine the order in which the contestants finished.

This is more of a word problem than a combinatorial one, and as the latter is not very straightforward and to brute force it I'll need a number of auxiliary functions. A way to list all permutations is first:

let rec selections = function
| [] -> []
| x :: xs ->
  (x, xs) :: List.fold_right (fun (y, ys) l ->
    (y, x :: ys) :: l) (selections xs) []

let rec permutations = function
| ([] | [_]) as l -> [l]
| l ->
  List.fold_right (fun (y, ys) ->
    List.fold_right (fun zs l -> (y :: zs) :: l)
    (permutations ys)) (selections l) []

The first condition asks for permutations having no fixed points, or derangements. I need a way to single derangements out:

let count_fixpoints l p =
  List.fold_left (fun s (x, y) ->
    if x = y then succ s else s) 0 (List.combine l p)

let is_derangement l p = count_fixpoints l p = 0

Lastly, in order to filter consecutive positions, I need a way to generate them and filter them out:

let intersect l m = List.filter (fun x -> List.mem x l) m

let rec pairup = function
| [] | [_] -> []
| x :: (y :: _ as xs) -> (x, y) :: pairup xs

The solution to the problem is a word-for-word translation of the problem's conditions:

let imo_1963_01 =
  let prediction  = ['D'; 'A'; 'E'; 'C'; 'B'] in
  let contestants = List.sort compare prediction in
  let all_pairs   = pairup contestants
  and pred_pairs  = pairup prediction in
  permutations contestants >>= fun p ->
  guard (is_derangement  contestants p
      && count_fixpoints prediction  p = 2) >>
  let pp = pairup p in
  guard (List.length (intersect all_pairs pp) = 0) >>
  guard (match intersect pred_pairs pp with
  | [(x, y); (z, t)] -> y <> z && t <> x
  | _ -> false) >>
  return p

that is, the solution is to be found among all the permutations of the contestants which are derangements and have exactly two positions in common with the prediction. Of these candidates they must have no pair in common with the pairs in the sorted list of contestants, and has to have two disjoint pairs in common with the prediction.

Some would argue, I'm sure, that monadic code is not purely functional, or that it is too rigidly laid out by the sequential nature of monadic binding. I think that it is ideal to solve these word problems since I find that the solution closely follows the conditions laid out in the statement. All in all I've left out less than 40 lines of perfectly obvious support code, and gave solutions with 5, 7 and 30-odd lines. It was a fun exercise.

The Tree Nursery

On my last post, Benedikt Grundmann asked if changing representations for trie nodes, in particular eliminating 'a option-boxed references, would represent a speed-up or not. I started tinkering with a direct representation:

type 'a t =
{ value : 'a option;
  split : char;
  lokid : 'a t;
  eqkid : 'a t;
  hikid : 'a t; }

(call this representation trie1). This forces me to use a sentinel node to represent empty trees:

let rec empty = { value = None; split = '\255'; lokid = empty; eqkid = empty; hikid = empty; }

let is_empty t = t == empty

(the split value is entirely arbitrary and never relied upon). This recursive structure means that I cannot use pattern matches to discriminate and deconstruct nodes, and all comparisons must be strictly by reference equality on the carefully preserved unique sentinel node. This represented a rough 10% speedup on my original code.

Meanwhile, Mauricio Fernández tried his hand at it, and came with a tantalizing result:

in the lapse between the present post and the preceding one, I implemented ternary trees with different constructors for the nodes with and without a value. In my tests, construction is 50% faster, and lookups are between 20% and 80% faster. The code is much less elegant than the one presented here, because it moreover includes some special-casing to handle the "" (empty string) key, which the above code chokes at.

(It's true that the empty string cannot be a key in both the original and the above implementation. I thought it obvious by the tradeoffs I listed in selecting the representation for the key, but truth is I never guarded the operations against it. It's a bug, not a feature). The numbers are too exciting to not try this approach. By carefully teasing apart each case, I arrived to the following data type:

type 'a t = E
          | L of 'a
          | B of      char * 'a t * 'a t * 'a t
          | K of 'a * char * 'a t * 'a t * 'a t

(call this representation trie2). The first constructor E represents an empty tree. The second constructor L represents a leaf with value α. The third constructor B represents a branching node without value. The fourth constructor K represents a marked, valued branching node. Functions over this type require eight matches in general, four each for the middle and four for the end of the key.

I'll present the benchmark results first, then the code. In every case I've run the benchmark three times in sequence from the command line and retained the last measurement. My test machine is an MSI U100 with an 1.6 GHz Intel Atom and 2 GB RAM running Mac OS X 10.5.7, and the code is compiled with ocamlopt -g -inline 10 -w Aelz -pp camlp4o version 3.11.0. Each run is comprised of three tests using the standard Map balanced-height tree, trie1 and trie2 ternary trees, with a Gc.full_major collection at the start of each. The test data sets are the 234936-word /usr/share/dict/words file in Mac OS X, and a 89546-word Spanish lexicon, with 3.91% overlap between both. The tests measure:

1. Insertion of all the English words into an empty tree via a fold_left
2. Retrieval of the list of all inserted words via a S.fold (with validation of length and sortedness outside the timing)
3. Retrieval of all the English words
4. Retrieval of all the Spanish words
5. Deletion of all the English words from the tree via a fold_left (with validation that the resulting tree is empty)

These are the results for the list of words in dictionary order:

Test of 89546-word set in dictionary order. In ops/s.

Data Structure	Insert	List All	Find All	Find Sp.	Delete All
Map	167,105.77	2,056,784.27	254,101.92	225,067.88	802,020.62
trie1	184,619.90	373,794.66	981,049.90	1,201,493.12	273,188.70
trie2	232,617.76	468,132.42	791,158.14	984,052.05	371,658.90

These are the results for the list of words in median order, that is the list with the median first followed by the left and right sublists both in median order:

Test of 89546-word set in median order. In ops/s.

Data Structure	Insert	List All	Find All	Find Sp.	Delete All
Map	168,873.28	2,100,585.38	251,426.31	224,190.06	318,687.33
trie1	270,556.35	352,670.96	1,669,707.66	2,139,124.18	707,988.31
trie2	329,939.38	462,198.78	1,385,399.36	1,692,227.60	804,071.63

These are the combined results, normalized so that the corresponding operations on Map are 1:

Test of 89546-word set both in dictionary and in median order. Relative to Map

Data Structure	Insert	List All	Find All	Find Sp.	Delete All
trie1 (dict)	1.10481	0.181737	3.86085	5.33836	0.340626
trie2 (dict)	1.39204	0.227604	3.11355	4.37225	0.463403
trie1 (median)	1.60213	0.167892	6.64094	9.54157	2.22158
trie2 (median)	1.95377	0.220033	5.51016	7.54818	2.52307

Some things are easy to see and to explain:

• The insertion code is faster for ternary trees than for the height-balanced tree, as there is no need to maintain expensive shape invariants
• The ternary trees highly favor lookups, especially by non-present keys
• Insertion in sorted order is pessimal for ternary trees. Even so, they remain very fast, as claimed by Bentley and Sedgewick
• The fold reconstructs the keys one character at a time by appending to the current prefix, which is quadratic on the prefix length. This accounts for the speed being a 20% of the corresponding Map.fold
• The shape of the ternary trees depend on the order of key insertion in a dramatic way
• The shape of height-balanced trees also depend heavily on the order of key insertion

Some other things are not so easy to understand. In particular, it is not easy to see why trie2 is 20%–30% faster than trie1 except for lookups where it is a 20% slower.

For comparison, here's trie1's lookup:

let lookup k t =
  let n = String.length k in
  if n == 0 then failwith "lookup" else
  let rec go i t =
    if is_empty t then None else
    let c = k.[i] in
    if t.split > c then go  i    t.lokid else
    if t.split < c then go  i    t.hikid else
    if i + 1   < n then go (i+1) t.eqkid else
                        t.value
  in go 0 t

And here's trie2's:

let lookup k t =
  let n = String.length k in
  let rec go i = function
  | E                             -> None
  | L  v              when i == n -> Some v
  | L  _                          -> None
  | B (   _, _, _, _) when i == n -> None
  | B (   c, l, q, h)             ->
    let c' = k.[i] in
    if  c' < c then go  i    l else
    if  c' > c then go  i    h else
                    go (i+1) q
  | K (v, _, _, _, _) when i == n -> Some v
  | K (_, c, l, q, h)             ->
    let c' = k.[i] in
    if  c' < c then go  i    l else
    if  c' > c then go  i    h else
                    go (i+1) q
  in go 0 t

(It's not that ugly). Twice the number of lines, a slew of cases, it should be obvious why the latter is slower, right? Then compare and contrast: here's trie1's insert:

let add k v t =
  let n = String.length k in
  if n == 0 then failwith "add" else
  let rec go i t =
    if is_empty t then
      if i+1 < n then { empty with split = k.[i]; eqkid = go (i+1) t }
                 else { empty with split = k.[i]; value = Some v     }
    else
    let c = k.[i] in
    if t.split > c then { t with lokid = go  i    t.lokid } else
    if t.split < c then { t with hikid = go  i    t.hikid } else
    if i + 1   < n then { t with eqkid = go (i+1) t.eqkid } else
                        { t with value = Some v           }
  in go 0 t

Here's trie2's

let add k v t =
  let n = String.length k in
  let rec go i = function
  | E                 when i == n -> L  v
  | E                             -> B (   k.[i], E, go (i+1) E, E)
  | L  _              when i == n -> L  v
  | L  v                          -> K (v, k.[i], E, go (i+1) E, E)
  | B (   c, l, q, h) when i == n -> K (v, c    , l, q         , h)
  | B (   c, l, q, h)             ->
    let c' = k.[i] in
    if  c' < c then B (   c, go i l,          q,      h) else
    if  c' > c then B (   c,      l,          q, go i h) else
                    B (   c,      l, go (i+1) q,      h)
  | K (_, c, l, q, h) when i == n -> K (v, c    , l, q         , h)
  | K (v, c, l, q, h)             ->
    let c' = k.[i] in
    if  c' < c then K (v, c, go i l,          q,      h) else
    if  c' > c then K (v, c,      l,          q, go i h) else
                    K (v, c,      l, go (i+1) q,      h)
  in go 0 t

(Now this is not as pretty). They are exactly in the same shape and relation one with another. Some things I've observed:

• OCaml is not as clever as it should be compiling or-patterns. Many of the branches in trie2's functions could be coalesced into those, but the speed hit is very consistently noticeable
• Refining the constructors for each code path pays off in algorithmic ways: the normalization required in remove can be restricted to those branches statically known to have modified children, and the sharing can be made explicit:

  let remove k t =
    let prune = function
    | B (   _, E, E, E) -> E
    | K (v, _, E, E, E) -> L v
    | t                 -> t
    in
    let n = String.length k in
    let rec go i t = match t with
    | E                             -> t
    | L  _              when i == n -> E
    | L  _                          -> t
    | B (   _, _, _, _) when i == n -> t
    | B (   c, l, q, h)             ->
      let c' = k.[i] in
      if  c' < c then prune (B (   c, go i l,          q,      h)) else
      if  c' > c then prune (B (   c,      l,          q, go i h)) else
                      prune (B (   c,      l, go (i+1) q,      h))
    | K (_, c, l, q, h) when i == n -> B (c, l, q, h)
    | K (v, c, l, q, h) ->
      let c' = k.[i] in
      if  c' < c then prune (K (v, c, go i l,          q,      h)) else
      if  c' > c then prune (K (v, c,      l,          q, go i h)) else
                      prune (K (v, c,      l, go (i+1) q,      h))
    in go 0 t
  

  Furthermore, the duality between add and remove is easy to see: the former's base cases introduce L and K nodes, the latter turns them into E and B nodes, respectively.

In any case, I look forward to read Mauricio's analysis and code. All this goes to show that these ternary trees are excellent data structures to represent string maps, as claimed by Bentley and Sedgewick.

Simple, Efficient Trie Maps

I became interested in implementing Bentley and Sedgewick's Ternary Trees by reading the challenge proposed on Programming Praxis (a great blog, by the way). I thought the algorithms would be a good fit for a purely functional data structure, and I am happy with the results. Here is my implementation.

I follow the Map.S signature, so that this can be a drop-in replacement for string-keyed maps, so that the interface (.mli file) is simply:

include Map.S with type key = string

For the implementation I restrict the key space to strings:

type key = string

Trees are optional references to nodes, and these are records containing the mapped value whenever this node represents a full key, the character suffix for the tree and the three branches:

type 'a t = 'a node option
and 'a node =
{ value : 'a option;
  split : char;
  lokid : 'a t;
  eqkid : 'a t;
  hikid : 'a t; }

In this I follow closely the paper, except for a crucial detail: in it, the authors make use of the terminator inherent to every C string, the null byte '\0'. This forces a design decision:

• emulate the termination character. This has the distinct disadvantage of having to prohibit it from occurring in string keys
• use a char option as the key. This forces boxing the key into a constructed type and complicates the logic in all the comparisons, which has a negative impact in the inner loops
• use a char list or string as the key. This has the advantage that it can be adapted to compress the search paths, but it requires generating substrings of everything
• guard on the string length. This is the choice I took. It complicates the comparison logic somewhat, but it has the advantage of making the node representation more regular and compact, and allowing the full range of characters in string keys

An empty tree is represented by a None reference:

let empty = None

Consequently, the test for an empty tree is a match on the option:

let is_empty = function None -> true | Some _ -> false

This means that I'll be forced to normalize trees on deletion, of which more later. The workhorse function is lookup:

let lookup k t =
  let n = String.length k in
  let rec go i = function
  | None   -> None
  | Some t ->
    let cmp = Char.compare k.[i] t.split in
    if cmp < 0 then go  i    t.lokid else
    if cmp > 0 then go  i    t.hikid else
    if i+1 < n then go (i+1) t.eqkid else
                    t.value
  in go 0 t

It is as simple as it can be given the choices I outlined above. The crucial bit is that the presence or absence of the value is the flag that indicates if this node represents a present or absent key in the tree. The signature's retrieval and membership functions are simply wrappers around this:

let find k t = match lookup k t with
| Some x -> x
| None   -> raise Not_found

let mem k t = match lookup k t with
| Some _ -> true
| None   -> false

Inserting a key-value pair in the tree is a bit more involved:

let add k v t =
  let n = String.length k in
  let rec go i = function
  | None when i+1 < n ->
    Some { value = None  ; split = k.[i]; lokid = None; eqkid = go (i+1) None; hikid = None }
  | None              ->
    Some { value = Some v; split = k.[i]; lokid = None; eqkid =          None; hikid = None }
  | Some t            ->
    let cmp = Char.compare k.[i] t.split in
    if cmp < 0 then Some { t with lokid = go  i    t.lokid } else
    if cmp > 0 then Some { t with hikid = go  i    t.hikid } else
    if i+1 < n then Some { t with eqkid = go (i+1) t.eqkid } else
                    Some { t with value = Some v }
  in go 0 t

The case when the end of a branch is reached is inlined, so that a descending branch keyed on the key's suffix is built until reaching its last character, which stores the keyed value. Otherwise, the suffix is compared to the node key, and the appropriate child is built. Note that the recursively built subtree replaces the corresponding child in the node; this is what makes this data structure purely functional and persistent. Removing a mapping is similar, but complicated by the need to normalize the tree so that an empty tree is represented by None:

let remove k t =
  let prune = function
  | { value = None; lokid = None; eqkid = None; hikid = None } -> None
  | t -> Some t
  in
  let n = String.length k in
  let rec go i = function
  | None   -> None
  | Some t ->
    let cmp = Char.compare k.[i] t.split in
    if cmp < 0 then prune { t with lokid = go  i    t.lokid } else
    if cmp > 0 then prune { t with hikid = go  i    t.hikid } else
    if i+1 < n then prune { t with eqkid = go (i+1) t.eqkid } else
                    prune { t with value = None             }
  in go 0 t

In this case, prune works somewhat as a smart constructor for node references.

The big advantage of ternary search trees is that, since they are modeled on three-way quicksort, they provide ordered key storage with relative insensitivity to the order of insertions. The big disadvantage is that they don't have a canonical form for a given key set, that is, their shape depends on the order in which the keys are inserted. This means that they provide a cheap fold but no easy way to compare trees:

let fold f t e =
  let rec go prefix e = function
  | None   -> e
  | Some t ->
    let e   = go prefix e t.lokid in
    let key = prefix ^ String.make 1 t.split in
    let e   = match t.value with None -> e | Some v -> f key v e in
    let e   = go key    e t.eqkid in
              go prefix e t.hikid
  in go "" e t

Note that the key is built character by character but only on the equal kids. Note also that the resulting fold is guaranteed to be ordered on the keys by the sequence in which the recursion is carried out. The rest of the iteration functions in the signature are trivial wrappers:

let iter f t = fold (fun k v () ->        f k v ) t ()
and map  f t = fold (fun k v    -> add k (f v  )) t empty
and mapi f t = fold (fun k v    -> add k (f k v)) t empty

As I indicated above, ternary search trees have no canonical form, so to compare two trees you have to somehow traverse them in synchronization. One easy but expensive way would be to turn both into key-value lists and compare those; this would require O(n min m) time and O(n + m) space. The best way would be to turn the fold into a continuation-passing iterator that lazily streams keys and traverse both in tandem. I leave it as an exercise to the reader:

let compare _ _ _ = failwith "Not comparable"

let equal eq t t' =
  0 == compare (fun v w -> if eq v w then 0 else 1) t t'