type color = Red | Black
type 'a rbtree = Leaf | Node of 'a * color * 'a rbtree * 'a rbtree

(* Balance a node by matching it against the four cases and returning a new,
 * well-structured node *)
let balance_node node = match node with
      Node (z, Black, Node (y, Red, Node (x, Red, a, b), c), d)
    | Node (z, Black, Node (x, Red, a, Node (y, Red, b, c)), d)
    | Node (x, Black, a, Node (z, Red, Node (y, Red, b, c), d))
    | Node (x, Black, a, Node (y, Red, b, Node (z, Red, c, d))) ->
        Node (y, Red, Node (x, Black, a, b), Node (z, Black, c, d))
    | Node (a, b, c, d) -> node
    | Leaf -> raise (Failure "Cannot balance a leaf")

(* 'a rbtree -> 'a -> 'a rbtree *)
let insert tree value =
    let rec insert_in_tree = function
        (* Insert new Node at leaf position *)
          Leaf -> Node (value, Red, Leaf, Leaf)
        | Node (v, color, left, right) ->
            if value = v then
                (* Value has already been inserted *)
                tree
            else if value < v then
                (* Inserted value is smaller, insert in left tree *)
                balance_node (Node (v, color, insert_in_tree left, right))
            else
                (* Inserted value is larger, insert in right tree *)
                balance_node (Node (v, color, left, insert_in_tree right))
    in
    match insert_in_tree tree with
    (* Color the root black *)
      Node (v, _, left, right) -> Node (v, Black, left, right)
    | Leaf -> raise (Failure "Error during insertion")

(* 'a -> 'a rbtree -> bool *)
let rec lookup value = function
      Leaf -> false
    | Node (v, _, left, right) -> v = value
        || (if value < v then lookup value left else lookup value right)
