Linux核心中的紅/黑樹狀結構

紅/黑樹狀結構是平衡二叉樹的一種，它有很好的性質，樹中的結點都是有序的，而且因為它本身就是平衡的，所以尋找也不會出現非常惡劣的情況，基於二叉樹的操作的時間複雜度是O(log(N))。Linux核心在管理vm_area_struct時就是採用了紅/黑樹狀結構來維護記憶體塊的。

先到include/linux/rbtree.h中看一下紅/黑樹狀結構的一些定義，如下：

[Copy to clipboard] [ - ]

CODE:

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_parent_color其實沒那麼簡單，Andrea Arcangeli在這裡使用了一個小的技巧，不過非常棒。正如名字所暗示，這個成員其實包含指向parent的指標和此結點的顏色！它是怎麼做到的呢？很簡單，對齊起了作用。既然是sizeof(long)大小的對齊，那麼在IA-32上，任何rb_node結構體的地址的低兩位肯定都是零，與其空著不用，還不如用它們表示顏色，反正顏色就兩種，其實一位就已經夠了。

這樣，提取parent指標只要把rb_parent_color成員的低兩位清零即可：

[Copy to clipboard] [ - ]

CODE:

#define rb_parent(r) ((struct rb_node *)((r)->rb_parent_color & ~3))

取顏色只要看最後一位即可：

[Copy to clipboard] [ - ]

CODE:

#define rb_color(r) ((r)->rb_parent_color & 1)

測試顏色和設定顏色也是水到渠成的事了。需要特別指出的是下面的一個內嵌函式：

[Copy to clipboard] [ - ]

CODE:

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

它把parent設為node的父結點，並且讓rb_link指向node。

我們把重點集中在lib/rbtree.c上，看看一些和紅/黑樹狀結構相關的重要演算法。開始之前我們一起回憶一下紅/黑樹狀結構的規則：

1. 每個結點要麼是紅色要麼是黑色；

2. 根結點必須是黑色；

3. 紅結點如果有孩子，其孩子必須都是黑色；

4. 從根結點到葉子的每條路徑必須包含相同數目的黑結點。

這四條規則可以限制一棵排序樹是平衡的。

__rb_rotate_left是把以root為根的樹中的node結點進行左旋，__rb_rotate_right是進行右旋。這兩個函數是為後面的插入和刪除服務，而不是為外部提供介面。

新插入的結點都設為葉子，染成紅色，插入後如果破壞了上述規則，通過調整顏色和旋轉可以恢複，二叉樹又重新平衡。插入操作的介面函數是

[Copy to clipboard] [ - ]

CODE:

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

它把已確定父結點的node結點融入到以root為根的紅/黑樹狀結構中，具體演算法的分析可以參考[1]中第14.3節，這裡的實現和書中的講解幾乎完全一樣。怎麼確定node的父結點應該在調用rb_insert_color之前通過手工迭帶完成。值得指出的一點是，雖然插入操作需要一個迴圈迭代，但是總的旋轉次數不會超過兩次！所以效率還是很樂觀的。

刪除操作多多少少都有點麻煩，它要先執行像普通二叉尋找樹的“刪除”，然後根據刪除結點的顏色來判斷是否執行進一步的操作。刪除的介面是：

[Copy to clipboard] [ - ]

CODE:

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

其實它並沒有真正刪除node，而只是讓它和以root為根的樹脫離關係，最後它還要判斷是否調用__rb_erase_color來調整。具體演算法的講解看參考[1]中第13.3和14.4節，__rb_erase_color對應書中的RB-DELETE-FIXUP，此處的實現和書上也基本上一致。

其餘的幾個介面就比較簡單了。

[Copy to clipboard] [ - ]

CODE:

struct rb_node *rb_first(struct rb_root *root);

在以root為根的樹中找出並返回最小的那個結點，只要從根結點一直向左走就是了。

[Copy to clipboard] [ - ]

CODE:

struct rb_node *rb_last(struct rb_root *root);

是找出並返回最大的那個，一直向右走。

[Copy to clipboard] [ - ]

CODE:

struct rb_node *rb_next(struct rb_node *node);

返回node在樹中的後繼，這個稍微複雜一點。如果node的右孩子不為空白，它只要返回node的右子樹中最小的結點即可；如果為空白，它要向上尋找，找到迭帶結點是其父親的左孩子的結點，返回父結點。如果一直上述到了根結點，返回NULL。

[Copy to clipboard] [ - ]

CODE:

struct rb_node *rb_prev(struct rb_node *node);

返回node的前驅，和rb_next中的操作對稱。

[Copy to clipboard] [ - ]

CODE:

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

用new替換以root為根的樹中的victim結點。

紅/黑樹狀結構介面使用的一個典型例子如下：

[Copy to clipboard] [ - ]

CODE:

static inline struct page * rb_search_page_cache(struct inode * inode,
                                                 unsigned long offset)
{
        struct rb_node * n = inode->i_rb_page_cache.rb_node;
        struct page * page;

        while (n)
        {
                page = rb_entry(n, struct page, rb_page_cache);

                if (offset < page->offset)
                        n = n->rb_left;
                else if (offset > page->offset)
                        n = n->rb_right;
                else
                        return page;
        }
        return NULL;
}

static inline struct page * __rb_insert_page_cache(struct inode * inode,
                                                   unsigned long offset,
                                                   struct rb_node * node)
{
        struct rb_node ** p = &inode->i_rb_page_cache.rb_node;
        struct rb_node * parent = NULL;
        struct page * page;

        while (*p)
        {
                parent = *p;
                page = rb_entry(parent, struct page, rb_page_cache);

                if (offset < page->offset)
                        p = &(*p)->rb_left;
                else if (offset > page->offset)
                        p = &(*p)->rb_right;
                else
                        return page;
        }

        rb_link_node(node, parent, p);

        return NULL;
}

static inline struct page * rb_insert_page_cache(struct inode * inode,
                                                 unsigned long offset,
                                                 struct rb_node * node)
{
        struct page * ret;
        if ((ret = __rb_insert_page_cache(inode, offset, node)))
                goto out;
        rb_insert_color(node, &inode->i_rb_page_cache);
out:
        return ret;
}

因為紅/黑樹狀結構的這些良好性質和實現中介面的簡易性，它被廣泛應用到核心編程中，大大提高了核心的效率。

參考資料：

1. Introduction to Algorithms, Thomas H. Cormen, Charles E. Leiserson, and Ronald L. Rivest, MIT Press.

2. Understanding the Linux Kernel, 3rd Edition, Daniel P. Bovet, Marco Cesati, O'Reilly.

3. Linux Kernel 2.6.19 source code.