Linux内核源码之红黑树注释

红黑树的原理理解起来并不困难，无非是左右旋转，涂色。然而要编码出一个高效简洁的实现代码却挺复杂。近日学习Linux源码，对其红黑树实现作了注释，与大家分享一下。我贴的代码是我边理解边“抄写”出来的，和源始代码可能会有少许出入，但不影响理解。本文着重注释插入删除操作，其它部分不作重点。若要了解红黑树的基础知识或是Linux红黑树的使用，请另行参考其它资料。

先看一下节点定义：

typedef struct st_rb_node {
/**
* 本结构体四字节对齐，因而其地址中低位两个bit永远是0。
* Linux内核开发人员非常爱惜内存，他们用其中一个空闲的位来存储颜色信息。
* parent_color成员实际上包含了父节点指针和自己的颜色信息。
*/
unsigned long parent_color;
#define RB_RED   0
#define RB_BLACK 1
struct st_rb_node *left, *right;
}__attribute__((aligned(sizeof(long)))) rb_node;

Linux的数据结构有个特点就是结构不包括数据，而是由数据来包括结构信息。由结构体获取数据信息是通过 CONTAINER_OF 这个宏来实现的，它利用了一些编译器的特性，有兴趣的可以参考Linux的链表源码。

头文件中其它主要内容如下（比较好理解，未作注释）：

#define rb_parent(r)    ((rb_node *)((r)->parent_color & ~3))
#define rb_color(r)     ((r)->parent_color & 1)
#define rb_is_red(r)    (!rb_color(r))
#define rb_is_black(r)  rb_color(r)
#define rb_set_red(r)   do { (r)->parent_color &= ~1; } while (0)
#define rb_set_black(r) do { (r)->parent_color |= 1; } while (0)

static inline void rb_set_parent(rb_node *node, rb_node *p) {
node->parent_color = (node->parent_color & 3) | (unsigned long) p;
}

static inline void rb_set_color(rb_node *node, int color) {
node->parent_color = (node->parent_color & ~1) | color;
}

#define RB_ROOT (rb_root) { NULL, }
#define rb_entry(ptr, type, member) container_of(ptr, type, member)

#define RB_EMPTY_ROOT(root) ((root)->node == NULL)
#define RB_EMPTY_NODE(node) (rb_parent(node) == node)
#define RB_CLEAR_NODE(node) (rb_set_parent(node, node))

void rb_insert_color(rb_node *node, rb_root *root);
void rb_erase(rb_node *node, rb_root *root);

rb_node * rb_next(const rb_node *node);
rb_node * rb_prev(const rb_node *node);
rb_node * rb_first(const rb_root *root);
rb_node * rb_last(const rb_root *root);

void rb_replace_node(rb_node *victim, rb_node *new, rb_root *root);

static inline void rb_link_node(rb_node *node, rb_node *p, rb_node **link) {
node->parent_color = (unsigned long) p;
node->left = node->right = NULL;
*link = node;
}

两个基本操作：左旋与右旋（未注释，仅为方便参考）

/**
* 对node进行左旋。有关概念可参考WIKI等资源。
*/
static void _rb_rotate_left(rb_node *node, rb_root *root) {
rb_node *right = node->right, *parent = rb_parent(node);

if ((node->right = right->left))
rb_set_parent(node->right, node);

right->left = node;
rb_set_parent(node, right);

if (parent) {
if (node == parent->left)
parent->left = right;
else
parent->right = right;
} else {
root->node = right;
}
rb_set_parent(right, parent);
}

/**
* 对node进行右旋。有关概念可参考WIKI等资源。
*/
static void _rb_rotate_right(rb_node *node, rb_root *root) {
rb_node *left = node->left, *parent = rb_parent(node);

if ((node->left = left->right))
rb_set_parent(node->left, node);

left->right = node;
rb_set_parent(node, left);

if (parent) {
if (node == parent->left)
parent->left = left;
else
parent->right = left;
} else {
root->node = left;
}
rb_set_parent(left, parent);
}

插入操作的实现（前提是节点已经插入在目标位置，下面的函数只负责修正树的性质）：

