0x00 前言
由于红黑树具有非常重要工程实践意义，很多基础工程中都包含有红黑树的实现。对比 paho.mqtt.c / nginx / libuv / linux 中红黑树的实现发现，Linux 内核中红黑树的实现部分最为经典，本文通过对 Linux 内核中红黑树的实现进行初步分析，并利用 Linux 内核中红黑树的接口，引用《算法导论》中数据结构扩展的一般方法，对红黑树扩展来实现顺序统计树。
0x01 Linux 内核中红黑树实现分析
① 结构定义
Linux 内核的红黑树数据结构定义如下：
struct rb_node {
unsigned long __rb_parent_color;
struct rb_node *rb_right;
struct rb_node *rb_left;
} __attribute__((aligned(sizeof(long))));
可以看出 parent 与 color 共用一个 field，__attribute__((aligned(sizeof(long)))) 可以保证该结构在内存中的地址至少是 4 字节对齐的，所以这个字段的最后两个 bit 可以用来表述红黑颜色，设置 parent 时使用如下方法：
#define rb_parent(r) ((struct rb_node *)((r)->__rb_parent_color & ~3))
设置 color 时使用如下方法：
static inline void rb_set_black(struct rb_node *rb)
{
rb->__rb_parent_color |= RB_BLACK;
}
这种精巧的节约内存占用的设计在 Linux 内核中经常出现，实际工程实践中具有很大的借鉴意义。
② 数据结构嵌入方式
在设计一个数据结构时，通常比较本位的直观思想是将该数据结构的负载作为结构的一个 field 来嵌入到数据结构中，例如 paho 中的红黑树定义：
typedef struct NodeStruct {
struct NodeStruct *parent; /**< pointer to parent tree node, in case we need it */
struct NodeStruct *child[2]; /**< pointers to child tree nodes 0 = left, 1 = right */
void* content; /**< pointer to element content */
size_t size; /**< size of content */
unsigned int red : 1;
} Node;
通过一个 void* content 来负载红黑树的实际数据，这种方式在负载只存在于一种数据结构中时比较方便直观，但对于负载存在于多种数据结构中时，反过来将数据结构嵌入到负载中往往比较方便。当然这样需要一些实现技巧，Linux 内核里最经典的 container_of 便可以应用在这种场景中。
#define container_of(ptr, type, member) ({ \
const typeof( ((type *)0)->member ) *__mptr = (ptr); \
(type *)( (char *)__mptr - offsetof(type,member) );})
#define rb_entry(ptr, type, member) container_of(ptr, type, member)
rb_entry 可以很方便的返回当前红黑树节点的负载指针。
③ 缓存
缓存思想无处不在，对于红黑树中最高频的计算 – leftmost 节点的计算，Linux 内核红黑树中将其缓存在 root 结构中：
struct rb_root_cached {
struct rb_root rb_root;
struct rb_node *rb_leftmost;
};
这样会增加 insert 和 remove 的操作复杂度，但对于迭代器这种使用场景却是非常高效的。
④ 扩展
Linux 内核红黑树的实现中最优雅的部分是可以很方便的对其进行扩展。这样我们可以利用红黑树查找操作 O(lgn) 时间复杂度的特性，充分的发挥想象力，去扩展很多应用场景。
首先，没有将红黑树 insert 和 remove 接口直接封装，将 insert 操作中的 link_node 和 rebalance 分开封装，在进行扩展时由使用者自己实现 insert 操作，这样可以同时操作 augmentd 数据。
在红黑树的 rebalance 操作中，节点的 augmented 数据也会受到影响，Linux 内核的红黑树在 rbtree_augmented.h 中通过 callback 的方式，由使用者自行定义 rebalance 操作中 augmented 数据的重新计算方式：
struct rb_augment_callbacks {
void (*propagate)(struct rb_node *node, struct rb_node *stop);
void (*copy)(struct rb_node *old, struct rb_node *new);
void (*rotate)(struct rb_node *old, struct rb_node *new);
};
在下一节中，我们将对 Linux 内核的红黑树进行一次扩展实践，来验证这种数据结构设计方式的优点。
0x02 红黑树扩展实践
《算法导论》第 14 章 数据结构的扩张 中使用红黑树举例来介绍数据结构扩张的一般方法：
- 选择一种基础数据结构；
- 确定基础数据结构中要维护的附加信息；
- 检验基础数据结构上的基本修改操作能否维护附加信息；
- 设计一些新的操作；
接下来我们通过 Linux 内核中红黑树的接口来实现一次对红黑树的扩展实践。顺序统计树(order-statistic tree) 可以在 O(lgn) 时间内计算集合中元素的秩，相比于朴素数据结构中遍历元素对比来统计秩的方式在时间复杂度上有指数级的提升。
https://github.com/aggresss/playground-algorithm/tree/main/rbtree/augment/order-statistic-tree
struct ostree_node {
uint32_t key;
uint32_t augmented;
struct rb_node rb;
};
static void augment_compute(struct rb_node *rb) {
if (!rb) {
return;
}
uint32_t augmented = 1;
if (rb->rb_left) {
augmented += rb_entry(rb->rb_left, struct ostree_node, rb)->augmented;
}
if (rb->rb_right) {
augmented += rb_entry(rb->rb_right, struct ostree_node, rb)->augmented;
}
rb_entry(rb, struct ostree_node, rb)->augmented = augmented;
}
static void augment_propagate(struct rb_node *rb, struct rb_node *stop) {
while (rb != stop) {
struct ostree_node *node = rb_entry(rb, struct ostree_node, rb);
augment_compute(&node->rb);
rb = rb_parent(&node->rb);
}
}
static void augment_copy(struct rb_node *rb_old, struct rb_node *rb_new) {
struct ostree_node *old = rb_entry(rb_old, struct ostree_node, rb);
struct ostree_node *new = rb_entry(rb_new, struct ostree_node, rb);
new->augmented = old->augmented;
}
static void augment_rotate(struct rb_node *rb_old, struct rb_node *rb_new) {
augment_compute(rb_old);
augment_compute(rb_new);
}
static const struct rb_augment_callbacks augment_callbacks = {
augment_propagate,
augment_copy,
augment_rotate};
void ostree_insert(struct ostree_node *node, struct rb_root_cached *root) {
struct rb_node **new = &root->rb_root.rb_node, *rb_parent = NULL;
uint32_t key = node->key;
struct ostree_node *parent;
while (*new) {
rb_parent = *new;
parent = rb_entry(rb_parent, struct ostree_node, rb);
parent->augmented++;
if (key < parent->key)
new = &parent->rb.rb_left;
else
new = &parent->rb.rb_right;
}
node->augmented = 1;
rb_link_node(&node->rb, rb_parent, new);
rb_insert_augmented(&node->rb, &root->rb_root, &augment_callbacks);
}
void ostree_remove(struct ostree_node *node, struct rb_root_cached *root) {
rb_erase_augmented(&node->rb, &root->rb_root, &augment_callbacks);
}
struct ostree_node *ostree_select(struct rb_root_cached *root, uint32_t rank) {
uint32_t node_rank = 0;
struct ostree_node *osnode = NULL;
struct rb_node *rbnode = root->rb_root.rb_node;
while (rbnode) {
node_rank = 1;
if (rbnode->rb_left) {
node_rank += rb_entry(rbnode->rb_left, struct ostree_node, rb)->augmented;
}
if (rank == node_rank) {
osnode = rb_entry(rbnode, struct ostree_node, rb);
break;
} else if (rank < node_rank) {
rbnode = rbnode->rb_left;
} else {
rank -= node_rank;
rbnode = rbnode->rb_right;
}
}
return osnode;
}
uint32_t ostree_rank(struct rb_root_cached *root, struct ostree_node *node) {
uint32_t rank;
struct rb_node *rbnode = &node->rb;
rank = 1;
if (rbnode->rb_left) {
rank += rb_entry(rbnode->rb_left, struct ostree_node, rb)->augmented;
}
while (rbnode != root->rb_root.rb_node) {
if (rbnode == rb_parent(rbnode)->rb_right) {
rank += 1;
if (rb_parent(rbnode)->rb_left) {
rank += rb_entry(rb_parent(rbnode)->rb_left, struct ostree_node, rb)->augmented;
}
}
rbnode = rb_parent(rbnode);
}
return rank;
}