lib/rbtree.c


1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84

#include "rbtree.h"

static inline int
s_direction(const struct rbtree_node *node)
{
    return node == node->parent->children[0] ? 0 : 1;
}

static struct rbtree_node*
s_rotate(struct rbtree_node *sub, int dir)
{
    int opp = 1 - dir;
    struct rbtree_node *sub_parent = sub->parent;
    struct rbtree_node *new_root = sub_parent->children[opp];
    struct rbtree_node *new_child = new_root->children[dir];

    sub_parent->children[opp] = new_child;
    if(new_child) new_child->parent = sub;

    new_root->children[dir] = sub;
    new_root->parent = sub_parent;
    sub->parent = new_root;

    if(sub_parent)
        sub_parent->children[s_direction(sub)] = new_root;
    return new_root;
}

void
rbtree_insert(struct rbtree_node *root, struct rbtree_node *node)
{
    struct rbtree_node *parent = root;
    int direction;
    while(1) {
        direction = parent->get_key(parent) < node->get_key(node);
        if(parent->children[direction] == 0) break;
        parent = parent->children[direction];
    }

    node->color = 1;
    parent->children[direction] = node;
    node->parent = parent;

    //Rebalance the tree
    do {
        //Tree is already balanced, return
        if(parent->color == 0) return;
        struct rbtree_node *gparent = parent->parent;

        //Parent is the root node, set parent to black.
        if(!gparent) {
            parent->color = 0;
            return;
        }

        direction = s_direction(parent);
        struct rbtree_node *uncle = gparent->children[1-direction];
        if(!uncle || uncle->color == 0) {
            //If the uncle node DNE or is black, the tree has to be rotated.
            if(node == parent->children[1-direction]) {
                //Rotate parent node so that node is the new parent.
                s_rotate(parent, direction);
                node = parent;
                parent = gparent->children[direction];
            }

            //Rotate the parent node so that the gparent is the new parent.
            s_rotate(gparent, 1-direction);
            parent->color = 0;
            gparent->color = 1;
            return;
        }

        //Both the parent and uncle are red. Push up the red and iterate up.
        parent->color = 0;
        uncle->color = 0;
        gparent->color = 1;
        node = gparent;
    } while((parent = node->parent));

    //Tree is balanced now.
}