/**
* node是新插入的结点，有着默认的红色。
* 本函数检查是否有违背红黑树性质的地方，并进行纠正。
*/
void rb_insert_color(rb_node *node, rb_root *root) {
rb_node *parent, *gp;

/**
* 如果有父节点且父节点是红色，进行树的调整以保证树的性质。
*/
while ((parent = rb_parent(node)) && rb_is_red(parent)) {
gp = rb_parent(parent);

if (parent == gp->left) {
/**
* 父节点是祖父节点左子的情况。
* "else"中的情况左右相反，这里只注释"if"里的代码。
*/
register rb_node *uncle = gp->right;
if (uncle && rb_is_red(uncle)) {
/**
* 如果有红叔，则将红叔与红父均涂黑，并将祖父节点涂红。
*/
rb_set_black(parent);
rb_set_black(uncle);
rb_set_red(gp);
/**
* 现在简单路径中黑节点个数仍然平衡，但祖父变成了红色，
* 我们不确定有没有造成父子均红的情况，所以需要对祖父节点进行下一轮修复。
*/
node = gp;
continue;
}

/**
* 现在是叔节点为空或为黑的情况。
*/
if (node == parent->right) {
/**
* 如果新节点是父节点的右子，对父节点进行左旋。
* 旋转后树仍然平衡，但新节点占了原父节点的位子。
* 这两个节点交换角色后，新的父节点是红的，其左子也是红的。
*/
_rb_rotate_left(parent, root);
register rb_node *tmp = node;
node = parent;
parent = tmp;
}
/**
* 此时父红左子红，树平衡但有连续红节点。
*
* 父涂黑，祖父涂红，再对祖父右旋，树即调整到合法状态。
*/
rb_set_black(parent);
rb_set_red(gp);
_rb_rotate_right(gp, root);
return;
} else {
register rb_node *uncle = gp->left;
if (uncle && rb_is_red(uncle)) {
rb_set_black(parent);
rb_set_black(uncle);
rb_set_red(gp);
node = gp;
continue;
}

if (node == parent->left) {
_rb_rotate_right(parent, root);
register rb_node *tmp = node;
node = parent;
parent = tmp;
}
rb_set_black(parent);
rb_set_red(gp);
_rb_rotate_left(gp, root);
return;
}
}
/**
* 若无父节点，只需将node涂黑；
* 若父节点为黑，插入红节点不影响树的性质。
* 循环体后直接将node涂黑，可以同时保证以上两点。
*/
rb_set_black(root->node);
}

最复杂的是删除操作，分为两部分。

删除第一步，移出节点：

