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;
}
參考檔案
- https://github.com/torvalds/linux/blob/v5.11/include/linux/rbtree_augmented.h
- https://www.kernel.org/doc/html/latest/core-api/rbtree.html
- https://stackoverflow.com/questions/17288746/red-black-nodes-struct-alignment-in-linux-kernel
- 那些演算法在哪里？
轉載請註明出處，本文鏈接：https://www.uj5u.com/qukuanlian/272548.html
標籤：區塊鏈