Naive implementation of red-black tree with OCaml

rbtree.ml

module RBTree : sig
  type 'a t
  val empty : 'a t
  val insert : 'a t -> 'a -> 'a t
  val exists : 'a t -> 'a -> bool
  val remove : 'a t -> 'a -> 'a t
  val fold_left : ('b -> 'a -> 'b) -> 'b -> 'a t -> 'b
  val fold_right : ('a -> 'b -> 'b) -> 'a t -> 'b -> 'b
  val balanced : 'a t -> bool
  val from_list : 'a list -> 'a t
  val to_list : 'a t -> 'a list
end =
struct
  type color = Red | Black

  type 'a t =
  | Leaf
  | Node of color * 'a t * 'a * 'a t

  let empty = Leaf

  let insert tree v =
    let makeBlack tree =
      match tree with
      | Node (c, l, v, r) -> Node (Black, l, v, r)
      | Leaf -> Leaf
    in
    let balance color left v right =
      match (color, left, v, right) with
      (*
            ((x))         (y)
             / \         /   \
           (y)  d     ((z)) ((x))
           / \     ->  / \   / \
         (z)  c       a   b c   d
         / \
        a   b
      *)
      | (Black, Node (Red, Node (Red, a, z, b), y, c), x, d) ->
        Node (Red, Node (Black, a, z, b), y, Node (Black, c, x, d))
      (*
            ((x))         (y)
             / \         /   \
           (z)  d     ((z)) ((x))
           / \     ->  / \   / \
          a  (y)      a   b c   d
             / \
            b   c
      *)
      | (Black, Node (Red, a, z, Node (Red, b, y, c)), x, d) ->
        Node (Red, Node (Black, a, z, b), y, Node (Black, c, x, d))
      (*
          ((z))           (y)
           / \           /   \
          a  (y)      ((z)) ((x))
             / \   ->  / \   / \
            b  (x)    a   b c   d
               / \
              c   d
      *)
      | (Black, a, z, Node (Red, b, y, Node (Red, c, x, d))) ->
        Node (Red, Node (Black, a, z, b), y, Node (Black, c, x, d))
      (*
          ((z))          ((y))
           / \           /   \
          a  (x)      ((z)) ((x))
             / \   ->  / \   / \
           (y)  d     a   b c   d
           / \
          b   c
      *)
      | (Black, a, z, Node (Red, Node (Red, b, y, c), x, d)) ->
        Node (Red, Node (Black, a, z, b), y, Node (Black, c, x, d))
      | _ -> Node (color, left, v, right)
    in
    let rec ins tree =
      match tree with
      | Node (_, _, x, _) when x = v -> tree
      | Node (c, l, x, r) when x > v -> balance c (ins l) x r
      | Node (c, l, x, r) when x < v -> balance c l x (ins r)
      | _ -> Node (Red, Leaf, v, Leaf)
    in
    makeBlack (ins tree)

  let rec exists tree v =
    match tree with
    | Node (_, _, value, _) when value = v -> true 
    | Node (_, left, value, _) when value > v -> exists left v
    | Node (_, _, value, right) when value < v -> exists right v
    | _ -> false

  let rec fold_left f init tree =
    match tree with
    | Node (_, left, x, right) -> fold_left f (f (fold_left f init left) x) right
    | Leaf -> init

  let rec fold_right f tree init =
    match tree with
    | Node (_, left, x, right) -> fold_right f left (f x (fold_right f right init))
    | Leaf -> init

  let balanced tree =
    let rec black_count tree =
      match tree with
      | Node (Red, Node (Red, _, _, _), _, _) -> None
      | Node (Red, _, _, Node (Red, _, _, _)) -> None
      | Node (color, l, _, r) -> begin
        match (black_count l, black_count r) with
          | (Some ln, Some rn) when ln = rn -> Some (if color = Black then ln + 1 else ln)
          | _ -> None
        end
      | Leaf -> Some 1
    in
    match black_count tree with
    | Some _ -> true
    | _ -> false

  let from_list l =
    List.fold_left (fun t v -> insert t v) empty l

  let to_list tree =
    fold_right (fun x xs -> x :: xs) tree []

  let remove tree v =
    from_list (List.filter (fun value -> value <> v) (to_list tree))
end

let () =
  let tree = RBTree.from_list [1; 2; 3; 4; 5; 6; 7] in
  assert (RBTree.exists tree 5);
  assert (not (RBTree.exists tree 10));
  assert (RBTree.balanced tree);
  assert (RBTree.balanced (RBTree.remove tree 5))