Red-Black Tree Implementation

rbtree.jl

# load("examples/tree.jl")
typealias Color Int8

# TODO: RBTreeSet
const BLACK = convert(Color, 0)
const RED = convert(Color, 1)
const NOCOLOR = convert(Color, 2)

# TODO: Use Two Colors for nodes
abstract RBTreeNode{K,V}
type Nil{K,V} <: RBTreeNode{K,V}
end

type BlackNode{K,V} <: RBTreeNode{K,V}
    key::  K
    value:: V
    left:: RBTreeNode{K,V}
    right::RBTreeNode{K,V}
    parent::RBTreeNode{K,V}
end
type RedNode{K,V} <: RBTreeNode{K,V}
    key::  K
    value:: V
    left:: RBTreeNode{K,V}
    right::RBTreeNode{K,V}
    parent::RBTreeNode{K,V}

    RedNode{K,V}(k::K, v::V) = new(k,v, Nil{K,V}(), Nil{K,V}(), Nil{K,V}())
    RedNode{K,V}(k::K, v::V, l, r, p) = new(k,v, l,r,p)
end
typealias NotRed Union(Nil, BlackNode)
typealias NotBlack Union(Nil, RedNode)

type RBTree{K,V}
    root:: RBTreeNode{K,V}
    nitems

    RBTree() = new(Nil{K,V}(), 0)
end
function show(io, tree::RBTree)
    print(io, "RBTree({")
    first = true
    for (k, v) = tree
        first || print(io, ',')
        first = false
        show(io, k)
        print(io, "=>")
        show(io, v)
    end
    print(io, "})")
end



replace_child(pn::Nil, cur, new) = Nothing
function replace_child(pn, cur, new)
    if left(pn) == cur
        set_left(pn, new)
    else
        set_right(pn, new)
    end
end

color(::Nil) = NOCOLOR # TODO: make this go away
color(::RedNode) = RED
make_red(::RBTree, n::RedNode) = n
function make_black{K,V}(tree::RBTree{K,V}, oldn::RedNode{K,V}) 
    newn = BlackNode{K,V}(key(oldn),value(oldn), left(oldn), right(oldn), parent(oldn))

    replace_child(parent(oldn), oldn, newn)
    set_parent(right(newn), newn)
    set_parent(left(newn), newn)

    if tree.root == oldn
        tree.root = newn
    end
    newn
end

color(::BlackNode) = BLACK
make_black(::RBTree, n::BlackNode) = n
function make_red{K,V}(tree::RBTree{K,V}, oldn::BlackNode{K,V}) 
    newn = RedNode{K,V}(key(oldn),value(oldn), left(oldn), right(oldn), parent(oldn))

    replace_child(parent(oldn), oldn, newn)
    set_parent(right(newn), newn)
    set_parent(left(newn), newn)

    if tree.root == oldn
        tree.root = newn
    end
    newn
end

parent{K,V}(n::BlackNode{K,V}) = n.parent
parent{K,V}(n::RedNode{K,V}) = n.parent
set_parent(n::Nil, k) = Nothing
set_parent{K,V}(n::BlackNode{K,V}, parent) = n.parent = parent
set_parent{K,V}(n::RedNode{K,V}, parent) = n.parent = parent

left{K,V}(n::BlackNode{K,V}) = n.left
set_left{K,V}(n::BlackNode{K,V}, left) = n.left = left
left{K,V}(n::RedNode{K,V}) = n.left
set_left{K,V}(n::RedNode{K,V}, left) = n.left = left

right{K,V}(n::BlackNode{K,V}) = n.right
set_right{K,V}(n::BlackNode{K,V}, right) = n.right = right
right{K,V}(n::RedNode{K,V}) = n.right
set_right{K,V}(n::RedNode{K,V}, right) = n.right = right

key{K,V}(n::BlackNode{K,V}) = n.key
set_key{K,V}(n::BlackNode{K,V}, nkey::K) = n.key = nkey
key{K,V}(n::RedNode{K,V}) = n.key
set_key{K,V}(n::RedNode{K,V}, nkey::K) = n.key = nkey

value{K,V}(n::BlackNode{K,V}) = n.value
set_value{K,V}(n::BlackNode{K,V}, nvalue::V) = n.value = nvalue
value{K,V}(n::RedNode{K,V}) = n.value
set_value{K,V}(n::RedNode{K,V}, nvalue::V) = n.value = nvalue


