Binary Search Tree in OCAML

BST.ml

open Core.Std

module type COMPARABLE = sig
    type t
    val compare : t -> t -> int
end

module type NODE = sig
    type 'a t
    val create: 'a -> 'a t
    val elem : 'a t -> 'a
    val root : 'a t -> 'a t option
    val left : 'a t -> 'a t option
    val right : 'a t -> 'a t option
    val set_root : 'a t -> 'a t option -> unit
    val set_left : 'a t -> 'a t option -> unit
    val set_right : 'a t -> 'a t option -> unit
end

module Node : NODE = struct
    type 'a t = {
        elem : 'a;
        mutable root : 'a t option;
        mutable left : 'a t option;
        mutable right : 'a t option
    }
    let create x = { elem = x; root = None; left = None; right = None }
    let elem t = t.elem
    let root t = t.root
    let left t = t.left
    let right t = t.right
    let set_left t x = t.left <- x
    let set_right t x = t.right <- x
    let set_root t x = t.root <- x
end

module Make (Comparable : COMPARABLE) = struct
    type t = { mutable root : Comparable.t Node.t option }
    let empty = { root = None }
    let singleton x = { root = Some (Node.create x) }
    let add t x =
        match t.root with
        | None -> t.root <- Some (Node.create x)
        | Some n ->
                let ( < ) x y = Comparable.compare x y < 0 in
                let nptr = ref n in
                let quit_loop = ref false in
                let new_node = Node.create x in
                while not !quit_loop do
                    if x < Node.elem !nptr then
                        if Option.is_none (Node.left !nptr) then
                            begin
                                Node.set_root new_node (Some !nptr);
                                Node.set_left !nptr (Some new_node);
                                quit_loop := true
                            end
                        else
                            nptr := Option.value_exn (Node.left !nptr)
                    else
                        if Option.is_none (Node.right !nptr) then
                            begin
                                Node.set_root new_node (Some !nptr);
                                Node.set_right !nptr (Some new_node);
                                quit_loop := true
                            end
                        else
                            nptr := Option.value_exn (Node.right !nptr)

                done
    let inorder_walk t ~f =
        let rec inorder_walk_node node =
            begin
                Option.value_map ~f:inorder_walk_node
                                 ~default:() (Node.left node);
                f (Node.elem node);
                Option.value_map ~f:inorder_walk_node
                                 ~default:() (Node.right node)
            end
        in Option.value_map t.root ~default:() ~f:inorder_walk_node
end

module IntTree = Make(Int)

let () = let tree = IntTree.singleton 20 in
         let ls = List.init 20 ~f:(fun n -> Random.int (n + 1)) in
         List.iter ls ~f:(IntTree.add tree);
         IntTree.inorder_walk tree ~f:(fun n -> printf "%i\n" n)