//! returns: 0
//! Red-black tree.
//! Implement a red-black tree (balanced BST) using a stack-allocated node pool.

constant POOL_SIZE: u32 = 128;
constant NIL: u32 = 0;

constant RED: u32 = 0;
constant BLACK: u32 = 1;

record RBNode: Copy {
    key: i32,
    color: u32,
    left: u32,
    right: u32,
    parent: u32,
}

record RBTree: Copy {
    pool: *unsafe mut [RBNode],
    poolNext: u32,
    root: u32,
    inorder: *unsafe mut [i32],
    inorderCount: u32,
}

unsafe fn allocNode(t: *unsafe mut RBTree, key: i32) -> u32 {
    let idx: u32 = t.poolNext;
    set t.poolNext += 1;
    set t.pool[idx] = RBNode { key, color: RED, left: NIL, right: NIL, parent: NIL };
    return idx;
}

unsafe fn rotateLeft(t: *unsafe mut RBTree, x: u32) {
    let y: u32 = t.pool[x].right;
    set t.pool[x].right = t.pool[y].left;
    if t.pool[y].left <> NIL {
        set t.pool[t.pool[y].left].parent = x;
    }
    set t.pool[y].parent = t.pool[x].parent;
    if t.pool[x].parent == NIL {
        set t.root = y;
    } else if x == t.pool[t.pool[x].parent].left {
        set t.pool[t.pool[x].parent].left = y;
    } else {
        set t.pool[t.pool[x].parent].right = y;
    }
    set t.pool[y].left = x;
    set t.pool[x].parent = y;
}

unsafe fn rotateRight(t: *unsafe mut RBTree, x: u32) {
    let y: u32 = t.pool[x].left;
    set t.pool[x].left = t.pool[y].right;
    if t.pool[y].right <> NIL {
        set t.pool[t.pool[y].right].parent = x;
    }
    set t.pool[y].parent = t.pool[x].parent;
    if t.pool[x].parent == NIL {
        set t.root = y;
    } else if x == t.pool[t.pool[x].parent].right {
        set t.pool[t.pool[x].parent].right = y;
    } else {
        set t.pool[t.pool[x].parent].left = y;
    }
    set t.pool[y].right = x;
    set t.pool[x].parent = y;
}

unsafe fn insertFixup(t: *unsafe mut RBTree, zArg: u32) {
    let mut z: u32 = zArg;
    while t.pool[t.pool[z].parent].color == RED {
        if t.pool[z].parent == t.pool[t.pool[t.pool[z].parent].parent].left {
            let y: u32 = t.pool[t.pool[t.pool[z].parent].parent].right;
            if t.pool[y].color == RED {
                set t.pool[t.pool[z].parent].color = BLACK;
                set t.pool[y].color = BLACK;
                set t.pool[t.pool[t.pool[z].parent].parent].color = RED;
                set z = t.pool[t.pool[z].parent].parent;
            } else {
                if z == t.pool[t.pool[z].parent].right {
                    set z = t.pool[z].parent;
                    rotateLeft(t, z);
                }
                set t.pool[t.pool[z].parent].color = BLACK;
                set t.pool[t.pool[t.pool[z].parent].parent].color = RED;
                rotateRight(t, t.pool[t.pool[z].parent].parent);
            }
        } else {
            let y: u32 = t.pool[t.pool[t.pool[z].parent].parent].left;
            if t.pool[y].color == RED {
                set t.pool[t.pool[z].parent].color = BLACK;
                set t.pool[y].color = BLACK;
                set t.pool[t.pool[t.pool[z].parent].parent].color = RED;
                set z = t.pool[t.pool[z].parent].parent;
            } else {
                if z == t.pool[t.pool[z].parent].left {
                    set z = t.pool[z].parent;
                    rotateRight(t, z);
                }
                set t.pool[t.pool[z].parent].color = BLACK;
                set t.pool[t.pool[t.pool[z].parent].parent].color = RED;
                rotateLeft(t, t.pool[t.pool[z].parent].parent);
            }
        }
    }
    set t.pool[t.root].color = BLACK;
}

unsafe fn insert(t: *unsafe mut RBTree, key: i32) {
    let z: u32 = allocNode(t, key);
    let mut y: u32 = NIL;
    let mut x: u32 = t.root;

    while x <> NIL {
        set y = x;
        if key < t.pool[x].key {
            set x = t.pool[x].left;
        } else {
            set x = t.pool[x].right;
        }
    }

    set t.pool[z].parent = y;
    if y == NIL {
        set t.root = z;
    } else if key < t.pool[y].key {
        set t.pool[y].left = z;
    } else {
        set t.pool[y].right = z;
    }

    insertFixup(t, z);
}

unsafe fn search(t: *unsafe RBTree, key: i32) -> bool {
    let mut x: u32 = t.root;
    while x <> NIL {
        if key == t.pool[x].key {
            return true;
        } else if key < t.pool[x].key {
            set x = t.pool[x].left;
        } else {
            set x = t.pool[x].right;
        }
    }
    return false;
}

unsafe fn inorderWalk(t: *unsafe mut RBTree, x: u32) {
    if x == NIL {
        return;
    }
    inorderWalk(t, t.pool[x].left);
    set t.inorder[t.inorderCount] = t.pool[x].key;
    set t.inorderCount += 1;
    inorderWalk(t, t.pool[x].right);
}

