rbtree.ooc

RBNodeColor: enum {
    black = 0
    red = 1
}

RBNode: class <K, V> {
    
    
    color := RBNodeColor black
    left, right, parent: This<K, V>
    key: K
    value: V
    
    key: func() -> K { key }
    value: func() -> V { value }
    
    init: func(=key, =value, =parent) {}
    
    grandparent: func -> This<K, V> {
        if (parent)
            return parent parent
        return null
    }
    
    uncle: func -> This<K, V> {
        gp := grandparent()
        if (gp == null)
            return null
        
        if (gp left == this)
            return gp right
        else
            return gp left
    }
    
    sibling: func -> This<K, V> {
        if (parent) {
            if (isLeft?())
                return parent right
            else
                return parent left
        }
        return null
    }
    
    isLeft?: func -> Bool { (parent && parent left == this) }
    isRight?: func -> Bool { (parent && parent right == this) }
    
    isBlack?: func -> Bool { (color == RBNodeColor black) }
    isRed?: func -> Bool { (color == RBNodeColor red) }
    
    hasParent?: func -> Bool { (parent != null) }
    isLeaf?: func -> Bool { (!left && !right) }
    hasLeft?: func -> Bool { (left != null) }
    hasRight?: func -> Bool { (right != null) }
    
    walk: func (fn: Func(RBNode<K,V>)) {        
        if (left) {
            left walk(fn)
        }
        fn(this)
        if (right) {
            right walk(fn)
        }
    }
    
}

