(* Simple 2–3 tree in OCaml.  Rather than indirecting through red-black
   trees I thought I’d see how much of a hassle it was to just implement
   2–3 trees directly. *)

(* A map is a 2–3 tree mapping keys to values; its leaf nodes contain
   no data, and its internal nodes have either 2 or 3 children,
   separated by respectively 1 or 2 key-value pairs. *)
type ('k, 'v) map = 
| Empty
| Two of ('k, 'v) map * ('k, 'v) branch
| Three of ('k, 'v) map * ('k, 'v) branch * ('k, 'v) branch
and ('k, 'v) branch = 'k * 'v * ('k, 'v) map
(* Internally, adding a key-value pair to a map produces an “outcome”.
   This may be a single updated map, or it may be a “split,” where a
   key-value pair propagates up to the parent.
*)
and ('k, 'v) outcome =
| Single of ('k, 'v) map
| Split of ('k, 'v) map * ('k, 'v) branch

let rec get m k = match m with
  | Empty -> None
  | Two(left, (k', v', right)) ->
    if k = k' then Some v' else get (if k < k' then left else right) k
  | Three(left, (k', v', mid), (k'', v'', right)) ->
    if k = k' then Some v' else
      if k < k' then get left k else
        if k = k'' then Some v'' else
          if k < k'' then get mid k else
            get right k

let rec put' m k v = match m with
  | Empty -> Split(Empty, (k, v, Empty))
  | Two(left, (k', v', right)) -> (
    if k = k' then Single(Two(left, (k, v, right))) else
      if k < k' then match put' left k v with
      | Single(left') -> Single(Two(left', (k', v', right)))
      | Split(left', (k'', v'', mid)) ->
        Single(Three(left', (k'', v'', mid), (k', v', right)))
      else match put' right k v with
      | Single(right') -> Single(Two(left, (k', v', right')))
      | Split(mid, (k'', v'', right')) ->
        Single(Three(left, (k', v', mid), (k'', v'', right')))
  )
  | Three(left, ((k', v', mid) as b1), ((k'', v'', right) as b2)) -> 
    if k = k' then Single(Three(left, (k', v, mid), b2)) else
      if k < k' then match put' left k v with
      | Single(left') -> Single(Three(left', b1, b2))
      | Split(left', b3) -> Split(Two(left', b3), (k', v', Two(mid, b2)))
      else if k = k'' then Single(Three(left, b1, (k'', v, right))) else
        if k < k'' then match put' mid k v with
        | Single(mid') -> Single(Three(left, (k', v', mid'), b2))
        | Split(mid', (kn, vn, right')) ->
          Split(Two(left, (k', v', mid')), (kn, vn, Two(right', b2)))
        else match put' right k v with
        | Single(right') -> Single(Three(left, b1, (k'', v'', right')))
        | Split(mid', b3) ->
          Split(Two(left, b1), (k'', v'', Two(mid', b3)))
   
(* Ugh, that’s nasty.  It could probably be improved with a
   dotted-pair structure that is a linked list of branches on a given
   level, which I think gets us to red-black trees.  At first I’d
   flubbed the above: the tree had the right contents (?) but wasn’t
   balancing, because I got the Empty base case wrong. 
   
   dictof [7, (); 6, (); 5, (); 4, (); 3, (); 2, (); 1, ()] was still
   wrong; it duplicates the 6 entry.  Longer such maps got some items
   out of order.  That was a bug in the last case.  Then `invert
   digits` was wrong; it lost 2, duplicated 6, and put 3 and 4 out of
   order.  This was a bug in the Single(mid') -> Single(Three(left,
   (k', v', mid'), b2)) case, which said k'' and v'' instead.
   
   Since then I haven’t found any more bugs.  Maybe it’s correct now?
   I don’t have a lot of confidence.

   Probably wrapping up k and v into a tuple would help avoid such
   bugs too (though not those two in particular.)

   *)

(* Splitting the root produces a new tree level *)
let put m k v = match put' m k v with
  | Single(tree) -> tree
  | Split(left, right) -> Two(left, right)

let rec items' m tl = match m with
  | Empty -> tl
  | Two(left, (k, v, right)) -> items' left ((k, v) :: items' right tl)
  | Three(left, (k, v, mid), (k', v', right)) ->
    items' left ((k, v) :: items' mid ((k', v') :: items' right tl))

let items m = items' m []

let rec dictof' m = function [] -> m | (k, v)::xs -> dictof' (put m k v) xs
let dictof xs = dictof' Empty xs

let rec depth = function
  | Empty -> 0
  | Two(left, (_, _, right)) -> 1 + max (depth left) (depth right)
  | Three(left, (_, _, mid), (_, _, right)) ->
    1 + max (depth left) (max (depth mid) (depth right))
   
let rec depthok = function
  | Empty -> true
  | Two(left, (_, _, right)) ->
    depth left = depth right && depthok left && depthok right
  | Three(left, (_, _, mid), (_, _, right)) ->
    depth left = depth right && depthok left && depthok right &&
    depth left = depth mid && depthok mid
   
and invert d = dictof (List.map (fun (k, v) -> (v, k)) (items d))

let digits = dictof ["one", 1; "two", 2; "three", 3; "four", 4; "five", 5;
                     "six", 6; "seven", 7; "eight", 8; "nine", 9; "zero", 0] ;;