unsafe fn countNodes(t: *unsafe RBTree, x: u32) -> u32 {
    if x == NIL {
        return 0;
    }
    return 1 + countNodes(t, t.pool[x].left) + countNodes(t, t.pool[x].right);
}

unsafe fn blackHeight(t: *unsafe RBTree, x: u32) -> i32 {
    if x == NIL {
        return 1;
    }
    let leftBH: i32 = blackHeight(t, t.pool[x].left);
    let rightBH: i32 = blackHeight(t, t.pool[x].right);

    if leftBH == -1 or rightBH == -1 {
        return -1;
    }
    if leftBH <> rightBH {
        return -1;
    }

    if t.pool[x].color == BLACK {
        return leftBH + 1;
    }
    return leftBH;
}

unsafe fn noRedRed(t: *unsafe RBTree, x: u32) -> bool {
    if x == NIL {
        return true;
    }
    if t.pool[x].color == RED {
        if t.pool[t.pool[x].left].color == RED {
            return false;
        }
        if t.pool[t.pool[x].right].color == RED {
            return false;
        }
    }
    if not noRedRed(t, t.pool[x].left) {
        return false;
    }
    return noRedRed(t, t.pool[x].right);
}

unsafe fn resetTree(t: *unsafe mut RBTree) {
    let mut i: u32 = 0;
    while i < POOL_SIZE {
        set t.pool[i] = RBNode { key: 0, color: BLACK, left: NIL, right: NIL, parent: NIL };
        set i += 1;
    }
    set t.poolNext = 1;
    set t.root = NIL;
    set t.inorderCount = 0;
}

unsafe fn testAscending(t: *unsafe mut RBTree) -> i32 {
    resetTree(t);

    let mut i: i32 = 0;
    while i < 32 {
        insert(t, i);
        set i += 1;
    }

    assert countNodes(t, t.root) == 32;
    assert t.pool[t.root].color == BLACK;

    let bh: i32 = blackHeight(t, t.root);
    assert bh <> -1;

    assert noRedRed(t, t.root);

    set t.inorderCount = 0;
    inorderWalk(t, t.root);
    assert t.inorderCount == 32;
    let mut j: u32 = 0;
    while j < 32 {
        assert t.inorder[j] == j as i32;
        set j += 1;
    }

    return 0;
}

unsafe fn testDescending(t: *unsafe mut RBTree) -> i32 {
    resetTree(t);

    let mut i: i32 = 31;
    while i >= 0 {
        insert(t, i);
        if i == 0 {
            break;
        }
        set i -= 1;
    }

    assert countNodes(t, t.root) == 32;
    assert blackHeight(t, t.root) <> -1;
    assert noRedRed(t, t.root);

    set t.inorderCount = 0;
    inorderWalk(t, t.root);
    let mut j: u32 = 0;
    while j < 32 {
        assert t.inorder[j] == j as i32;
        set j += 1;
    }

    return 0;
}

unsafe fn testRandom(t: *unsafe mut RBTree) -> i32 {
    resetTree(t);

    let mut inserted: [bool; 48] = [false; 48];
    let mut seed: u32 = 42;
    let mut count: u32 = 0;

    while count < 48 {
        set seed = (seed * 1103515245 + 12345) & 0x7FFFFFFF;
        let val: u32 = seed % 48;
        if not inserted[val] {
            insert(t, val as i32);
            set inserted[val] = true;
            set count += 1;
        }
    }

    assert countNodes(t, t.root) == 48;
    assert blackHeight(t, t.root) <> -1;
    assert noRedRed(t, t.root);

    let mut i: i32 = 0;
    while i < 48 {
        assert search(t, i);
        set i += 1;
    }

    assert not search(t, -1);
    assert not search(t, 48);
    assert not search(t, 100);

    set t.inorderCount = 0;
    inorderWalk(t, t.root);
    assert t.inorderCount == 48;
    let mut j: u32 = 0;
    while j < 48 {
        assert t.inorder[j] == j as i32;
        set j += 1;
    }

    return 0;
}

unsafe fn testHeight(t: *unsafe mut RBTree) -> i32 {
    resetTree(t);

    let mut i: i32 = 0;
    while i < 63 {
        insert(t, i);
        set i += 1;
    }

    let bh: i32 = blackHeight(t, t.root);
    assert bh >= 3;
    assert bh <= 7;

    return 0;
}

@default unsafe fn main() -> i32 {
    let mut pool: [RBNode; 128] = [RBNode { key: 0, color: 1, left: 0, right: 0, parent: 0 }; 128];
    let mut inorder: [i32; 128] = [0; 128];

    let mut t: RBTree = RBTree {
        pool: &mut pool[..],
        poolNext: 1,
        root: NIL,
        inorder: &mut inorder[..],
        inorderCount: 0,
    };

    let r1: i32 = testAscending(&mut t);
    if r1 <> 0 { return 10 + r1; }

    let r2: i32 = testDescending(&mut t);
    if r2 <> 0 { return 20 + r2; }

    let r3: i32 = testRandom(&mut t);
    if r3 <> 0 { return 30 + r3; }

    let r4: i32 = testHeight(&mut t);
    if r4 <> 0 { return 50 + r4; }

    return 0;
}