type key == int;;

sort color == {a:int | 0 <= a <= 1};;

datatype tree with (color, nat, nat, nat) = (* color, black height, violation, size *)
    E(0, 0, 0, 0)
  | {cl:color}{cr:color}{bh:nat}{sl:nat}{sr:nat}
    B(0, bh+1, 0, sl+sr+1) of tree(cl, bh, 0, sl) * key * tree(cr, bh, 0, sr)
  | {cl:color}{cr:color}{bh:nat}{sl:nat}{sr:nat}
    R(1, bh, cl+cr, sl+sr+1) of tree(cl, bh, 0, sl) * key * tree(cr, bh, 0, sr)
;;

exception Item_already_exists;;

let balance = function
    (R(R(a, x, b), y, c), z, d) -> R(B(a, x, b), y, B(c, z, d))
  | (R(a, x, R(b, y, c)), z, d) -> R(B(a, x, b), y, B(c, z, d))
  | (a, x, R(R(b, y, c), z, d)) -> R(B(a, x, b), y, B(c, z, d))
  | (a, x, R(b, y, R(c, z, d))) -> R(B(a, x, b), y, B(c, z, d))
(*
  | (a, x, b) -> B(a, x, b)
*)
  | ((R(B _, _, B _) as tl), z, (R(B _, _, B _) as tr)) -> B(tl, z, tr)
  | ((R(B _, _, B _) as tl), z, (R(B _, _, E) as tr)) -> B(tl, z, tr)
  | ((R(B _, _, B _) as tl), z, (R(E, _, B _) as tr)) -> B(tl, z, tr)
  | ((R(B _, _, B _) as tl), z, (R(E, _, E) as tr)) -> B(tl, z, tr)
  | ((R(B _, _, B _) as tl), z, (B _ as tr)) -> B(tl, z, tr)
  | ((R(B _, _, B _) as tl), z, (E as tr)) -> B(tl, z, tr)

  | ((R(B _, _, E) as tl), z, (R(B _, _, B _) as tr)) -> B(tl, z, tr)
  | ((R(B _, _, E) as tl), z, (R(B _, _, E) as tr)) -> B(tl, z, tr)
  | ((R(B _, _, E) as tl), z, (R(E, _, B _) as tr)) -> B(tl, z, tr)
  | ((R(B _, _, E) as tl), z, (R(E, _, E) as tr)) -> B(tl, z, tr)
  | ((R(B _, _, E) as tl), z, (B _ as tr)) -> B(tl, z, tr)
  | ((R(B _, _, E) as tl), z, (E as tr)) -> B(tl, z, tr)

  | ((R(E, _, B _) as tl), z, (R(B _, _, B _) as tr)) -> B(tl, z, tr)
  | ((R(E, _, B _) as tl), z, (R(B _, _, E) as tr)) -> B(tl, z, tr)
  | ((R(E, _, B _) as tl), z, (R(E, _, B _) as tr)) -> B(tl, z, tr)
  | ((R(E, _, B _) as tl), z, (R(E, _, E) as tr)) -> B(tl, z, tr)
  | ((R(E, _, B _) as tl), z, (B _ as tr)) -> B(tl, z, tr)
  | ((R(E, _, B _) as tl), z, (E as tr)) -> B(tl, z, tr)

  | ((R(E, _, E) as tl), z, (R(B _, _, B _) as tr)) -> B(tl, z, tr)
  | ((R(E, _, E) as tl), z, (R(B _, _, E) as tr)) -> B(tl, z, tr)
  | ((R(E, _, E) as tl), z, (R(E, _, B _) as tr)) -> B(tl, z, tr)
  | ((R(E, _, E) as tl), z, (R(E, _, E) as tr)) -> B(tl, z, tr)
  | ((R(E, _, E) as tl), z, (B _ as tr)) -> B(tl, z, tr)
  | ((R(E, _, E) as tl), z, (E as tr)) -> B(tl, z, tr)

  | ((B _ as tl), z, (R(B _, _, B _) as tr)) -> B(tl, z, tr)
  | ((B _ as tl), z, (R(B _, _, E) as tr)) -> B(tl, z, tr)
  | ((B _ as tl), z, (R(E, _, B _) as tr)) -> B(tl, z, tr)
  | ((B _ as tl), z, (R(E, _, E) as tr)) -> B(tl, z, tr)
  | ((B _ as tl), z, (B _ as tr)) -> B(tl, z, tr)
  | ((B _ as tl), z, (E as tr)) -> B(tl, z, tr)

  | ((E as tl), z, (R(B _, _, B _) as tr)) -> B(tl, z, tr)
  | ((E as tl), z, (R(B _, _, E) as tr)) -> B(tl, z, tr)
  | ((E as tl), z, (R(E, _, B _) as tr)) -> B(tl, z, tr)
  | ((E as tl), z, (R(E, _, E) as tr)) -> B(tl, z, tr)
  | ((E as tl), z, (B _ as tr)) -> B(tl, z, tr)
  | ((E as tl), z, (E as tr)) -> B(tl, z, tr)
withtype {cl:color}{cr:color}{bh:nat}{vl: nat}{vr:nat | vl+vr <= 1}{sl:nat}{sr:nat}
         tree(cl, bh, vl, sl) * key * tree(cr, bh, vr, sr) ->
         [c:color] tree(c, bh+1, 0, sl+sr+1)
;;

let insert x t =
  let rec ins = function
      E -> R(E, x, E)
    | B(a, y, b) ->
      if x < y then balance(ins a, y, b)
      else if y < x then balance(a, y, ins b)
           else raise Item_already_exists
    | R(a, y, b) ->
      if x < y then R(ins a, y, b)
      else if y < x then R(a, y, ins b)
           else raise Item_already_exists
   withtype {c:color}{bh:nat}{s:nat}
            tree(c, bh, 0, s) ->  [c':color][v:nat | v <= c] tree(c', bh, v, s+1) in
   match ins t with
     R(a, y, b) -> B(a, y, b)
(*
   | t -> t
*)
   | B _ as t -> t
   | E as t -> t
withtype {c:color}{bh:nat}{s:nat} key -> tree(c, bh, 0, s) -> [bh:nat] tree(0, bh, 0, s+1)
;;