Intrusive red-black trees in C

rb_tree.c

#include "rb_tree.h"

#define PARENT_BIT 1
#define COLOR_BIT 2
#define PARENT_MASK (PARENT_BIT | COLOR_BIT)

#define rb_is_red(n) ((n) != (void *)0 && ((n)->parent & COLOR_BIT))
#define rb_is_black(n) ((n) == (void *)0 || !((n)->parent & COLOR_BIT))
#define rb_set_red(n) ((n)->parent |= COLOR_BIT)
#define rb_set_black(n) ((n)->parent &= ~COLOR_BIT)

#define rb_links(n) ((struct rb_node **)(n))
#define rb_parent(n) ((struct rb_node *)((n)->parent & ~PARENT_MASK))
#define rb_set_parent(n, p) \
    ((n)->parent = ((uintptr_t)(p) | ((n)->parent & PARENT_MASK)))

static void rb_link_fixup(struct rb_node *n);
static void rb_unlink_fixup(struct rb_node *p, int dir);
static void rb_replace(struct rb_node *target, struct rb_node *node);
static void rb_rotate(struct rb_node *p, int dir);

void rb_link(struct rb_node *node, struct rb_node **plink)
{ 
    *plink = node;
    node->left = node->right = (void *)0;
    node->parent = (uintptr_t)&plink[-((uintptr_t)plink[1] & PARENT_BIT)];
    node->parent |= PARENT_BIT;
    rb_set_red(node);
    rb_link_fixup(node);
}

static void rb_link_fixup(struct rb_node *n)
{
    struct rb_node *g, *p, *u;
    while ((p = rb_parent(n)) && (p->parent & PARENT_MASK)) {
        /* Red violation while parent is red */
        if (rb_is_black(p))
            return;

        /* Grandparent becomes red */
        g = rb_parent(p);
        rb_set_red(g);

        /* (Double-)rotate if uncle is black */
        u = rb_links(g)[p == g->left];
        if (rb_is_black(u)) {
            int dir = (n != p->left);
            if (p != rb_links(g)[dir])
                rb_rotate(p, (dir = !dir));
            rb_set_black(rb_links(g)[dir]);
            rb_rotate(g, !dir);
            return;
        }

        /* Recolor (uncle is red) */
        rb_set_black(u);
        rb_set_black(p);
        n = g;
    }

    /* New root */
    rb_set_black(n);
}

void rb_unlink(struct rb_node *node)
{
    struct rb_node *n, *p, *c;

    /*
     * Replace a node with two children with its in-order successor to simplify
     * balancing (the successor has at most 1 child). Alternatively, replace the
     * node with its predecessor. We try to be clever and choose the replacement
     * node based on which branch has the lowest red node. TODO: benchmark plz
     */
    if (node->left && node->right) {
        struct rb_node *pred = node->left;
        struct rb_node *succ = node->right;
        int plen = rb_is_black(pred);
        int slen = rb_is_black(succ);
        while (pred->right || succ->left) {
            if (pred->right) {
                pred = pred->right;
                plen = rb_is_red(pred) ? 0 : plen+1;
            }
            if (succ->left) {
                succ = succ->left;
                slen = rb_is_red(succ) ? 0 : slen+1;
            }
        }
        n = (plen < slen) ? pred : succ;
        c = rb_links(n)[n == succ];
    } else {
        n = node;
        c = rb_links(n)[!n->left];
    }

    p = rb_parent(n);
    if ((rb_links(p)[p->left != n] = c)) {
        /* Unlink a node with only one child */
        c->parent = n->parent;
    } else if (rb_is_black(n)) {
        /* Unlink a black node with no children */
        rb_unlink_fixup(p, !!p->left);
    }

    /* Restore unlinked successor/predecessor */
    if (node != n)
        rb_replace(node, n);
}

static void rb_unlink_fixup(struct rb_node *p, int dir)
{
    while ((p->parent & PARENT_MASK)) {
        struct rb_node *s = rb_links(p)[!dir];

        if (rb_is_red(s)) {
            rb_rotate(p, dir);
            rb_set_black(s);
            rb_set_red(p);
            s = rb_links(p)[!dir];
        } 

        if (rb_is_black(s->left) && rb_is_black(s->right)) {
            rb_set_red(s);
            if (rb_is_black(p)) {
                struct rb_node *g = rb_parent(p);
                dir = (p != g->left);
                p = g;
                continue;
            }
        } else {
            if (rb_is_black(rb_links(s)[!dir])) {
                rb_rotate(s, !dir);
                rb_set_red(s);
                s = rb_links(p)[!dir];
            }

            rb_rotate(p, dir);
            s->parent = (s->parent & ~COLOR_BIT) | (p->parent & COLOR_BIT);
            rb_set_black(rb_links(s)[!dir]);
        }

        rb_set_black(p);
        return;
    }
}

