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

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

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

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

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

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

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

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

fn insertFixup(t: *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 {
                t.pool[t.pool[z].parent].color = BLACK;
                t.pool[y].color = BLACK;
                t.pool[t.pool[t.pool[z].parent].parent].color = RED;
                z = t.pool[t.pool[z].parent].parent;
            } else {
                if z == t.pool[t.pool[z].parent].right {
                    z = t.pool[z].parent;
                    rotateLeft(t, z);
                }
                t.pool[t.pool[z].parent].color = BLACK;
                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 {
                t.pool[t.pool[z].parent].color = BLACK;
                t.pool[y].color = BLACK;
                t.pool[t.pool[t.pool[z].parent].parent].color = RED;
                z = t.pool[t.pool[z].parent].parent;
            } else {
                if z == t.pool[t.pool[z].parent].left {
                    z = t.pool[z].parent;
                    rotateRight(t, z);
                }
                t.pool[t.pool[z].parent].color = BLACK;
                t.pool[t.pool[t.pool[z].parent].parent].color = RED;
                rotateLeft(t, t.pool[t.pool[z].parent].parent);
            }
        }
    }
    t.pool[t.root].color = BLACK;
}

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

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

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

    insertFixup(t, z);
}

fn search(t: *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 {
            x = t.pool[x].left;
        } else {
            x = t.pool[x].right;
        }
    }
    return false;
}

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

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

fn blackHeight(t: *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;
}

fn noRedRed(t: *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);
}

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

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

    let mut i: i32 = 0;
    while i < 32 {
        insert(t, i);
        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);

    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;
        j += 1;
    }

    return 0;
}

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

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

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

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

    return 0;
}

fn testRandom(t: *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 {
        seed = (seed * 1103515245 + 12345) & 0x7FFFFFFF;
        let val: u32 = seed % 48;
        if not inserted[val] {
            insert(t, val as i32);
            inserted[val] = true;
            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);
        i += 1;
    }

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

    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;
        j += 1;
    }

    return 0;
}

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

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

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

    return 0;
}

@default 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;
}