Linux核心紅黑樹原理與使用

紅黑樹（Red-Black Tree，RBT）是一種平衡的二叉查詢樹，前面的紅黑樹原理與實現這篇文章中詳細介紹了紅黑樹的細節。在Linux的核心原始碼中已經給我們實現了一棵紅黑樹，我們可以方便地拿過來進行使用。

本文將參考Linux核心的原始碼和文件資料，介紹Linux核心中紅黑樹的實現細節及使用方法。

簡介

Linux有很多地方用到了紅黑樹，比如高精度計時器使用紅黑樹樹組織定時請求，EXT3檔案系統也使用紅黑樹樹來管理目錄，虛擬儲存管理系統也有用紅黑樹樹進行VMAs（Virtual Memory Areas）的管理。

本文參考的Linux核心版本為linux-2.6.39.4，可以從官網https://www.kernel.org/pub/linux/kernel/v2.6/上進行下載。其中關於紅黑樹的檔案位置為：

• 標頭檔案： linux-2.6.39.4\include\linux\rbtree.h
• 實現程式碼：linux-2.6.39.4\lib\rbtree.c
• 文件說明：linux-2.6.39.4\Documentation\rbtree.txt

結構定義

Linux核心紅黑樹的實現與傳統的實現方式有些不同，它對針對核心對速度的需要做了優化。每一個rb_node節點是嵌入在用RB樹進行組織的資料結構中，而不是用rb_node指標進行資料結構的組織。

Linux核心中紅黑樹節點的定義如下，其中rb_node是節點型別，而rb_root是僅包含一個節點指標的類，用來表示根節點。

struct rb_node
{
	unsigned long  rb_parent_color;
#define	RB_RED		0
#define	RB_BLACK	1
	struct rb_node *rb_right;
	struct rb_node *rb_left;
} __attribute__((aligned(sizeof(long))));

struct rb_root
{
	struct rb_node *rb_node;
};

粗略一看，這裡似乎沒有定義顏色的域，但這就是這裡紅黑樹實現的一個巧妙的地方。rb_parent_color這個域其實同時包含了顏色資訊以及父親節點的指標，因為該域是一個long的型別，需要大小為sizeof(long)的對齊，那麼在一般的32位機器上，其後兩位的數值永遠是0，於是可以拿其中的一位來表示顏色。事實上，這裡就是使用了最低位來表示顏色資訊。
明白了這點，那麼以下關於父親指標和顏色資訊的操作都比較容易理解了，其本質上都是對rb_parent_color的位進行操作。

#define rb_parent(r)   ((struct rb_node *)((r)->rb_parent_color & ~3)) //低兩位清0
#define rb_color(r)   ((r)->rb_parent_color & 1)                       //取最後一位
#define rb_is_red(r)   (!rb_color(r))                                  //最後一位為0？
#define rb_is_black(r) rb_color(r)                                     //最後一位為1？
#define rb_set_red(r)  do { (r)->rb_parent_color &= ~1; } while (0)    //最後一位置0
#define rb_set_black(r)  do { (r)->rb_parent_color |= 1; } while (0)   //最後一位置1

static inline void rb_set_parent(struct rb_node *rb, struct rb_node *p) //設定父親
{
	rb->rb_parent_color = (rb->rb_parent_color & 3) | (unsigned long)p;
}
static inline void rb_set_color(struct rb_node *rb, int color)          //設定顏色
{
	rb->rb_parent_color = (rb->rb_parent_color & ~1) | color;
}

然後是幾個巨集定義：

#define RB_ROOT	(struct rb_root) { NULL, }                         //初始根節點指標
#define rb_entry(ptr, type, member) container_of(ptr, type, member)//包含ptr的結構體指標
#define RB_EMPTY_ROOT(root) ((root)->rb_node == NULL)              //判斷樹是否空
#define RB_EMPTY_NODE(node) (rb_parent(node) == node)              //判斷節點是否空，父親是否等於自身
#define RB_CLEAR_NODE(node) (rb_set_parent(node, node))            //設定節點為空，父親等於自身

這裡需要注意的是container_of本身也是個巨集，其定義在kernel.h中：

#define container_of(ptr, type, member) ({                \
    const typeof( ((type *)0)->member ) *__mptr = (ptr);  \
    (type *)( (char *)__mptr - offsetof(type,member) );})

而其中的offsetof則定義在stddef.h中：

