GADT balanced AVL trees

avl.ml

type z = Zero
type 'a s = Succ of 'a

type _ nat =
  | Z : z nat
  | S : 'n nat -> 'n s nat

type (_, _, _) balance =
  | Eq    : ('h  , 'h  , 'h  ) balance
  | Left  : ('h  , 'h s, 'h s) balance
  | Right : ('h s, 'h s, 'h  ) balance

type ('a, _) avl =
  | Leaf : (_, z) avl
  | Branch
     : ('hl, 'h, 'hr) balance
     * ('a, 'hl) avl
     * 'a
     * ('a, 'hr) avl
    -> ('a, 'h s) avl

(* {{{ Utils *)

let rec to_list
  : type a h. a list -> (a, h) avl -> a list
  = fun acc -> function
    | Leaf -> acc
    | Branch (_, left, pivot, right) ->
      to_list (pivot :: to_list acc right) left

let to_list t = to_list [] t

(* }}} *)

(* {{{ Insert *)

type (_, _) ins =
  | Ok : ('a, 'h) avl -> ('a, 'h) ins
  | Overflow : ('a, 'h s) avl -> ('a, 'h) ins

let single elt = Overflow (Branch (Eq, Leaf, elt, Leaf))

let ins_left
  : type a hl h hr
  . (hl, h, hr) balance -> (a, hl) ins -> a -> (a, hr) avl -> (a, h s) ins
  = fun balance left pivot right ->
  match left with
  | Ok left -> Ok (Branch (balance, left, pivot, right))
  | Overflow (Branch (balance_left, a, p, b) as left) ->
    match balance with
    | Left -> Ok (Branch (Eq, left, pivot, right))
    | Eq -> Overflow (Branch (Right, left, pivot, right))
    | Right ->
      match balance_left with
      | Right -> Ok (Branch (Eq, a, p, Branch (Eq, b, pivot, right)))
      | Eq -> Overflow (Branch (Left, a, p, Branch (Right, b, pivot, right)))
      | Left ->
        let bal bl br x q y =
          Branch (Eq, Branch (bl, a, p, x),
                      q,
                      Branch (br, y, pivot, right)) in
        Ok (match b with
            | Branch (Eq,    x, q, y) -> bal Eq    Eq   x q y
            | Branch (Left,  x, q, y) -> bal Right Eq   x q y
            | Branch (Right, x, q, y) -> bal Eq    Left x q y)

let ins_right
  : type a hl h hr
  . (hl, h, hr) balance -> (a, hl) avl -> a -> (a, hr) ins -> (a, h s) ins
  = fun balance left pivot right ->
  match right with
  | Ok right -> Ok (Branch (balance, left, pivot, right))
  | Overflow (Branch (balance_right, a, p, b) as right) ->
    match balance with
    | Right -> Ok (Branch (Eq, left, pivot, right))
    | Eq -> Overflow (Branch (Left, left, pivot, right))
    | Left ->
      match balance_right with
      | Left -> Ok (Branch (Eq, Branch (Eq, left, pivot, a), p, b))
      | Eq -> Overflow (Branch (Right, Branch (Left, left, pivot, a), p, b))
      | Right ->
        let bal bl br x q y =
          Branch (Eq, Branch (bl, left, pivot, x),
                      q,
                      Branch (br, y, p, b)) in
        Ok
          (match a with
           | Branch (Eq,    x, q, y) -> bal Eq    Eq   x q y
           | Branch (Left,  x, q, y) -> bal Right Eq   x q y
           | Branch (Right, x, q, y) -> bal Eq    Left x q y)

let rec cons
  : type a h . a -> (a, h) avl -> (a, h) ins
  = fun elt -> function
    | Leaf -> single elt
    | Branch (balance, left, pivot, right) ->
      ins_left balance (cons elt left) pivot right

let rec snoc
  : type a h . a -> (a, h) avl -> (a, h) ins
  = fun elt -> function
    | Leaf -> single elt
    | Branch (balance, left, pivot, right) ->
      ins_right balance left pivot (snoc elt right)

let rec insert
  : type a h . (a -> a -> int) -> a -> (a, h) avl -> (a, h) ins
  = fun compare elt -> function
    | Leaf -> single elt
    | Branch (balance, left, pivot, right) ->
      if compare elt pivot <= 0
      then ins_left  balance (insert compare elt left) pivot right
      else ins_right balance left pivot (insert compare elt right)

(* }}} *)

(* {{{ Pop *)

type (_, _) rem =
  | Uk : ('a, 'h) avl -> ('a, 'h) rem
  | Underflow : ('a, 'h) avl -> ('a, 'h s) rem

let rem_left
  : type a hl h hr
  . (hl, h, hr) balance -> (a, hl) rem -> a -> (a, hr) avl -> (a, h s) rem
  = fun balance left pivot right ->
    match left with
      | Uk left -> Uk (Branch (balance, left, pivot, right))
      | Underflow left ->
        match balance with
        | Eq -> Uk (Branch (Left, left, pivot, right))
        | Right -> Underflow (Branch (Eq, left, pivot, right))
        | Left ->
          match right with
            | Branch (Eq, a, p, b) ->
              Uk (Branch (Right, Branch (Left, left, pivot, a), p, b))
            | Branch (Left, a, p, b) ->
              Underflow (Branch (Eq, Branch (Eq, left, pivot, a), p, b))
            | Branch (Right, center, q, r) ->
              let bal bl br a p b =
                Branch (Eq, Branch (bl, left, pivot, a),
                            p,
                            Branch (br, b, q, r)) in
              Underflow
                (match center with
                 | Branch (Eq,    a, p, b) -> bal Eq    Eq   a p b
                 | Branch (Left,  a, p, b) -> bal Right Eq   a p b
                 | Branch (Right, a, p, b) -> bal Eq    Left a p b)

let rem_right
  : type a hl h hr
  . (hl, h, hr) balance -> (a, hl) avl -> a -> (a, hr) rem -> (a, h s) rem
  = fun balance left pivot right ->
    match right with
      | Uk right -> Uk (Branch (balance, left, pivot, right))
      | Underflow right ->
        match balance with
        | Eq -> Uk (Branch (Right, left, pivot, right))
        | Left -> Underflow (Branch (Eq, left, pivot, right))
        | Right ->
          match left with
            | Branch (Eq, a, p, b) ->
              Uk (Branch (Left, a, p, Branch (Right, b, pivot, right)))
            | Branch (Right, a, p, b) ->
              Underflow (Branch (Eq, a, p, Branch (Eq, b, pivot, right)))
            | Branch (Left, r, q, center) ->
              let bal bl br a p b =
                Branch (Eq, Branch (bl, r, q, a),
                            p,
                            Branch (br, b, pivot, right)) in
              Underflow
                (match center with
                 | Branch (Eq,    a, p, b) -> bal Eq    Eq   a p b
                 | Branch (Left,  a, p, b) -> bal Right Eq   a p b
                 | Branch (Right, a, p, b) -> bal Eq    Left a p b)

let rec pop_front
  : type a h . (a, h) avl -> a * (a, h) rem
  = function
    | Leaf -> raise Not_found
    | Branch (Eq,   Leaf, pivot, right) -> pivot, Underflow right
    | Branch (Left, Leaf, pivot, right) -> pivot, Underflow right
    | Branch (balance, (Branch _ as left), pivot, right) ->
      let elt, left = pop_front left in
      elt, rem_left balance left pivot right

let rec pop_back
  : type a h . (a, h) avl -> a * (a, h) rem
  = function
    | Leaf -> raise Not_found
    | Branch (Eq,    left, pivot, Leaf) -> pivot, Underflow left
    | Branch (Right, left, pivot, Leaf) -> pivot, Underflow left
    | Branch (balance, left, pivot, (Branch _ as right)) ->
      let elt, right = pop_back right in
      elt, rem_right balance left pivot right

(* }}} *)

(* {{{ Remove *)

let rec remove
  : type a h
  . (a -> int) -> (a, h) avl -> a * (a, h) rem
  = fun select -> function
    | Leaf -> raise Not_found
    | Branch (balance, left, pivot, right) ->
      match select pivot with
      | c when c < 0 ->
        let elt, left = remove select left in
        elt, rem_left balance left pivot right
      | c when c > 0 ->
        let elt, right = remove select right in
        elt, rem_right balance left pivot right
      | _ (* 0 *) ->
        pivot, match balance with
        | Eq ->
          begin match left with
            | Leaf -> Underflow right
            | _ ->
              let pivot, left = pop_back left in
              match left with
              | Uk left -> Uk (Branch (Eq, left, pivot, right))
              | Underflow left -> Uk (Branch (Left, left, pivot, right))
          end
        | Right ->
          let pivot, left = pop_back left in
          begin match left with
          | Uk left -> Uk (Branch (Right, left, pivot, right))
          | Underflow left -> Underflow (Branch (Eq, left, pivot, right))
          end
        | Left ->
          let pivot, right = pop_front right in
          begin match right with
          | Uk right -> Uk (Branch (Left, left, pivot, right))
          | Underflow right -> Underflow (Branch (Eq, left, pivot, right))
          end

(* }}} *)

(* {{{ Append *)

type (_, _, _) add =
  | Add_z : (z, 'a, 'a) add
  | Add_s : ('a, 'b, 'c) add -> ('a s, 'b, 'c s) add

type ('hl, 'hr, 'delta, 'max) diff =
  | Same_height : ('a, 'a, z, 'a) diff
  | Greater_right : ('d, 'a, 'b) add -> ('a, 'b, 'd, 'b) diff
  | Greater_left  : ('d, 'b, 'a) add -> ('a, 'b, 'd, 'a) diff

let rec append
  : type a hl h hr d
  . (hl, hr, d, h) diff -> (a, hl) avl -> (a, hr) avl -> (a, h) ins
  = fun diff left right ->
    match diff, left, right with
    | Greater_right Add_z, _, _ -> append Same_height left right
    | Greater_left  Add_z, _, _ -> append Same_height left right
    | Same_height, _, _ ->
      let pivot, left = pop_back left in
      begin match left with
        | Uk left -> Overflow (Branch (Eq, left, pivot, right))
        | Underflow left -> Overflow (Branch (Left, left, pivot, right))
      end
    | Greater_right (Add_s diff), _, Branch (bal, center, pivot, right) ->
      let merge' diff bal =
        ins_left bal (append diff left center) pivot right in
      let merge diff = merge' (Greater_right diff) in
      begin match bal, diff with
      | Left, Add_z -> merge' (Greater_left (Add_s Add_z)) Eq
      | Left, Add_s diff -> merge diff bal
      | Eq,    _ -> merge diff bal
      | Right, _ -> merge diff bal
      end
    | Greater_left (Add_s diff), Branch (bal, left, pivot, center), _ ->
      let merge' diff bal =
        ins_right bal left pivot (append diff center right) in
      let merge diff = merge' (Greater_left diff) in
      begin match bal, diff with
      | Right, Add_z -> merge' (Greater_right (Add_s Add_z)) Eq
      | Right, Add_s diff -> merge diff bal
      | Eq,   _ -> merge diff bal
      | Left, _ -> merge diff bal
      end

(* }}} *)

(* {{{ Wrapper *)

type 'a t = Avl : 'h nat * ('a, 'h) avl -> 'a t

let empty = Avl (Z, Leaf)

let is_empty = function
  | Avl (_, Leaf) -> true
  | _ -> false

let to_list (Avl (_, t)) = to_list t

let wrap_ins h = function
  | Ok t -> Avl (h, t)
  | Overflow t -> Avl (S h, t)

let cons elt (Avl (h, t)) = wrap_ins h (cons elt t)
let snoc elt (Avl (h, t)) = wrap_ins h (snoc elt t)
let insert compare elt (Avl (h, t)) = wrap_ins h (insert compare elt t)

let wrap_rem
  : type a h. h nat -> (a, h) rem -> a t
  = fun h -> function
  | Uk t -> Avl (h, t)
  | Underflow t -> match h with S h -> Avl (h, t)

let pop_front (Avl (h, t)) =
  let elt, t = pop_front t in
  elt, wrap_rem h t

let pop_back (Avl (h, t)) =
  let elt, t = pop_back t in
  elt, wrap_rem h t

let remove select (Avl (h, t)) =
  let elt, t = remove select t in
  elt, wrap_rem h t


let rec add_nat
  : type a . a nat -> (a, z, a) add
  = function
    | Z -> Add_z
    | S a -> Add_s (add_nat a)

let rec add_zs
  : type a b c . (a, b, c) add -> (a, b s, c s) add
  = function Add_z -> Add_z | Add_s p -> Add_s (add_zs p)

type (_, _) nat_diff = Diff : 'm nat * ('a, 'b, _, 'm) diff -> ('a, 'b) nat_diff

let rec nat_diff
  : type a b . a nat -> b nat -> (a, b) nat_diff
  = fun a b ->
    match a, b with
    | Z, Z -> Diff (Z, Same_height)
    | Z, b -> Diff (b, Greater_right (add_nat b))
    | a, Z -> Diff (a, Greater_left  (add_nat a))
    | S a, S b ->
      let Diff (mx, p) = nat_diff a b in
      let mx = S mx in
      match p with
      | Same_height -> Diff (mx, Same_height)
      | Greater_right p -> Diff (mx, Greater_right (add_zs p))
      | Greater_left  p -> Diff (mx, Greater_left  (add_zs p))

let append (Avl (hl, left)) (Avl (hr, right)) =
  let Diff (mx, d) = nat_diff hl hr in
  wrap_ins mx (append d left right)

(* }}} *)