A functional red-black tree written in F#.

RedBlackTree.fs

module rec RedBlackTree

type Color =
    | NegBlack
    | Red
    | Black
    | DoubleBlack

let private redder = function
    | NegBlack | Red -> NegBlack
    | Black -> Red
    | DoubleBlack -> Black
let private blacker = function
    | NegBlack -> Red
    | Red -> Black
    | Black | DoubleBlack -> DoubleBlack
    
type Tree<'a when 'a : comparison> =
    | Leaf
    | BlackLeaf
    | Branch of color : Color * left : Tree<'a> * element : 'a * right : Tree<'a>
    
let private blacken = function
    | Leaf | BlackLeaf -> Leaf
    | Branch (_, l, x, r) -> Branch (Black, l, x, r)

let private redderTree = function
    | Leaf | BlackLeaf -> Leaf
    | Branch(c, l, x, r) -> Branch(redder c, l, x, r)
    
let private blackerTree = function
    | Leaf | BlackLeaf -> BlackLeaf
    | Branch(c, l, x, r) -> Branch(blacker c, l, x, r)

let private isDoubleBlack = function
    | BlackLeaf | Branch(color = DoubleBlack) -> true
    | _ -> false

let private balance = function
    | Branch(Black | DoubleBlack as color, Branch(Red, a, x, Branch(Red, b, y, c)), z, d)
    | Branch(Black | DoubleBlack as color, Branch(Red, Branch(Red, a, x, b), y, c), z, d)
    | Branch(Black | DoubleBlack as color, a, x, Branch(Red, b, y, Branch(Red, c, z, d)))
    | Branch(Black | DoubleBlack as color, a, x, Branch(Red, Branch(Red, b, y, c), z, d)) ->
        Branch(redder color, Branch(Black, a, x, b), y, Branch(Black, c, z, d))
    | Branch(DoubleBlack, a, x, Branch(NegBlack, Branch(Black, b, y, c), z, Branch(Black, d, e, f))) ->
        Branch(Black, Branch(Black, a, x, b), y, balance(Branch(Black, c, z, Branch(Red, d, e, f))))
    | Branch(DoubleBlack, Branch(NegBlack, (Branch(color = Black) as tree), x, Branch(Black, b, y, c)), z, d) ->
        Branch(Black, balance(Branch(Black, redderTree tree, x, b)), y, Branch(Black, c, z, d))
    | tree -> tree

let private bubble = function
    | Branch(color, left, elem, right) when isDoubleBlack left || isDoubleBlack right ->
        balance(Branch(blacker color, redderTree left, elem, redderTree right))
    | tree -> balance tree
    
let max = function
    | Branch(_, _, x, Leaf) -> x
    | Branch(_, _, _, right) -> max right
    | _ -> invalidOp "max of leaf"

let removeMax = function
    | Branch(color, left, elem, Leaf) ->
        remove(Branch(color, left, elem, Leaf))
    | Branch(color, left, elem, right) ->
        bubble(Branch(color, left, elem, removeMax right))
    | _ -> invalidOp "remove max of leaf"

let remove = function
    | Leaf | BlackLeaf -> Leaf
    | Branch(Red, Leaf, _, Leaf) -> Leaf
    | Branch(Black, Leaf, _, Leaf) -> BlackLeaf
    | Branch(Black, Leaf, _, Branch(Red, a, x, b))
    | Branch(Black, Branch(Red, a, x, b), _, Leaf) ->
        Branch(Black, a, x, b)
    | Branch(color, left, _, right) ->
        let maxVal = max left
        let removedLeft = removeMax left
        bubble(Branch(color, removedLeft, maxVal, right))

let delete x tree =
    let rec rem = function
        | Leaf | BlackLeaf -> Leaf
        | Branch(color, left, y, right) as t ->
            if x = y then remove t
            else if x < y then bubble(Branch(color, rem left, y, right))
            else bubble(Branch(color, left, y, rem right))
    
    blacken (rem tree)

let insert x tree =
    let rec ins = function
        | Leaf | BlackLeaf -> Branch(Red, Leaf, x, Leaf)
        | Branch(color, left, y, right) as t ->
            if x = y then t
            else if x < y then balance(Branch(color, ins left, y, right))
            else balance(Branch(color, left, y, ins right))
    
    blacken (ins tree)

let contains x = function
    | Leaf | BlackLeaf -> false
    | Branch(_, _, y, _) when x = y -> true
    | Branch(_, left, y, right) ->
        if x < y then contains x left
        else contains x right