linux内核的红黑树

linux内核的红黑树

很多博文都介绍过红黑树，linux内核对红黑树的实现无疑是经典中的经典，那么就来看看这经典中的经典是怎么样的。

1. 回忆树的基本知识

什么是树，有很多博文中会罗列几条规则，满足那几条规则的数据结构就称为树，这种定义方式是很数学化的表达方式。这里抓出来两条最重要的规则：

1)树的节点有且仅有一个父节点；
2)任何非页节点都可以分为若干不相交的子树。
其实这两条规则都是表述同一个含义，就是树是不相交的，节点之间不能构成回环。
常用到的是二叉树，二叉树对节点孩子的个数进行了约束，即节点的孩子个数不超过两个。二叉树的一个重要应用是排序与查找，所以就有了排序二叉树或者说查找二叉树。排序二叉树的左子树的节点的值都小于根节点，右子树的的节点的值都大于根节点。在红黑树的左旋和右旋操作中要用到该性质，下图所示，即为一个左旋转。问题是为什么D可以作为A的右孩子，因为D总是比A大的，这就是二叉树性质的运用。

图1 二叉树性质的运用
二叉树在不幸的时候会退化为链表结构，为了这个有了平衡二叉树，红黑树就是我们常用的平衡二叉树之一。是的，今天的主角就是红黑树。

2. 什么是红黑树

linux源码(linux 5.13)中对红黑树的定义是如下的描述，那么黑节点的孩子可以为黑吗？这是完全有可能的。

1) A node is either red or black
2) The root is black //根节点为黑
3) All leaves (NULL) are black //所有的叶节点为黑
4) Both children of every red node are black //所有红节点的孩子为黑
5) Every simple path from root to leaves contains the same number
   of black nodes.

3. linux源码中红黑树的实现

我们从使用者的角度来看红黑树是如何实现的。在linux中，双向链表、哈希链表、红黑树的实现都有一个共同点，那就是数据结构只用于连接，而不存放数据。需要使用这些连接的数据结构，其中必须包含相应连接的成员。不论是链表还是树，都有根节点。红黑树的根节点，就是一个指向节点的指针。

struct rb_root {
	struct rb_node *rb_node;
};
struct rb_node {
	unsigned long  __rb_parent_color;
	struct rb_node *rb_right;
	struct rb_node *rb_left;
} __attribute__((aligned(sizeof(long))));//最低的4位始终为0，可作其它用途

定义一个树的根节点使用宏RB_ROOT，rb_node指针被初始化为NULL。RB_ROOT返回的是一个rb_root实例，由于one_tree是数组，所以右边需要加花括号。one_tree是数组名，是用于赋值给指针。root_tree就是指向红黑树根节点的指针，当有很多红黑树时，只需把root_tree当做一个数组看待。

struct rb_root one_tree[1] = { RB_ROOT };
struct rb_root *root_tree = one_tree;

在红黑树中查找的样例如下所示，linux内核并没有采用提供回调函数的方式让使用者直接使用，而是需要开发者自己定义。在如下的流程中，开发者可以改变比较大小的方式，还可以根据需要增加一些其它的处理流程。红黑树的查找与普通的二叉排序树并没有区别。

struct mytype *my_search(struct rb_root *root, char *string)
  {
  	struct rb_node *node = root->rb_node; //rb_node是一个指针

  	while (node) {
  	//参数的第三个参数node是结构体mytype中的成员名称
  		struct mytype *data = container_of(node, struct mytype, node);
		int result;

		result = strcmp(string, data->keystring);

		if (result < 0)
  			node = node->rb_left;
		else if (result > 0)
  			node = node->rb_right;
		else
  			return data;
	}
	return NULL;
  }

有差别的是在红黑树中插入节点。与查找不同，一开始定义的指针new是一个指向根节点的指针的指针。*new的含义即是指向节点的指针。为什么要用指向指针的指针呢？目的是为了在通过函数传参时能修改指针的值。rb_link_node将data->node中的数据成员进行了修改，__rb_parent_color的颜色修改为其父节点的颜色，左孩子和右孩子都设置为NULL，并将new指针指向node指针。那么new指针的含义是什么，它可能是左孩子，也可能是右孩子，更确切的含义是当前节点要插入的位置。

int my_insert(struct rb_root *root, struct mytype *data)
  {
  	struct rb_node **new = &(root->rb_node), *parent = NULL;

  	/* Figure out where to put new node */
  	while (*new) {
  		struct mytype *this = container_of(*new, struct mytype, node);
  		int result = strcmp(data->keystring, this->keystring);

		parent = *new;
  		if (result < 0)
  			new = &((*new)->rb_left);
  		else if (result > 0)
  			new = &((*new)->rb_right);
  		else
  			return FALSE;
  	}

  	/* Add new node and rebalance tree. */
  	rb_link_node(&data->node, parent, new);
  	rb_insert_color(&data->node, root);

	return TRUE;
  }

rb_insert_color函数是红黑树相对于普通二叉排序树最大的差别。这个函数中平衡二叉树的方式主要有3种：a)改变颜色，0表示red，1表示black。对于对齐的节点，低4位总是0，所以默认的情况下是红色的节点。那么对于父节点是black的情况，直接插入即可，无需改变颜色。如果父节点是红色则需要改变颜色。b)左旋转，将一个节点的右孩子变为父节点，节点本身变为左孩子。c)右旋转，将一个节点的做孩子变为父节点，节点本身变为右孩子。