RBTree: class <K, V> {
    root: RBNode<K, V> = null
    comparator: Func<K> (K, K) -> Int
    
    init: func(comp: Func<K> (K, K) -> Int) {
        comparator = comp
    }
    
    replaceNode: func (oldn, newn: RBNode<K, V>) {
        if (!oldn hasParent?()) {
            root = newn
        } else {
            if (oldn isLeft?())
                oldn parent left = newn
            else
                oldn parent right = newn
        }
        if (newn)
            newn parent = oldn parent
    }
    
    contains: func (key: K) -> Bool {
        return __findNode(key) ? true : false
    }
    
    get: func (key: K) -> V {
        node := __findNode(key)
        if (node)
            return node value
        else
            return null
    }
    
    put: func (key: K, value: V) {
        node := RBNode<K, V> new(key, value, null)
        parent: RBNode<K, V> = null
        follower := root
        cmp: Int = 0
        while (follower) {
            parent = follower
            cmp = comparator(key as K, follower key as K)
            "put - comparison result: %d" format(cmp) println()
            if (cmp < 0) {
                follower = follower left
            } else if (cmp > 0) {
                follower = follower right
            } else {
                follower value = value
                return
            }
        }
        
        if (parent) {
            node parent = parent
            cmp := comparator(node key as K, parent key as K)
            if (cmp < 0)
                parent left = node
            else
                parent right = node
        } else {
            root = node
        }
        
        node color = RBNodeColor red
        
        while (node) {
            // mostly based on the wikipedia page and this page:
            // http://www.cs.auckland.ac.nz/software/AlgAnim/red_black.html
            if (node parent == null) {
                // case 1
                node color = RBNodeColor black
                return
            } else if (node parent isBlack?()) {
                // case 2
                return
            } else {
                // case 3
                uncle := node uncle()

                if (uncle && uncle isRed?()) {
                    node parent color = RBNodeColor black
                    uncle color = RBNodeColor black
                    gpar := node grandparent()
                    gpar color = RBNodeColor red
                    node = gpar
                    continue
                } else {
                    // case 4
                    gpar := node grandparent()
                    
                    if (node isRight?() && node parent isLeft?()) {
                        __rotateLeft(node parent)
                        node = node left
                    } else if (node isLeft?() && node parent isRight?()) {
                        __rotateRight(node parent)
                        node = node right
                    }
                
                    // case 5
                    gpar = node grandparent()
                    node parent color = RBNodeColor black
                    gpar color = RBNodeColor red
                    if (node isLeft?() && node parent isLeft?()) {
                        __rotateRight(gpar)
                    } else if (node isRight?() && node parent isRight?()) {
                        __rotateLeft(gpar)
                    }
                }
            }
            return
        }
    }
    
    remove: func (key: K) -> Bool {
        n := __findNode(key)
        if (n) {
            if (!n isLeaf?()) {
                // Can't use 'n right isLeaf?()' here because the right node might be null
                child := n hasLeft?() ? n left : n right
                replaceNode(n, child)
                if (n isBlack?()) {
                    if (child isRed?())
                        child color = RBNodeColor black
                    else
                        __deleteNode(child)
                }
            } else {
                replaceNode(n, null)
            }
            return true
        }
        return false
    }
    
    __deleteNode: func (n: RBNode<K, V>) {
        if (n parent != null) {
            // case 2
            s := n sibling()
            if (s isRed?()) {
                n parent color = RBNodeColor red
                s color = RBNodeColor black
                if (n isLeft?())
                    __rotateLeft(n parent)
                else
                    __rotateRight(n parent)
            }
            // case 3
            if ((n parent isBlack?()) &&
                (s isBlack?()) &&
                (s left isBlack?()) &&
                (s right isBlack?())) {
                s color = RBNodeColor red
                __deleteNode(n parent)
            } else {
                // case 4
                if ((n parent isRed?()) &&
                    (s isBlack?()) &&
                    (s left isBlack?()) &&
                    (s right isBlack?())) {
                    s color = RBNodeColor red
                    n parent color = RBNodeColor black
                } else {
                    // case 5
                    if (s isBlack?()) {
                        if ((n isLeft?()) &&
                            (s right isBlack?()) &&
                            (s left isRed?())) {
                            s color = RBNodeColor red
                            s left color = RBNodeColor black
                            __rotateRight(s)
                        } else if ((n isRight?()) &&
                            (s left isBlack?()) &&
                            (s right isRed?())) {
                            s color = RBNodeColor red
                            s right color = RBNodeColor black
                            __rotateLeft(s)
                        }
                    }
                    // case 6
                    s color = n parent color
                    n parent color = RBNodeColor black
                    if (n isLeft?()) {
                        s right color = RBNodeColor black
                        __rotateLeft(n parent)
                    } else {
                        s left color = RBNodeColor black
                        __rotateRight(n parent)
                    }
                }
            }
        }
    }
    
    __findNode: func (key: K) -> RBNode<K, V> {
        node := root
        while (node) {
            cmp := comparator(key as K, node key as K)
            "findNode - comparison result: %d" format(cmp) println()
            if (cmp == 0)
                return node
            else if (cmp < 0)
                node = node left
            else
                node = node right
        }
        return null
    }
    
    __rotateLeft: func (node: RBNode<K, V>) {
        right := node right
        node right = right left
        if (node right)
            node right parent = node
        
        right parent = node parent
        
        if (node parent == null) {
            root = right
        } else {
            if (node isLeft?()) {
                node parent left = right
            } else { // right
                node parent right = right
            }
        }
        
        right left = node
        node parent = right
    }
    
    __rotateRight: func (node: RBNode<K, V>) {
        left := node left
        node left = left right
        if (node left)
            node left parent = node
        
        left parent = node parent
        
        if (node parent == null) {
            root = left
        } else {
            if (node isLeft?()) {
                node parent left = left
            } else { // right
                node parent right = left
            }
        }
        
        left right = node
        node parent = left
    }
}

assert: func (test: Bool) {
    Exception new("FAILURE") throw()
}

walkFn: func(n: RBNode<Int, Int>) {
    color: String = n isBlack?() ? "black" : "red"
    if (n isLeft?())
        "left: %d %s" format(n key(), color) println()
    else if (n isRight?())
        "right: %d %s" format(n key(), color) println()
    else
        "root: %d %s" format(n key(), color) println()
    "value: %d" format(n value()) println()
}

compareFn: func<K> (left, right: K) -> Int {
    result: Int = (left - right)
    "comparing %d and %d - result: %d" format(left, right, result) println()
    return result
}

main: func {
    woop := RBTree<Int, Int> new(compareFn)
    woop put(1, 100)
    woop put(2, 200)
    woop put(3, 300)
    woop put(4, 400)
    woop put(5, 500)
    woop put(6, 600)
    woop put(7, 700)
    woop put(8, 800)
    woop put(9, 900)
    woop put(10, 1000)
    woop put(7, 12345)
    woop root walk(walkFn)
    "%d" format(woop get(6)) println()
    
    woop remove(6)
    "value at 6: %d" format(woop get(6)) println()
    "contains 6: %d" format(woop contains(6)) println()
}