#define offsetof(TYPE, MEMBER) ((size_t) &((TYPE *)0)->MEMBER)

offsetof巨集取得member成員在type物件中相對於物件首地址的偏移量，具體是通過把0強制轉化成為type型別指標，然後引用成員member，此時得到的指標大小即為偏移量（因為物件首地址為0）。
container_of巨集取得包含ptr的資料結構的指標，具體是把ptr轉化為type物件中member型別的指標，然後減去member型別在type物件的偏移量得到type物件的首地址。

紅黑樹操作

接下來的__rb_rotate_left和__rb_rotate_right就是對紅黑樹進行的左旋和右旋操作。注意，程式碼中的第一個if語句中是=而不是==，意思是先賦值，然後再對該值判斷是否為空，如果不為空的情況下才設定該節點的父親。這樣程式碼顯得非常簡潔，但如果以為是==的比較，則可能會感到困惑，不夠他這裡也使用了兩個小括號進行提示，因為一般情況只需一個括號即可。

void __rb_rotate_left(struct rb_node *node, struct rb_root *root);
void __rb_rotate_right(struct rb_node *node, struct rb_root *root);

而rb_insert_color則是把新插入的節點進行著色，並且修正紅黑樹使其達到平衡，其效果就是前文的insertFixup的效果。

void rb_insert_color(struct rb_node *, struct rb_root *);

插入節點時需要把新節點指向其父親節點，這可以通過rb_link_node函式完成：

void rb_link_node(struct rb_node * node, struct rb_node * parent, struct rb_node ** rb_link);

刪除節點則通過rb_erase進行，然後通過__rb_erase_color進行紅黑樹的修正。

void rb_erase(struct rb_node *, struct rb_root *);
void __rb_erase_color(struct rb_node *node, struct rb_node *parent, struct rb_root *root);

可以通過呼叫rb_replace_node來替換一個節點，但是替換完成後並不會對紅黑樹做任何調整，所以如果新節點的值與被替換的值有所不同時，可能會出現問題。

void rb_replace_node(struct rb_node *old, struct rb_node *new, struct rb_root *tree);

另外有幾個進行紅黑樹遍歷的函式，其原理均非常簡單，本質上就是這裡的求後繼、前驅、最小值、最大值的函式實現，不過這裡的程式碼實現非常簡潔和巧妙。

extern struct rb_node *rb_next(const struct rb_node *); //後繼
extern struct rb_node *rb_prev(const struct rb_node *); //前驅
extern struct rb_node *rb_first(const struct rb_root *);//最小值
extern struct rb_node *rb_last(const struct rb_root *); //最大值

實際使用

Linux核心中的紅黑樹實現非常巧妙，我們可以在自己的程式中進行使用，不過要稍微進行修改具體的方法如下：

1. 拷貝rbtree.h和rbtree.c到工程目錄下。
2. 修改rbtree.h：刪除兩個#include語句，新增stddef.h中的NULL和offsetof巨集定義，新增kernel.h中的container_of巨集定義。
3. 修改rbtree.c：把兩個#include語句替換成#include “rbtree.h”，刪除所有刪除所有的EXPORT_SYMBOL巨集。
4. 可以開始使用，參考linux-2.6.39.4\Documentation\rbtree.txt文件。

使用核心中的rbtree原始碼，需要自己實現插入和搜尋的關鍵程式碼，下面提供一些簡單的例子，雖然內容差異很大，但是其基本思想是不變的，可以很容易改成需要的程式碼。

首先是搜尋節點，基本思想就是根據二叉查詢樹的查詢過程進行：

struct mytype *my_search(struct rb_root *root, char *string)
{
	struct rb_node *node = root->rb_node;
	while (node)
	{
		struct mytype *data = container_of(node, struct mytype, node);
		int 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;
}

然後是插入節點，需要在插入一個資料之前先要查詢到適合插入的位置，然後將節點加入到樹中並將樹調整到平衡狀態：

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;
}

最後是刪除節點，可以直接使用核心介面直接進行：

struct mytype *data = mysearch(&mytree, "walrus");
if (data)
{
	rb_erase(&data->node, &mytree);
	myfree(data);
}

另外如果要遍歷一棵紅黑樹，可以使用核心提供的介面進行，而不需要自己實現：

struct rb_node *node;
for (node = rb_first(&mytree); node; node = rb_next(node))
	printk("key=%s\n", rb_entry(node, struct mytype, node)->keystring);