function _insert_or_modify{K,V}(tree::RBTree, val::V, key::K)
    if tree.nitems == 0
        tree.root = make_black(tree, RedNode{K,V}(key,val)) # TODO nonsens
        tree.nitems += 1
    else
        _insert_or_modify(tree, tree.root, RedNode{K,V}(key,val))
    end
end
function _insert_or_modify(tree::RBTree, t::RBTreeNode, r::RBTreeNode)
    if key(r) == key(t)
        set_value(t, value(r))
    elseif key(r) < key(t)
        if nil(left(t))
            set_left(t, r)
            r = make_red(tree, r)
            set_parent(r, t)
            tree.nitems += 1
            _just_inserted(tree, r, t)
        else
            _insert_or_modify(tree, left(t), r)
        end
    else
        if nil(right(t))
            set_right(t, r)
            r = make_red(tree, r)
            set_parent(r, t)
            tree.nitems += 1
            _just_inserted(tree, r, t)
        else
            _insert_or_modify(tree, right(t), r)
        end
    end
end


nil(t::Nil) = true
nil(k::RBTreeNode) = false

function uncle{K,V}(par::RBTreeNode{K,V})
    gp = parent(par)
    if nil(gp)
        return gp
    end
    if par == left(gp)
        return right(gp)
    end
    left(gp)
end
function sibling(n::RBTreeNode, par::RBTreeNode)
    if n == left(par)
        return right(par)
    else
        return left(par)
    end
end

function del{K,V}(tree::RBTree{K,V}, key::K)
    n = find(tree, convert(K,key))
    _delete_node(tree, n)
end

function _delete_node(tree::RBTree, n::RBTreeNode)
    if nil(left(n)) && nil(right(n))
        _delete_leaf(tree, n)
    elseif nil(left(n)) || nil(right(n))
        _delete_node_with_one_child(tree, n)
    else
        BST_delete_node_with_two_childs(tree, n)
    end
end

function BST_delete_leaf{K,V}(tree::RBTree{K,V}, n::RBTreeNode{K,V})
    if left(parent(n)) == n
        set_left(parent(n), Nil{K,V}()) # TODO: rather remove left?
    else
        set_right(parent(n), Nil{K,V}()) # TODO: rather remove right
    end
    tree.nitems -= 1
end
function BST_delete_node_with_one_child{K,V}(tree::RBTree{K,V}, n::RBTreeNode{K,V})
    child = nil(left(n)) ? right(n) : left(n)

    replace_child(parent(n), n, child)
    set_parent(child, parent(n))
    tree.nitems -= 1
end
function BST_delete_node_with_two_childs{K,V}(tree::RBTree{K,V}, n::RBTreeNode{K,V})
    # child = rand(1)[1] < 0.5 ? _smallest_succ(right(n)) : _biggest_succ(left(n)) # TODO
    child = _smallest_succ(right(n))

    set_key(n, key(child))
    set_value(n, value(child))
    _delete_node(tree, child)
end

_delete_leaf(tree::RBTree, n::RedNode) = BST_delete_leaf(tree, n)
function _delete_leaf(tree::RBTree, n::BlackNode)
    _delete_case1(tree, n)
    BST_delete_leaf(tree, n)
end

function _delete_node_with_one_child{K,V}(tree::RBTree{K,V}, M::RBTreeNode{K,V})
    C = nil(left(M)) ? right(M) : left(M)

    make_black(tree, C)
    BST_delete_node_with_one_child(tree, M)
end

_delete_case1(tree::RBTree, n::RBTreeNode) = _delete_case2(tree, n, parent(n), sibling(n, parent(n)))
_delete_case2(::RBTree, n::RBTreeNode, par::Nil) = Nothing # No parent: Done

function _delete_case2(tree::RBTree, n::RBTreeNode, par::RBTreeNode, s::RedNode) # Sibling is red
    par = make_red(tree, par)
    make_black(tree, s)
    if left(par) == n
        _rotate_left(tree, par)
    else
        _rotate_right(tree, par)
    end
    s = sibling(n, par)
    _delete_case3(tree, n, par, s, left(s))
end
_delete_case2(tree::RBTree, n::RBTreeNode, par::RBTreeNode, s::RBTreeNode) = _delete_case3(tree, n, par, s, left(s))

