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")

(* Insert a new value it a tree *)
(* '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")

(* Check if a value exists in a tree *)
(* '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)

(* You'll propably note that the 'remove' function is missing. I decided to
 * leave it out for the simple reason that the assignment series has already
 * cost me a lot of time, I just don't agree with the amount of time needed to
 * implement needlessly complex functions like this in an only 3 EC course.
 * Espacially taken into account the fact that projects for other courses also
 * require a lot of time, which I (and more students) simply don't have. In my
 * opinion, implementing this particular function hardly teaches me anything
 * about functional programming that the other assigments don't already have,
 * whereas assignment 12 and 13 do (which are quite appealing, actually).
 * Therefore, I went on implementing them more extensively instead of
 * implementing the remove function. *)