/**
* 删除node节点（其实只是将其与树脱离），并修正树。
* 红黑树的节点删除相对比插入更复杂，加之Linux的源码风格也不易懂，
* 此段代码需认真体会，可以多画画图，对各种情况标记、对照一下。
*/
void rb_erase(rb_node *node, rb_root *root) {
rb_node *child, *parent;
int color;

/**
* 这里其实分成了两个大分枝。
* 如果node有任一空子，则记录child并跳至条件语句后的语句，具体位置参考注释；
* 如果两子均不为空，则进入"else"分枝。
*/
if (!node->left) {
child = node->right;
} else if (!node->right) {
child = node->left;
} else {
/**
* 乾坤大挪移，将要删除的节点转换成只有一个子节点或无子节点的节点。
* 具体方法是，找到它的下一个节点，也即先找其右节点，再循环找左节点，
* 直到没有左子节点。找到的节点一定是比node"大"且紧邻node的节点。
* 将此找到的节点放到node位置，并保持node的原有颜色，树的性质并不会受影响。
* 此时问题转化成了删除找到的没有左子的节点。
* （提示：也可以找左子的右节点，循环下去，得到的将是“紧邻”的小于node的节点）
*/
rb_node *old = node, *left;
node = node->right;
while ((left = node->left) != NULL)
node = left;

/**
* 找到节点后，将其“大挪移”过去。
* 首先用old的父亲或root指向新找到的节点node.
*/
if (rb_parent(old)) {
if (rb_parent(old)->left == old)
rb_parent(old)->left = node;
else
rb_parent(old)->right = node;
} else {
root->node = node;
}

/**
* 记录找到的node节点相关信息。对于这种删除操作，node节点移至old位置后会沿用old的颜色属性，
* 树的性质不受影响，但原node位置将失去一个节点，我们要记录节点的相关信息以便做节点的删除工作。
*/
child = node->right;
parent = rb_parent(node);
color = rb_color(node);

if (parent == old) {
/**
* 参考前面while循环，如果parent==old,则找到的单子节点只能是old节点的右子。
* 这里保存的parent可以理 解为将用作node原来子节点的亲父节点。对于old是node右子的情况，
* node取代old后，原node的父节点还是node.
*/
parent = node;
} else {
/**
* node节点的删除处理。主要是将相应的父、子联结起来。
* 注意我们的“乾坤大挪移”可能会比较绕，删的数据是old节点，但我们将node转移过去占了old的位置，
* 所以对于节点结构来说，我们删除的是node原来的位置。
*/
if (child)
rb_set_parent(child, parent);
parent->left = child;

node->right = old->right;
rb_set_parent(old->right, node);
}

/**
* node移过去后继续使用old节点的颜色和父节点属性。
* parent_color同时保存了父节点信息和自己的颜色信息。
* 关于parent_color请参考头文件中的宏。
*/
node->parent_color = old->parent_color;
/**
* 之前已将old的parent的指向old子节点的指针指向了node，
* 并且node也已经连接了old的右子（如果node与old是子与父的关系，则node保留原右子即可）
* 上面一句完成了node指向新父亲即设置颜色的操作，现在只剩下左子了，把它搞定！
*/
node->left = old->left;
rb_set_parent(old->left, node);

/**
* 该删的该改的都已搞定，现在要检查颜色了。
*/
goto color;
}

/**
* 这里的条件语句和goto语句实在是让人晕头转向。现在对应的是本函数最上面的if语句中前两个case。
*
* 如果node有任一空子，也即node无子或有单子。这里跳过了前面的“乾坤大挪移”，同时相比挪移的情况，
* 这里的区别是node可能是右子，也可能是左子。但这里的删除操作要简单些。
*
* 思考：挪移与不挪移的情况，最终都是把问题转到了一个无子或只有单子的节点上，
* 为什么不能让两种情况执行一段共同的删除代码呢？
* 我想应该是挪移本身已经使node从原位置消失了，两者删除操作有差别，加之Linux内核对速度的狂热，
* 就分情况写成了现在这种很绕的代码。
*/
parent = rb_parent(node);
color = rb_color(node);

if (child)
rb_set_parent(child, parent);
if (parent) {
if (parent->left = node)
parent->left = child;
else
parent->right = child;
} else {
root->node = child;
}

color:
/**
* color是删除节点的颜色，删除的节点的层次正好在child参数与parent参数中间。
* 如果删除的是红节点，对树没有影响，所以只有删除的是黑节点时才修正颜色。
*/
if (color == RB_BLACK)
_rb_erase_color(child, parent, root);
}

第二步，简单移出如果破坏了树的性质，则需要修复：