function _delete_case3(tree::RBTree, n::RBTreeNode, par::BlackNode, s::NotRed, ls)
    @assert(left(s) == ls)
    @assert (nil(s) || color(s) == BLACK)
    if ((nil(left(s)) || color(left(s)) == BLACK) &&
        (nil(right(s)) || color(right(s)) == BLACK)
    )
        make_red(tree, s)
        _delete_case1(tree, parent(n))
    else
        _delete_case5(tree, n, par, s)
    end
end
function _delete_case3(tree::RBTree, n::RBTreeNode, par::RedNode, s::NotRed, ls)
# function _delete_case3(tree::RBTree, n::RBTreeNode, par::RedNode, s::NotRed, ls::NotRed)
    @assert(left(s) == ls)
    @assert(nil(s) || color(s) == BLACK)
    # @assert(nil(ls) || color(ls) == BLACK)
    if ((nil(left(s)) || color(left(s)) == BLACK) &&
        (nil(right(s)) || color(right(s)) == BLACK)
    )
        print("key(n): ", key(n), "\n")
        make_red(tree, s)
        make_black(tree, parent(n))
    else
        _delete_case5(tree, n, par, s)
    end
end
function _delete_case3(tree::RBTree, n::RBTreeNode, par::RedNode, s::RBTreeNode) 

    _delete_case5(tree, n, par, s)
end

_delete_case5(tree::RBTree, n::RBTreeNode, par::RBTreeNode, s::RedNode) = _delete_case6(tree, n, par, s)
function _delete_case5(tree::RBTree, n::RBTreeNode, par::RBTreeNode, s::RBTreeNode)
    if (n == left(par) &&
        (nil(right(s)) || color(right(s)) == BLACK) &&
        (!nil(left(s)) && color(left(s)) == RED))
            s = make_red(tree, s)
            make_black(tree, left(s))
            _rotate_right(tree, s)
    elseif (n == right(par) &&
        (nil(left(s)) || color(left(s)) == BLACK) &&
        (!nil(right(s)) && color(right(s)) == RED))
            s = make_red(tree, s)
            make_black(tree, right(s))
            _rotate_left(tree, s)
    end
    _delete_case6(tree, n, par, sibling(n,par))
end
function _delete_case6(tree::RBTree, n::RBTreeNode, par::BlackNode, s::RBTreeNode)
    s = make_black(tree, s)
    _delete_case7(tree, n, par, s)
end
function _delete_case6(tree::RBTree, n::RBTreeNode, par::RedNode, s::RBTreeNode)
    s = make_red(tree, s)
    _delete_case7(tree, n, par, s)
end

function _delete_case7(tree::RBTree, n::RBTreeNode, par::RBTreeNode, s::RBTreeNode)
    par = make_black(tree, par)

    if n == left(par)
        make_black(tree, right(s))
        _rotate_left(tree, par)
    else
        make_black(tree, left(s))
        _rotate_right(tree, par)
    end
end

_biggest_succ(n::RBTreeNode) = nil(right(n)) ? n : _biggest_succ(right(n))
_smallest_succ(tree::Nil) = tree
_smallest_succ(n::RBTreeNode) = nil(left(n)) ? n : _smallest_succ(left(n))

function _rotate_left(tree::RBTree, P::RBTreeNode)
    R = right(P)
    GP = parent(P)
    B = left(R)

    # Link P as new l of R
    set_parent(P, R)
    set_left(R, P)

    # Link R as new child of GP
    set_parent(R, GP)
    if !nil(GP)
        if left(GP) == P
            set_left(GP, R)
        else
            set_right(GP, R)
        end
    else # No grand-parent -> P was the root
        tree.root = make_black(tree, R)
    end

    # Link P as new parent of left(R)
    set_right(P, B)
    if !nil(B)
        set_parent(B, P)
    end
end

function _rotate_right(tree::RBTree, P::RBTreeNode)
    L = left(P)
    GP = parent(P)
    B = right(L)

    # Link P as new r of L
    set_parent(P, L)
    set_right(L, P)

    # Link L as new child of GP
    set_parent(L, GP)
    if !nil(GP)
        if left(GP) == P
            set_left(GP, L)
        else
            set_right(GP, L)
        end
    else # No grand-parent -> P was the root
        tree.root = make_black(tree, L)
    end

    # Link P as new parent of right(L)
    set_left(P, B)
    if !nil(B)
        set_parent(B, P)
    end
end