static void rb_replace(struct rb_node *target, struct rb_node *node)
{
    struct rb_node *parent = rb_parent(target);
    node->parent = target->parent;
    rb_links(parent)[parent->left != target] = node;
    if ((node->left = target->left))
        rb_set_parent(node->left, node);
    if ((node->right = target->right))
        rb_set_parent(node->right, node);
}

void rb_cache(struct rb_tree *tree, struct rb_node *node)
{
    struct rb_node *parent = rb_parent(node);
    int insert = (parent == (void *)tree)
               ? (tree->root == node)
               : (parent->left == node || parent->right == node);
    if (insert) {
        struct rb_node *leftmost = tree->cache.leftmost;
        struct rb_node *rightmost = tree->cache.rightmost;
        if (node->right == leftmost || leftmost->left == node)
            tree->cache.leftmost = node;
        if (node->left == rightmost || rightmost->right == node)
            tree->cache.rightmost = node;
        tree->cache.count++;
    } else {
        if (node == tree->cache.leftmost)
            tree->cache.leftmost = rb_next(node);
        if (node == tree->cache.rightmost)
            tree->cache.rightmost = rb_prev(node);
        tree->cache.count--;
    }
}

struct rb_node *rb_tree_iter(struct rb_tree *tree, int dir)
{
    if (!tree->cache.count) {
        struct rb_node *prev, *node = tree->root;
        for (prev = node; node; prev = node, node = rb_links(node)[dir]);
        return prev;
    }
    return !dir ? tree->cache.leftmost : tree->cache.rightmost;
}

struct rb_node *rb_node_iter(struct rb_node *node, int dir)
{
    if (!rb_links(node)[dir]) {
        struct rb_node *parent;
        while ((parent = rb_parent(node)) && (parent->parent & PARENT_MASK)) {
            if (node != rb_links(parent)[dir])
                return parent;
            node = parent;
        }
        return (void *)0;
    }
    node = rb_links(node)[dir];
    while (rb_links(node)[!dir]) node = rb_links(node)[!dir];
    return node;
}

/*
 * Binary tree rotations.
 *
 *        g                     .         g
 *       / \                    |\       / \
 *      p   u  .----------------' \     n   u
 *     / \     | left  rotation    )   / \
 *    s   n    '----------------. /   p   2
 *       / \                    |/   / \
 *      1   2                   '   s   1
 *
 *        g                     .         g
 *       / \                    |\       / \
 *      p   u  .----------------' \     n   u
 *     / \     | right rotation    )   / \
 *    n   s    '----------------. /   1   p
 *   / \                        |/       / \
 *  1   2                       '       2   s
 */
static void rb_rotate(struct rb_node *p, int dir)
{
    /* Parent links */
    struct rb_node *n = rb_links(p)[!dir];
    struct rb_node *g = rb_parent(p);
    rb_set_parent(n, g);
    rb_set_parent(p, n);

    /* Left/right links */
    if ((rb_links(p)[!dir] = rb_links(n)[dir]))
        rb_set_parent(rb_links(n)[dir], p);
    rb_links(n)[dir] = p;
    rb_links(g)[p != g->left] = n;
}

rb_tree.h

/*
 * Red-black tree (intrusive self-balancing binary search tree).
 */

#ifndef RB_TREE_H
#define RB_TREE_H

#include <stdint.h>

struct rb_node {
    struct rb_node *left;
    struct rb_node *right;
    uintptr_t parent;
} /*__attribute__((aligned(4)))*/ ;

#define rb_entry(ptr, type, member) \
    ((type *)((uintptr_t)(ptr) - (uintptr_t)&((type *)0)->member))

void rb_link(struct rb_node *node, struct rb_node **plink);
void rb_unlink(struct rb_node *node);

struct rb_tree {
    struct rb_node *root;

    struct { /* Must be zero-initialized even if caching is unused! */
        struct rb_node *leftmost;
        struct rb_node *rightmost;
        unsigned long count;
    } cache;
};

#define RB_TREE { (void *)0, { (void *)0, (void *)0, 0 } }

/* Call after each rb_link()/rb_unlink() to update tree->cache */
void rb_cache(struct rb_tree *tree, struct rb_node *node);

struct rb_node *rb_tree_iter(struct rb_tree *tree, int dir);
#define rb_first(t) (rb_tree_iter((t), 0))
#define rb_last(t) (rb_tree_iter((t), 1))

struct rb_node *rb_node_iter(struct rb_node *node, int dir);
#define rb_prev(n) (rb_node_iter((n), 0))
#define rb_next(n) (rb_node_iter((n), 1))

#endif