/**
* 前置环境：在node与parent之前，刚刚删除了一个黑色节点。现在树很可能不平衡，
* node与parent也可能红色冲突。
* 本函数进行树的性质的修正，以使树恢复平衡。在一些情况下问题会转移到上一层节点，
* 则须对上一层节点进行递归检查与修正。本函数中的while循环实际上实现了这种递归。
*
* 提示：这是红黑树里最绕的地方，看不明白可以多画画图。
*/
static void _rb_erase_color(rb_node *node, rb_node *parent, rb_root *root) {
/**
* other用来保存兄弟节点，这是树的修正过程中一个重要的参考节点。
*/
rb_node *other;
/**
* 循环条件：node不是红节点且node不是根节点。
* 解释：对于红节点或根节点，直接涂黑即可解决问题。
*/
while ((!node || rb_is_black(node)) && node != root->node) {
/**
* 为方便程序处理各种情况，这里和节点插入一样将问题分成了左右对称的两类。
* if-else两个分枝里代码逻辑完全相同，只是左右相反。所以我们只研究node是parent左子的情况。
*
* 在开始之前，我们先总结一下当前状态：
* 1:因为删除的是黑色节点，所以node与parent都有可能是红色节点。
* 2:node与parent之间少了一个黑色节点，则所有通过node的路径都少了一个黑色节点，不妨画图时用-1标出来；
* 但node的兄弟节点（node一定有兄弟，可以根据删除前树的平衡性质来反推）高度并未变化，可以记作0。
*
* 提示：在进行旋转、涂色等操作时，可以画图观查平衡状态的变化。
*/
if (parent->left == node) {
other = parent->right;
if (rb_is_red(other)) {
/**
* 如果兄弟节点是红色，则父节点是黑色，交换父、兄节点的颜色，并对父节点进行左旋。
* 旋转后，兄节点占了老的父节点位置，且和老的父节点颜色相同，所以不会向上造成颜色冲突。
* 我们仍然以老的父节点为父节点来看，现在的状态是：
* 父节点右子保持平衡，只有经过node的路径少了一个黑色节点。
* 现在问题和之前相似，但node有了一个黑色的兄弟（还有一个红色父亲）。
*/
rb_set_black(other);
rb_set_red(parent);
_rb_rotate_left(parent, root);
/**
* other指向新的兄弟节点。other现在必然是一个黑色节点，而不会是空。这一点可以根据旋转之前树的性质反证。
*/
other = parent->right;
}
/**
* 此时状态：
* node有黑色兄弟，父亲可能黑也可能红。
*/
if ((!other->left || rb_is_black(other->left)) &&
(!other->right || rb_is_black(other->right))) {
/**
* 如果other没有红色子节点，那我们就可以把other涂红，并向上转移问题。
* other涂红的后果是，other分枝少了一个黑节点，与node分枝保持了平衡，但parent整体少了一个黑色节点。
* 细心的人可能会发现，如果父亲是红色的，父亲与兄弟有颜色冲突，直接向上转移能纠正吗？当然是可以的。
* 向上一层之后，while里的黑节点判断失败，会直接执行while后面的语句，直接将parent涂黑，则树恢复平衡。
*/
rb_set_red(other);
node = parent;
parent = rb_parent(node);
} else {
/**
* 现在黑兄弟有红子节点，父亲颜色未知。
*/
if (!other->right || rb_is_black(other->right)) {
/**
* 如果黑兄弟右节点为空或为黑，则左节点一定是红的，我们想办法把它调整为右子为红。
* 至于为什么，看后面就知了。
* other->left与other交换颜色，对other进行右旋，other指向新的兄弟。根据右旋转的特点，
* 可知现在other仍然是黑色，且它有了一个红色右子。同时other分枝高度不变，颜色也没有冲突。
*/
rb_set_black(other->left);
rb_set_red(other);
_rb_rotate_right(other, root);
other = parent->right;
}
/**
* 此时状态：黑兄弟有红色右子节点。
* 不管parent是什么颜色，把other涂成父亲的颜色（之后旋转，other占据父亲的位置，向上没有颜色冲突），
* 把父亲涂黑，把黑兄的other涂黑，这时node分枝高度可能有变化也可能没变化，other分枝多了一个黑节点。
* 现在对父亲进行左旋转。旋转后的情况是右边分枝（原other右子）少了一个黑节点，重归平衡；
* 左边分枝则增加了一个黑节点，也恢复了平衡。此时也没有颜色冲突
*/
rb_set_color(other, rb_color(parent));
rb_set_black(parent);
rb_set_black(other->right);
_rb_rotate_left(parent, root);
/**
* 树已平衡，node置为根节点，并跳出循环。
* 疑问：这里为什么还要检查根节点呢？
*/
node = root->node;
break;
}
} else {
other = parent->left;
if (rb_is_red(other)) {
rb_set_black(other);
rb_set_red(parent);
_rb_rotate_right(parent, root);
other = parent->left;
}
if ((!other->left || rb_is_black(other->left)) &&
(!other->right || rb_is_black(other->right))) {
rb_set_red(other);
node = parent;
parent = rb_parent(node);
} else {
if (!other->left || rb_is_black(other->left)) {
rb_set_black(other->right);
rb_set_red(other);
_rb_rotate_left(other, root);
other = parent->left;
}
rb_set_color(other, rb_color(parent));
rb_set_black(parent);
rb_set_black(other->left);
_rb_rotate_right(parent, root);
node = root->node;
break;
}
}
}
/**
* 对于红节点或根节点，直接涂黑即可解决问题。
*/
if (node)
rb_set_black(node);
}

匆忙之中难免有疏漏，欢迎指正。