_just_inserted(tree::RBTree, n::RBTreeNode, ::Nil) = make_black(tree, n) # Case 1: root node
_just_inserted(::RBTree, ::RBTreeNode, ::BlackNode) = Nothing  # Parent is black. All done
_just_inserted(tree::RBTree, n::RBTreeNode, par::RedNode) = _insert_case3(tree, n, par, uncle(par))

function _insert_case3(tree::RBTree, n::RBTreeNode, par::RedNode, u::RedNode) # Uncle is red
    par = make_black(tree, par)
    u = make_black(tree, u)
    gp = make_red(tree, parent(par))
    _just_inserted(tree, gp, parent(gp))
end
# Uncle is nil or black
_insert_case3(tree::RBTree, n::RBTreeNode, par::RedNode, u::RBTreeNode) = _insert_case4(tree, n, par)

function _insert_case4(tree::RBTree, n::RBTreeNode, par::RedNode)
    gp = parent(par)

    if (n == right(par) && parent(n) == left(gp))
        _rotate_left(tree, parent(n))
        n = left(n)
        par = parent(n)
        gp = parent(par)
    elseif n == left(par) && parent(n) == right(gp)
        _rotate_right(tree, parent(n))
        n = right(n)
        par = parent(n)
        gp = parent(par)
    end

    par = make_black(tree, par)
    gp = make_red(tree, gp)
    if n == left(par)
        _rotate_right(tree, gp)
    else
        _rotate_left(tree, gp)
    end
end

function find{K,V}(tree::RBTree{K,V}, k::K)
    return find(tree.root, k)
end
find{K,V}(n::Nil{K,V}, key::K) = throw(KeyError(key))
function find{K,V}(n::RBTreeNode{K,V}, k::K)
    if (key(n) == k)
        return n
    end
    if k < key(n)
        return find(left(n), k)
    else
        return find(right(n), k)
    end
end

## Debug functions {{{
dlabel(t::Nil) = "EmptyNode"
dlabel(t::RBTreeNode) = strcat('"', string(key(t)), "_", string(value(t)), '"')

function dot(t::RBTree)
    nodes, edges = dot(t.root)
    join(append_any(["digraph {",
        "EmptyNode [label=\"\", shape=point, clr=\"0 0 0\"];"],
        nodes,
        edges,
        ["}"]), "\n")
end
function dot(t::Nil)
    [], []
end
function dot(t::RBTreeNode)
    lnodes, ledges = dot(left(t))
    rnodes, redges = dot(right(t))
    mylabel = dlabel(t)

    cname = "#ff0000"
    if color(t) == BLACK
        cname = "#b0b0b0"
    end

    edges = String[]
    if !nil(left(t))
        push(edges, strcat("$mylabel -> ", dlabel(left(t)), "[color=yellow];"))
    end
    if !nil(right(t))
        push(edges, strcat("$mylabel -> ", dlabel(right(t)), "[color=red];"))
    end
    if !nil(parent(t))
        push(edges, strcat("$mylabel -> ", dlabel(parent(t)), "[color=green];"))
    end

    append_any(["$mylabel [shape=box, color=\"$cname\", style=filled];"], lnodes, rnodes),
    append_any(ledges, redges, edges)
end
## End Debug functions }}}
function assign{K,V}(tree::RBTree{K,V}, val, key)
    key = convert(K, key)
    val = convert(V, val)
    _insert_or_modify(tree, val, key)
    val
end

function ref{K,V}(tree::RBTree{K,V}, key)
    value(find(tree, convert(K,key)))
end
start(Nil) = Nothing
done(Nil, ::Any) = true

function start{K,V}(tree::RBTree{K,V})
    return _smallest_succ(tree.root)
end
function done(::RBTree, state::RBTreeNode)
    nil(state)
end
function next{K,V}(tree::RBTree{K,V}, state::RBTreeNode{K,V})
    nn = state
    if nil(right(nn))
        while !nil(parent(nn)) && key(state) >= key(nn)
            nn = parent(nn)
        end
    else
        nn = _smallest_succ(right(nn))
    end
    return (key(state), value(state)), key(nn) > key(state) ? nn : Nil{K,V}()
end

function has{K,V}(tree::RBTree{K,V}, key::K)
    try
        find(tree, convert(K, key))
        true
    catch
        false
    end
end
length{K,V}(tree::RBTree{K,V}) = tree.nitems

## TODO: promoting rules?