static __always_inline void
__rb_insert(struct rb_node *node, struct rb_root *root,
	    void (*augment_rotate)(struct rb_node *old, struct rb_node *new))
{
	// 节点在没有上色之前__rb_parent_color表示指向父节点的指针，且低位未被标记
	struct rb_node *parent = rb_red_parent(node), *gparent, *tmp;

	while (true) {
		/*
		 * Loop invariant: node is red.
		 */
		if (unlikely(!parent)) {
			/*
			 * The inserted node is root. Either this is the
			 * first node, or we recursed at Case 1 below and
			 * are no longer violating 4).
			 */
			// 根节点直接设置为黑色，1表示黑色，__rb_parent_color被设置为1.
			// 因为根节点没有父节点，所以第二个参数为NULL
			rb_set_parent_color(node, NULL, RB_BLACK);
			break;
		}

		/*
		 * If there is a black parent, we are done.
		 * Otherwise, take some corrective action as,
		 * per 4), we don't want a red root or two
		 * consecutive red nodes.
		 */
		// 父节点为黑色，无需做任何处理，颜色呢？
		if(rb_is_black(parent))
			break;

		// 父节点为红色才能走到此处为红色意味着低位为0.
		gparent = rb_red_parent(parent);

		tmp = gparent->rb_right;
		if (parent != tmp) {	/* parent == gparent->rb_left */
			if (tmp && rb_is_red(tmp)) {// 父节点存在兄弟
				/*
				 * Case 1 - node's uncle is red (color flips).
				 *
				 *       G            g
				 *      / \          / \
				 *     p   u  -->   P   U
				 *    /            /
				 *   n            n
				 *
				 * However, since g's parent might be red, and
				 * 4) does not allow this, we need to recurse
				 * at g.
				 */
				// 父节点及其兄弟的parent都设置为黑，parent和兄弟不能直接访问其父节点了
				rb_set_parent_color(tmp, gparent, RB_BLACK);
				rb_set_parent_color(parent, gparent, RB_BLACK);
				node = gparent;
				parent = rb_parent(node);//最低两位被设为0，parent已经被设置为祖父的parent

				rb_set_parent_color(node, parent, RB_RED);
				continue;
			}

			tmp = parent->rb_right;
			if (node == tmp) {
				/*
				 * Case 2 - node's uncle is black and node is
				 * the parent's right child (left rotate at parent).
				 *
				 *      G             G
				 *     / \           / \
				 *    p   U  -->    n   U
				 *     \           /
				 *      n         p
				 *
				 * This still leaves us in violation of 4), the
				 * continuation into Case 3 will fix that.
				 */
				tmp = node->rb_left;
				WRITE_ONCE(parent->rb_right, tmp);
				WRITE_ONCE(node->rb_left, parent);
				if (tmp)
					rb_set_parent_color(tmp, parent,
							    RB_BLACK);
				rb_set_parent_color(parent, node, RB_RED);
				augment_rotate(parent, node);
				parent = node;
				tmp = node->rb_right;
			}

			/*
			 * Case 3 - node's uncle is black and node is
			 * the parent's left child (right rotate at gparent).
			 *
			 *        G           P
			 *       / \         / \
			 *      p   U  -->  n   g
			 *     /                 \
			 *    n                   U
			 */
			WRITE_ONCE(gparent->rb_left, tmp); /* == parent->rb_right */
			WRITE_ONCE(parent->rb_right, gparent);
			if (tmp)
				rb_set_parent_color(tmp, gparent, RB_BLACK);
			__rb_rotate_set_parents(gparent, parent, root, RB_RED);
			augment_rotate(gparent, parent);
			break;
		} else {
			tmp = gparent->rb_left;
			if (tmp && rb_is_red(tmp)) {
				/* Case 1 - color flips */
				rb_set_parent_color(tmp, gparent, RB_BLACK);
				rb_set_parent_color(parent, gparent, RB_BLACK);
				node = gparent;
				parent = rb_parent(node);
				rb_set_parent_color(node, parent, RB_RED);
				continue;
			}

			tmp = parent->rb_left;
			if (node == tmp) {
				/* Case 2 - right rotate at parent */
				tmp = node->rb_right;
				// 非锁的方式读时确保不是致命的
				WRITE_ONCE(parent->rb_left, tmp);
				WRITE_ONCE(node->rb_right, parent);
				if (tmp)
					rb_set_parent_color(tmp, parent,
							    RB_BLACK);
				rb_set_parent_color(parent, node, RB_RED);
				augment_rotate(parent, node);
				parent = node;
				tmp = node->rb_left;
			}

			/* Case 3 - left rotate at gparent */
			WRITE_ONCE(gparent->rb_right, tmp); /* == parent->rb_left */
			WRITE_ONCE(parent->rb_left, gparent);
			if (tmp)
				rb_set_parent_color(tmp, gparent, RB_BLACK);
			__rb_rotate_set_parents(gparent, parent, root, RB_RED);
			augment_rotate(gparent, parent);
			break;
		}
	}
}