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
    }
    
    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) }
    
    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) -> Int
    
    init: func(comp: Func (K, K) -> Int) {
        comparator = comp
    }
    
    get: func (key: K) -> V {
        node := root
        while (node) {
            cmp := comparator(key as K, node key as K)
            "get - comparison result: %d" format(cmp) println()
            if (cmp == 0)
                return node value
            else if (cmp < 0)
                node = node left
            else
                node = node right
        }
        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
        }
    }
    
    __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
    }
}

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 (left, right: 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(5)) println()
}