AVL Tree in OCaml

open Core;;

type 'a avl_tree =
    | Empty
    | Node of int * 'a * 'a avl_tree * 'a avl_tree;;

let avl_height = function
    | Empty -> 0
    | Node (h, _, _, _) -> h;;

let max_height l r =
    Int.max (avl_height l) (avl_height r);;

let rotate_to_left tr =
    match tr with
    | Empty -> tr
    | Node (_, x, t1, r) ->
            match r with
            | Empty -> tr
            | Node(_, y, t2, t3) ->
                    let tr = avl_height t3 in
                    let tl = 1 + max_height t1 t2 in
                    Node(1 + Int.max tr tl, y, Node(tl, x, t1, t2), t3);;

let rotate_to_right tr =
    match tr with
    | Empty -> tr
    | Node(_, y, l, t3) ->
            match l with
            | Empty -> tr
            | Node(_, x, t1, t2) ->
                    let tl = avl_height t1 in
                    let tr = 1 + max_height t2 t3 in
                    Node (1 + Int.max tl tr, x, t1, Node(tr, y, t2, t3));;

let rec raw_insert tr w =
    match tr with | Empty -> Node(1, w, Empty, Empty) | Node(_, v, l, r) as t->
            if Int.(w < v) then
                let newl = raw_insert l w in
                Node (1 + max_height newl r, v, newl, r)
            else if Int.(w > v) then
                let newr = raw_insert r w in
                Node(1 + max_height newr l, v, l, newr)
            else t;;

type balance_state = Left | Right | Neutral

let check_balance l r =
    let balance = (avl_height l) - (avl_height r) in
    if Int.(balance < -1) then Right
    else if Int.(balance > 1) then Left
    else Neutral;;

let insert tree w =
    let tr = raw_insert tree w in
    match tr with
    | Empty -> tr
    | Node(h, v, l, r) ->
            match check_balance l r with
            | Neutral -> tr
            | Right   ->
                    (match r with
                    | Empty -> tr
                    | Node(rh, rv, t1, t2) ->
                            if Int.(w < rv) then
                                rotate_to_left (Node(h, v, l, Node(rh, rv, rotate_to_right t1, t2)))
                            else
                                rotate_to_left (Node(h, v, l, Node(rh, rv, t1, rotate_to_left t2))))
            | Left ->
                    (match l with
                    | Empty -> tr
                    | Node(lh, lv, t3, t4) ->
                            if Int.(w < lv) then
                                rotate_to_right (Node(h, v, Node(lh, lv, rotate_to_right t3, t4), r))
                            else
                                rotate_to_right (Node(h, v, Node(lh, lv, t3, rotate_to_left t4), r)));;