紅黑樹(shù)源碼這個(gè)東西我這些年前前后后讀了好幾遍每次都有新收獲。第一次翻開(kāi) Linux 內(nèi)核的 lib/rbtree.c 時(shí)說(shuō)實(shí)話看得很吃力滿屏的 __rb_parent_color、____rb_erase_color光看名字就覺(jué)得頭大。但當(dāng)我真正把插入、刪除、旋轉(zhuǎn)、變色這條線捋順之后再去翻 GCC 標(biāo)準(zhǔn)庫(kù)里那份幾百行的 stl_tree.h發(fā)現(xiàn)原來(lái) STL 里 map、set 的底層 rb_tree 源碼就是同一套思想換了件衣裳。這篇文章我不打算講什么高深理論就把 rb_tree 源碼怎么讀、內(nèi)核版本和 STL 版本各自的門(mén)道、以及怎么把內(nèi)核那份 rbtree 搬到自己的工程里使用全部攤開(kāi)講清楚。適合三類人看一是準(zhǔn)備面試、需要手撕紅黑樹(shù)的同學(xué)二是做內(nèi)核驅(qū)動(dòng)或者嵌入式開(kāi)發(fā)想在工程里管理大量有序節(jié)點(diǎn)的朋友三是純粹想提升源碼閱讀功力的開(kāi)發(fā)者。讀完你會(huì)發(fā)現(xiàn)紅黑樹(shù)沒(méi)有傳說(shuō)中那么可怕真正牛的其實(shí)是工程實(shí)現(xiàn)里那些看似不起眼的細(xì)節(jié)。1. 為什么 rb_tree 源碼值得一讀再讀1.1 紅黑樹(shù)在真實(shí)世界的藏身之處很多人對(duì)紅黑樹(shù)的印象停留在《算法導(dǎo)論》第13章覺(jué)得它只是一個(gè)考試重點(diǎn)、面試考點(diǎn)。但實(shí)際上紅黑樹(shù)是現(xiàn)代軟件系統(tǒng)里最常用的平衡樹(shù)結(jié)構(gòu)藏得比你想象中深得多。先說(shuō)最常見(jiàn)的C 的 std::map、std::set、std::multimap、std::multiset在 GCC 的 libstdc 里底層就是 _Rb_tree。你用 map 存鍵值對(duì)的時(shí)候每一次插入、刪除、查找背后都在和紅黑樹(shù)的旋轉(zhuǎn)函數(shù)打交道。再看 Linux 內(nèi)核CFS 調(diào)度器用來(lái)管理可運(yùn)行進(jìn)程的就緒隊(duì)列用的是紅黑樹(shù)虛擬內(nèi)存管理里進(jìn)程的 VMA虛擬內(nèi)存區(qū)域是以紅黑樹(shù)組織的高精度定時(shí)器、epoll 的事件管理也都有紅黑樹(shù)的身影。CFS 調(diào)度器每次要選出 vruntime 最小的進(jìn)程實(shí)際上就是紅黑樹(shù)的“最左節(jié)點(diǎn)”查找時(shí)間復(fù)雜度 O(log n)。再看中間件領(lǐng)域Redis 的有序集合 zset當(dāng)成員數(shù)量超過(guò)閾值并且元素是 skiplist 編碼時(shí)底層有一層 dict skiplist但在某些實(shí)現(xiàn)和相關(guān)的有序結(jié)構(gòu)里紅黑樹(shù)同樣大量出現(xiàn)。Nginx 的定時(shí)器管理早期版本用的就是紅黑樹(shù)通過(guò) key 值直接定位到最近的超時(shí)事件?？梢哉f(shuō)從數(shù)據(jù)庫(kù)的索引思想到網(wǎng)絡(luò)框架的事件管理紅黑樹(shù)的應(yīng)用幾乎無(wú)處不在。所以讀 rb_tree 源碼不只是為了應(yīng)付面試而是你在真實(shí)工程里遲早要面對(duì)的東西。當(dāng)你需要在幾十萬(wàn)個(gè)節(jié)點(diǎn)里快速插入、刪除、查找有序數(shù)據(jù)時(shí)手寫(xiě)鏈表性能不夠用 AVL 樹(shù)旋轉(zhuǎn)太頻繁紅黑樹(shù)就是那個(gè)性能和實(shí)現(xiàn)復(fù)雜度平衡得最好的選擇。1.2 兩個(gè)經(jīng)典源碼流派內(nèi)核版與 STL 版市面上能讀到的 rb_tree 源碼大致分兩個(gè)流派。第一個(gè)流派是 Linux 內(nèi)核的實(shí)現(xiàn)文件位置在 lib/rbtree.c 和 include/linux/rbtree.h。這套實(shí)現(xiàn)的風(fēng)格極其克制為了節(jié)省內(nèi)存把節(jié)點(diǎn)的父指針和顏色塞進(jìn)了同一個(gè)無(wú)符號(hào)長(zhǎng)整型字段里旋轉(zhuǎn)和刪除修復(fù)函數(shù)用循環(huán)而不是遞歸大量使用宏和內(nèi)聯(lián)函數(shù)。內(nèi)核里所有紅黑樹(shù)節(jié)點(diǎn)都是嵌入到自定義結(jié)構(gòu)體中的通過(guò) container_of 宏找到宿主結(jié)構(gòu)這種“侵入式”設(shè)計(jì)使得節(jié)點(diǎn)管理非常高效不需要額外的內(nèi)存池。第二個(gè)流派是 GCC libstdc 里的 _Rb_tree也就是 stl_tree.h。這套實(shí)現(xiàn)更貼近教材上的經(jīng)典偽代碼顏色用 _Rb_tree_color 枚舉節(jié)點(diǎn)有獨(dú)立的內(nèi)存分配策略模板化程度很高。如果你讀的是《STL源碼剖析》這本書(shū)里的 rb_tree那對(duì)應(yīng)的大概是 SGI STL 的早期實(shí)現(xiàn)結(jié)構(gòu)上更清晰一些適合入門(mén)讀。我的建議是入門(mén)先讀教材理解算法進(jìn)階讀內(nèi)核源碼學(xué)工程技巧最后再回頭啃 STL 的模板實(shí)現(xiàn)去理解 C 泛型設(shè)計(jì)。兩條路線都能讓你對(duì) rb_tree 源碼的理解上一個(gè)臺(tái)階。2. 動(dòng)手讀源碼前必須吃透的底層原理2.1 五條不變量與“近似平衡”的精髓紅黑樹(shù)之所以叫紅黑樹(shù)是因?yàn)槊總€(gè)節(jié)點(diǎn)多了一個(gè)顏色屬性非紅即黑。它通過(guò)五條不變量來(lái)維持樹(shù)的平衡節(jié)點(diǎn)只有紅、黑兩種顏色。根節(jié)點(diǎn)是黑色的。葉子節(jié)點(diǎn)NIL 空節(jié)點(diǎn)是黑色的。紅色節(jié)點(diǎn)的兩個(gè)子節(jié)點(diǎn)必須是黑色的不能出現(xiàn)連續(xù)的紅色節(jié)點(diǎn)。從任意一個(gè)節(jié)點(diǎn)出發(fā)到它所有葉子節(jié)點(diǎn)的路徑上黑色節(jié)點(diǎn)的數(shù)量必須相同。第五條也就是常說(shuō)的“黑高相等”。為什么這五條規(guī)則就能讓樹(shù)保持相對(duì)平衡這里有個(gè)經(jīng)典的推導(dǎo)結(jié)論在一棵紅黑樹(shù)中最長(zhǎng)路徑的長(zhǎng)度不會(huì)超過(guò)最短路徑的兩倍。最短路徑自然就是全黑路徑而最長(zhǎng)路徑由于不能有連續(xù)紅色節(jié)點(diǎn)只能是“黑-紅-黑-紅”交替所以紅色節(jié)點(diǎn)的數(shù)量被限制住了。這個(gè)“最長(zhǎng)不超過(guò)最短兩倍”的松散平衡就是紅黑樹(shù)降低旋轉(zhuǎn)頻率的關(guān)鍵。作為對(duì)比AVL 樹(shù)嚴(yán)格要求左右子樹(shù)高度差不超過(guò) 1這種高度平衡在查找時(shí)確實(shí)更優(yōu)但代價(jià)是插入刪除時(shí)的旋轉(zhuǎn)次數(shù)明顯更多。紅黑樹(shù)犧牲了一點(diǎn)點(diǎn)查找性能畢竟是 O(log n) 級(jí)別的松平衡常數(shù)差別不大換來(lái)了插入刪除時(shí)的重平衡操作顯著減少。這就是為什么工程上 map、set、內(nèi)核調(diào)度器普遍選擇紅黑樹(shù)而不是 AVL 樹(shù)的核心原因。2.2 旋轉(zhuǎn)、變色與插入刪除的整體流程讀源碼之前腦子里必須先把旋轉(zhuǎn)和變色這招學(xué)會(huì)。旋轉(zhuǎn)是紅黑樹(shù)調(diào)整結(jié)構(gòu)的基本單位分左旋和右旋兩種。左旋就是某個(gè)節(jié)點(diǎn)下沉為左子節(jié)點(diǎn)它的右孩子上升為父節(jié)點(diǎn)右旋方向相反。旋轉(zhuǎn)過(guò)程會(huì)改變樹(shù)的結(jié)構(gòu)但不會(huì)改變中序遍歷的順序所以它不破壞二叉搜索樹(shù)的有序性。這正是旋轉(zhuǎn)能用來(lái)調(diào)整平衡而不打亂數(shù)據(jù)順序的根本原因。插入的整體流程是這樣的先按照普通二叉搜索樹(shù)的規(guī)則把新節(jié)點(diǎn)放到合適的位置然后把新節(jié)點(diǎn)染成紅色再沿著父節(jié)點(diǎn)向上修復(fù)。為什么要染紅因?yàn)椴迦胍粋€(gè)紅色節(jié)點(diǎn)只可能破壞“不能有連續(xù)紅色節(jié)點(diǎn)”這一條規(guī)則而不會(huì)破壞“黑高相等”。如果插入的是黑色節(jié)點(diǎn)那么從根到這條新路徑上的黑色節(jié)點(diǎn)數(shù)就會(huì)比別的路徑多 1整棵子樹(shù)的黑高全部不匹配修復(fù)起來(lái)極其麻煩。所以“默認(rèn)染紅”是工程上的最優(yōu)選擇。插入修復(fù)分幾種情況如果父節(jié)點(diǎn)是黑色直接結(jié)束什么也不用做如果父節(jié)點(diǎn)是紅色就得看叔父節(jié)點(diǎn)的顏色。叔父是紅色就做顏色翻轉(zhuǎn)父和叔變黑、祖父變紅然后把祖父當(dāng)成新節(jié)點(diǎn)繼續(xù)向上檢查叔父是黑色就通過(guò)旋轉(zhuǎn)來(lái)調(diào)整分成左左、左右、右右、右左四種情況本質(zhì)是先旋轉(zhuǎn)成“一條線”再旋轉(zhuǎn)加變色。理解了這四種情況你會(huì)發(fā)現(xiàn)它們只是對(duì)稱變換記住一種就行。刪除的流程要復(fù)雜一個(gè)量級(jí)。因?yàn)閯h掉一個(gè)節(jié)點(diǎn)之后如果它原本是黑色那么某些路徑上的黑高就會(huì)少 1出現(xiàn)“雙黑”問(wèn)題。刪除修復(fù)的核心就是處理這個(gè)“雙黑”節(jié)點(diǎn)。如果兄弟節(jié)點(diǎn)是紅色先通過(guò)旋轉(zhuǎn)把兄弟變黑如果兄弟是黑色且它的兩個(gè)子節(jié)點(diǎn)都是黑色就把兄弟染紅讓雙黑向上冒泡如果兄弟是黑色且它的右子節(jié)點(diǎn)是紅色就可以直接旋轉(zhuǎn)加變色收尾。這些情況我在后面的源碼拆解里會(huì)詳細(xì)展開(kāi)。2.3 為什么新節(jié)點(diǎn)必須染紅一個(gè)容易被忽略的底層邏輯我見(jiàn)過(guò)不少人讀紅黑樹(shù)源碼讀到 rb_insert_color 里有個(gè)顏色翻轉(zhuǎn)的分支時(shí)就很困惑為什么新節(jié)點(diǎn)一開(kāi)始要是紅色的干脆讓它是黑色的不是少一次修復(fù)嗎這里面的關(guān)鍵在“黑高”這個(gè)概念上。紅黑樹(shù)的所有葉子節(jié)點(diǎn)是 NIL 空節(jié)點(diǎn)它們都是黑色。當(dāng)你插入一個(gè)黑色節(jié)點(diǎn)時(shí)從根到這條新路徑上的黑色節(jié)點(diǎn)數(shù)量會(huì)比同一棵子樹(shù)下其他路徑多出一個(gè)這就直接破壞了第五條不變量。而第五條不變量的破壞是“結(jié)構(gòu)性”的它會(huì)波及到所有包含這條路徑的祖先節(jié)點(diǎn)修復(fù)時(shí)往往需要一路上溯到根節(jié)點(diǎn)代價(jià)極大。相比之下插入紅色節(jié)點(diǎn)影響的只是局部是否存在連續(xù)紅節(jié)點(diǎn)這是可以沿著祖先鏈一路檢查、通過(guò)變色和旋轉(zhuǎn)快速修復(fù)的。簡(jiǎn)單說(shuō)紅色節(jié)點(diǎn)的問(wèn)題是“點(diǎn)狀”的黑色節(jié)點(diǎn)的問(wèn)題是“面狀”的。內(nèi)核源碼里那幾行顏色判斷就是建立在這個(gè)取舍之上。3. 內(nèi)核 rbtree 源碼里的工程智慧3.1 struct rb_node 的低 bit 顏色存儲(chǔ)打開(kāi) include/linux/rbtree.h第一眼看到的就是這個(gè)結(jié)構(gòu)體struct rb_node { unsigned long __rb_parent_color; struct rb_node *rb_right; struct rb_node *rb_left; } __attribute__((aligned(sizeof(long))));很多初讀內(nèi)核源碼的人會(huì)愣住父指針和顏色怎么塞在一個(gè)字段里答案就在對(duì)齊上。在 64 位系統(tǒng)里malloc 或 kmalloc 返回的地址通常是 16 字節(jié)甚至是更嚴(yán)格的對(duì)齊這意味著rb_node 指針的低 4 位二進(jìn)制必然是 0。既然最低一位本來(lái)就沒(méi)人用那就拿它來(lái)存顏色約定最低位為 0 表示紅色最低位為 1 表示黑色。這樣設(shè)計(jì)的好處是實(shí)實(shí)在在的。一個(gè) rb_node 在 64 位系統(tǒng)下只占 24 字節(jié)兩個(gè)指針加一個(gè) unsigned long如果你單獨(dú)加一個(gè) int 字段存顏色整個(gè)結(jié)構(gòu)體可能因?yàn)樘畛渲苯幼兂?32 字節(jié)內(nèi)存開(kāi)銷多出三分之一。內(nèi)核里可能有幾十萬(wàn)個(gè) rb_node 節(jié)點(diǎn)共存這個(gè)節(jié)省非?？捎^。對(duì)應(yīng)的輔助內(nèi)聯(lián)函數(shù)也很有趣。讀父節(jié)點(diǎn)指針時(shí)要把最低位遮掉#define rb_parent(r) ((struct rb_node *)((r)-__rb_parent_color ~3))注意這里遮掉了低 2 位而不是只遮 1 位。早期內(nèi)核只用了 1 位存顏色后來(lái)為了給 rb_augmented增強(qiáng)紅黑樹(shù)留空間統(tǒng)一遮到 3。這個(gè)宏還專門(mén)加了 READ_ONCE/WRITE_ONCE 之類的內(nèi)存屏障處理在并發(fā)場(chǎng)景下保證讀到的一致值這套細(xì)節(jié)在用戶態(tài)實(shí)現(xiàn)里往往被簡(jiǎn)化掉。3.2 rb_insert_color 的插入修復(fù)源碼拆解內(nèi)核里插入新節(jié)點(diǎn)第一步是主動(dòng)建立鏈接關(guān)系把新節(jié)點(diǎn)接進(jìn)樹(shù)里然后調(diào)用 rb_insert_color 修復(fù)顏色。rb_link_node 做的事情很簡(jiǎn)單static inline void rb_link_node(struct rb_node *node, struct rb_node *parent, struct rb_node **rb_link) { node-__rb_parent_color (unsigned long)parent; node-rb_left node-rb_right NULL; *rb_link node; }注意這里新節(jié)點(diǎn)的 __rb_parent_color 就是父節(jié)點(diǎn)指針本身顏色位是 0也就是紅色。這正印證了上面說(shuō)的“默認(rèn)染紅”。接下來(lái)看插入修復(fù)的主邏輯我用偽代碼把內(nèi)核的核心流程還原一下void rb_insert_color(struct rb_node *node, struct rb_root *root) { struct rb_node *parent rb_red_parent(node), *gparent, *tmp; while (parent) { // 父節(jié)點(diǎn)是黑色整棵樹(shù)合法直接返回 if (!rb_is_red(parent)) return; // 走到這里父節(jié)點(diǎn)是紅色需要處理連續(xù)紅節(jié)點(diǎn) gparent rb_red_parent(parent); if (parent gparent-rb_left) { tmp gparent-rb_right; if (tmp rb_is_red(tmp)) { // 叔父是紅色變色上溯 rb_set_parent_color(tmp, gparent, RB_BLACK); rb_set_parent_color(parent, gparent, RB_BLACK); node gparent; continue; } // 叔父是黑色先調(diào)整方向再旋轉(zhuǎn) if (parent-rb_right node) { // 左右情況先左旋 tmp parent-rb_right; parent-rb_right node-rb_left; ... // 交換 parent 和 node 的角色 } // 左左情況右旋 變色 rb_set_parent_color(parent, gparent, RB_BLACK); rb_set_parent_color(gparent, parent, RB_RED); __rb_rotate_left(gparent, root); return; } else { // 對(duì)稱處理右邊 } } // 根節(jié)點(diǎn)強(qiáng)制染黑 WRITE_ONCE(root-rb_node-__rb_parent_color, RB_BLACK); }這段代碼最值得玩味的是那個(gè) while 循環(huán)的出口條件。循環(huán)里如果一直遇到叔父紅色就會(huì)一路把祖父染紅、自己上溯到祖父節(jié)點(diǎn)繼續(xù)上一層的判斷。直到某次父節(jié)點(diǎn)是黑色循環(huán)退出或者一路沖上根節(jié)點(diǎn)最后一句強(qiáng)制把根染黑。根節(jié)點(diǎn)為什么必須是黑的因?yàn)槿绻羌t的它的兩個(gè)子節(jié)點(diǎn)如果是紅的就會(huì)形成連續(xù)紅節(jié)點(diǎn)而且根節(jié)點(diǎn)染黑不改變?nèi)魏温窂降暮诟咚宰詈笠坏辣ｋU(xiǎn)就是無(wú)條件讓根變黑。你可能注意到我這里刻意省略了旋轉(zhuǎn)的部分細(xì)節(jié)因?yàn)樾D(zhuǎn)的具體代碼比較長(zhǎng)。內(nèi)核把旋轉(zhuǎn)封裝成了 __rb_rotate_left 和 __rb_rotate_right 兩個(gè)函數(shù)它們的共同特點(diǎn)是傳入 root 指針因?yàn)樾D(zhuǎn)可能會(huì)導(dǎo)致新的子樹(shù)根節(jié)點(diǎn)需要把新的根節(jié)點(diǎn)掛回全局根節(jié)點(diǎn)。這也是后來(lái)很多人移植內(nèi)核 rbtree 時(shí)最容易出錯(cuò)的地方忘記更新 root 指針。3.3 ____rb_erase_color 刪除修復(fù)源碼拆解刪除在 rb_erase 里分成兩部分。第一部分是找后繼節(jié)點(diǎn)、把值搬過(guò)去、摘除物理節(jié)點(diǎn)第二部分是如果摘除的節(jié)點(diǎn)是黑色就調(diào)用 ____rb_erase_color 修復(fù)黑高。找后繼的邏輯在 rb_next它的實(shí)現(xiàn)非常精妙struct rb_node *rb_next(const struct rb_node *node) { if (RB_EMPTY_NODE(node)) return NULL; if (node-rb_right) { // 有右子樹(shù)找右子樹(shù)的最左節(jié)點(diǎn) node node-rb_right; while (node-rb_left) node node-rb_left; return (struct rb_node *)node; } // 沒(méi)有右子樹(shù)向上找第一個(gè)“從左子樹(shù)走上來(lái)的祖先” while (rb_parent(node) node rb_parent(node)-rb_right) node rb_parent(node); return rb_parent(node); }這里有一個(gè)內(nèi)核特有的宏 RB_EMPTY_NODE它判斷的是 node 的 __rb_parent_color 是否等于 node 自身。如果節(jié)點(diǎn)被移除后內(nèi)核會(huì)把它的父指針指向自己形成一種特殊標(biāo)記表示這個(gè)節(jié)點(diǎn)不再屬于任何樹(shù)。這個(gè)設(shè)計(jì)在用戶態(tài)很少見(jiàn)到但本質(zhì)上是一種廉價(jià)的狀態(tài)標(biāo)記。刪除修復(fù)的核心是 ____rb_erase_color。這個(gè)函數(shù)代碼很長(zhǎng)我強(qiáng)烈建議你配合注釋讀原文。它的核心思路是把“被刪節(jié)點(diǎn)是黑色”造成的問(wèn)題抽象成當(dāng)前子樹(shù)少了一個(gè)黑色節(jié)點(diǎn)我們需要通過(guò)旋轉(zhuǎn)和變色讓其他路徑“勻”一個(gè)黑色節(jié)點(diǎn)過(guò)來(lái)。我把四種情況總結(jié)成一張速查表場(chǎng)景兄弟節(jié)點(diǎn)顏色兄弟的子節(jié)點(diǎn)情況處理動(dòng)作情況1紅色任意旋轉(zhuǎn)把兄弟變成黑色轉(zhuǎn)化成情況2/3/4情況2黑色兩個(gè)都是黑色兄弟染紅問(wèn)題向上冒泡一層情況3黑色左子紅、右子黑右旋兄弟子樹(shù)把紅色節(jié)點(diǎn)轉(zhuǎn)到右側(cè)情況4黑色右子紅左旋兄弟上位完成收尾這個(gè)表看起來(lái)簡(jiǎn)單實(shí)際寫(xiě)代碼時(shí)每一步都要仔細(xì)更新父子關(guān)系。內(nèi)核源碼里的寫(xiě)法是先把兄弟節(jié)點(diǎn)取出來(lái)根據(jù)兄弟和侄子們的顏色分支出處理最后通過(guò) __rb_rotate_set_parents 這個(gè)函數(shù)一段一段地調(diào)整。真正的工程代碼里全是位運(yùn)算和內(nèi)聯(lián)函數(shù)讀的時(shí)候要有耐心。3.4 rb_augmented 與區(qū)間管理新版內(nèi)核 rbtree 還有一個(gè)增強(qiáng)特性rb_augmented。它允許每個(gè)節(jié)點(diǎn)額外維護(hù)一些聚合信息比如子樹(shù)里的最大范圍、最小值等。以虛擬內(nèi)存管理為例內(nèi)核通過(guò)紅黑樹(shù)維護(hù)進(jìn)程的 VMA 區(qū)域每個(gè)節(jié)點(diǎn)存一個(gè) [start, end] 區(qū)間而聚合信息可以讓內(nèi)核在查找“包含某地址的 VMA”時(shí)快速跳過(guò)整棵不可能命中的子樹(shù)優(yōu)化查找效率。增強(qiáng) rbtree 的使用方式是在調(diào)用 rb_insert_augmented 和 rb_erase_augmented 時(shí)傳入一個(gè) rb_augment_callbacks 結(jié)構(gòu)體里面包含 rotate 和 propagate 兩個(gè)回調(diào)函數(shù)。rotate 負(fù)責(zé)在旋轉(zhuǎn)時(shí)更新子樹(shù)聚合信息propagate 負(fù)責(zé)在節(jié)點(diǎn)上溯時(shí)把新的聚合值傳遞給祖先。這套機(jī)制讓普通紅黑樹(shù)變成了一種“可擴(kuò)展平衡區(qū)間樹(shù)”思想很值得借鑒。4. 徒手實(shí)現(xiàn)一個(gè)迷你 rb_tree 源碼4.1 數(shù)據(jù)結(jié)構(gòu)定義讀源碼是一回事真正自己寫(xiě)一遍又是另一回事。我在工程里用過(guò)內(nèi)核版 rbtree也在面試前手寫(xiě)過(guò)迷你版。下面給出一個(gè)可直接運(yùn)行的迷你 C 實(shí)現(xiàn)核心方便你對(duì)照源碼理解。typedef enum { RB_RED 0, RB_BLACK 1 } rb_color; typedef struct rb_node { struct rb_node *parent; struct rb_node *left; struct rb_node *right; rb_color color; int key; // 以便驗(yàn)證實(shí)際工程中通常是結(jié)構(gòu)體內(nèi)的數(shù)據(jù) } rb_node; typedef struct { rb_node *root; } rb_tree;這個(gè)定義簡(jiǎn)化了內(nèi)核的低 bit 顏色存儲(chǔ)直接用一個(gè) enum 字段邏輯更清晰。如果你想進(jìn)階挑戰(zhàn)完全可以按照內(nèi)核的方式把 parent 和 color 合并成一個(gè) unsigned long那樣內(nèi)存效率更高也更貼近真實(shí)工程。4.2 左旋右旋與插入實(shí)現(xiàn)左旋操作的要點(diǎn)是拿到當(dāng)前節(jié)點(diǎn)的右孩子讓右孩子上位當(dāng)前節(jié)點(diǎn)成為右孩子的左孩子同時(shí)處理右孩子的左子樹(shù)掛到當(dāng)前節(jié)點(diǎn)的右孩子位置。寫(xiě)成代碼static void rb_rotate_left(rb_tree *tree, rb_node *node) { rb_node *r node-right; node-right r-left; if (r-left) r-left-parent node; r-parent node-parent; if (node-parent NULL) tree-root r; else if (node node-parent-left) node-parent-left r; else node-parent-right r; r-left node; node-parent r; }右旋完全對(duì)稱不再贅述。插入時(shí)先按普通 BST 規(guī)則找位置然后執(zhí)行修復(fù)函數(shù)。修復(fù)函數(shù)對(duì)照內(nèi)核版可以精簡(jiǎn)成這樣static void rb_insert_fixup(rb_tree *tree, rb_node *node) { while (node ! tree-root node-parent-color RB_RED) { if (node-parent node-parent-parent-left) { rb_node *uncle node-parent-parent-right; if (uncle uncle-color RB_RED) { node-parent-color RB_BLACK; uncle-color RB_BLACK; node-parent-parent-color RB_RED; node node-parent-parent; } else { if (node node-parent-right) { node node-parent; rb_rotate_left(tree, node); } node-parent-color RB_BLACK; node-parent-parent-color RB_RED; rb_rotate_right(tree, node-parent-parent); } } else { // 對(duì)稱處理 } } tree-root-color RB_BLACK; }這里最關(guān)鍵的細(xì)節(jié)是判斷叔父顏色時(shí)如果叔父是 NULL也當(dāng)作黑色處理。很多初寫(xiě)者在這里只判斷非空就訪問(wèn) color結(jié)果空指針解引用實(shí)際 NIL 節(jié)點(diǎn)是黑色NULL 就是黑。4.3 刪除與修復(fù)刪除修復(fù)的迷你版實(shí)現(xiàn)我建議直接參考內(nèi)核的 ____rb_erase_color 改寫(xiě)。核心是維護(hù)一個(gè)節(jié)點(diǎn) x表示“雙黑”節(jié)點(diǎn)所在位置。這里給出修復(fù)主循環(huán)的骨架static void rb_delete_fixup(rb_tree *tree, rb_node *node, rb_node *parent) { while (node ! tree-root (node NULL || node-color RB_BLACK)) { if (node parent-left) { rb_node *sibling parent-right; if (sibling sibling-color RB_RED) { sibling-color RB_BLACK; parent-color RB_RED; rb_rotate_left(tree, parent); sibling parent-right; } if ((sibling-left NULL || sibling-left-color RB_BLACK) (sibling-right NULL || sibling-right-color RB_BLACK)) { sibling-color RB_RED; node parent; parent parent-parent; } else { if (sibling-right NULL || sibling-right-color RB_BLACK) { sibling-left-color RB_BLACK; sibling-color RB_RED; rb_rotate_right(tree, sibling); sibling parent-right; } sibling-color parent-color; parent-color RB_BLACK; sibling-right-color RB_BLACK; rb_rotate_left(tree, parent); node tree-root; break; } } else { // 對(duì)稱處理 } } if (node) node-color RB_BLACK; }刪除修復(fù)是最容易寫(xiě)錯(cuò)的部分因?yàn)榍闆r分支多且每一輪循環(huán)后 node 和 parent 的指向都會(huì)變化。我的經(jīng)驗(yàn)是寫(xiě)完之后一定要配合隨機(jī)插入刪除的驗(yàn)證程序跑一輪 fuzz光靠肉眼檢查基本不可能保證正確。4.4 中序遍歷驗(yàn)證正確性寫(xiě)完插入刪除第一件事不是看平衡而是驗(yàn)證中序遍歷是否嚴(yán)格有序。中序遍歷可以這樣寫(xiě)static void rb_inorder(rb_node *node, void (*visit)(rb_node *)) { if (!node) return; rb_inorder(node-left, visit); visit(node); rb_inorder(node-right, visit); }如果遍歷輸出的 key 序列是遞增的說(shuō)明 BST 結(jié)構(gòu)沒(méi)有被破壞。緊接著再檢查紅黑性質(zhì)根節(jié)點(diǎn)是黑、沒(méi)有連續(xù)紅節(jié)點(diǎn)、每條路徑黑高相等。這三個(gè)檢查加在一起基本能保證實(shí)現(xiàn)沒(méi)有結(jié)構(gòu)性問(wèn)題。5. 把內(nèi)核 rbtree 搬到自己的工程里5.1 移植步驟與注意事項(xiàng)內(nèi)核的 rbtree 實(shí)現(xiàn)質(zhì)量很高很多人想直接在用戶態(tài)或者嵌入式工程里用。直接復(fù)制 rbtree.c 和 rbtree.h 肯定編譯不過(guò)因?yàn)槔锩娉涑庵鴥?nèi)核特有的宏。我踩過(guò)不少坑整理一下移植步驟第一步把內(nèi)核頭文件里依賴的宏統(tǒng)一替換掉。常見(jiàn)的有 BUG_ON 替換成 assert、WARN_ON 替換成打印、READ_ONCE/WRITE_ONCE 直接去掉或者換成普通賦值、unlikely/likely 直接刪掉。第二步處理 struct rb_node 的對(duì)齊屬性。內(nèi)核里用attribute((aligned(sizeof(long))))用戶態(tài)也要保留因?yàn)榈?bit 顏色方案依賴指針對(duì)齊。實(shí)際使用中 malloc 返回的地址在對(duì)齊上沒(méi)問(wèn)題但如果你的節(jié)點(diǎn)是從一個(gè) char 數(shù)組緩沖區(qū)偏移出來(lái)的就要小心偏移量是否對(duì)齊。第三步把 rb_root 初始化為 RB_ROOT就像初始化任何數(shù)據(jù)結(jié)構(gòu)一樣。沒(méi)有這句話根節(jié)點(diǎn)指向隨機(jī)地址插入第一個(gè)節(jié)點(diǎn)可能直接崩潰。第四步確保你的宿主結(jié)構(gòu)體里 rb_node 成員是第一個(gè)或者偏移足夠?qū)R。推薦把 rb_node 放在結(jié)構(gòu)體第一個(gè)字段這樣 from 指針就是 rb_node 本身省去 container_of 的偏移計(jì)算。5.2 用戶態(tài)驗(yàn)證與隨機(jī) fuzz移植完成以后寫(xiě)一個(gè)驗(yàn)證程序。我一般這么干隨機(jī)插入一百萬(wàn)個(gè)整數(shù)插入過(guò)程中每十萬(wàn)次中序遍歷一次檢查是否有序然后隨機(jī)刪除一半節(jié)點(diǎn)刪除過(guò)程再檢查黑高和連續(xù)紅約束最后整棵樹(shù)刪空確保沒(méi)有內(nèi)存泄漏。這個(gè) fuzz 程序看起來(lái)不起眼但它是檢驗(yàn)紅黑樹(shù)實(shí)現(xiàn)是否正確的試金石。我自己手寫(xiě)迷你版的時(shí)候就是在隨機(jī)刪除那一步暴露出問(wèn)題刪除修復(fù)里情況3轉(zhuǎn)情況4時(shí)忘了更新兄弟節(jié)點(diǎn)指針導(dǎo)致后續(xù)判斷用了舊節(jié)點(diǎn)樹(shù)結(jié)構(gòu)直接錯(cuò)亂。這種bug靠讀代碼極難發(fā)現(xiàn)靠 fuzz 一跑就現(xiàn)形。6. 常見(jiàn)問(wèn)題與排查技巧實(shí)錄6.1 三個(gè)高頻問(wèn)題定位思路移植和使用 rb_tree 過(guò)程中下面這些問(wèn)題出現(xiàn)頻率最高問(wèn)題一插入后根節(jié)點(diǎn)變了樹(shù)變得混亂。多半是更新根節(jié)點(diǎn)失敗。旋轉(zhuǎn)函數(shù)里要判斷 node-parent 是否為 NULL如果是則應(yīng)該把新的子樹(shù)根節(jié)點(diǎn)賦給 tree-root。寫(xiě)代碼時(shí)最容易漏掉這個(gè)判斷或者判斷寫(xiě)反了。問(wèn)題二低 bit 顏色方案里節(jié)點(diǎn)顏色讀取異常。檢查一下你取的節(jié)點(diǎn)地址是否真的對(duì)齊到 4 字節(jié)。如果自己寫(xiě)內(nèi)存池或者對(duì)象池分配的地址可能不是 page 對(duì)齊的這時(shí)最低位不是可靠的。穩(wěn)妥做法是頁(yè)面級(jí)對(duì)齊或者干脆用獨(dú)立的 color 字段。問(wèn)題三STL 里 map 刪除了一個(gè)迭代器后續(xù)遍歷崩潰。這不是紅黑樹(shù)實(shí)現(xiàn)問(wèn)題是迭代器失效。紅黑樹(shù)刪除一個(gè)節(jié)點(diǎn)后只有指向被刪節(jié)點(diǎn)的迭代器失效其他迭代器不失效——這恰恰是 STL 選擇紅黑樹(shù)而不是數(shù)組或鏈表作為底層結(jié)構(gòu)的優(yōu)勢(shì)之一。但如果你在遍歷過(guò)程中刪除了當(dāng)前迭代器應(yīng)該先自增保存下一個(gè)迭代器再刪除。6.2 調(diào)試紅黑樹(shù)的三板斧最后分享我調(diào)試紅黑樹(shù)的三板斧。第一板斧是開(kāi)一個(gè) debug 開(kāi)關(guān)每次插入刪除后校驗(yàn)整棵樹(shù)的紅黑性質(zhì)第二板斧是隨機(jī) fuzz同時(shí)記錄每次操作的 key方便崩潰時(shí)用同樣的隨機(jī)種子復(fù)現(xiàn)第三板斧是寫(xiě)一個(gè) dump 函數(shù)輸出整棵樹(shù)的括號(hào)表示或者縮進(jìn)結(jié)構(gòu)用眼睛直觀地看樹(shù)形。尤其是 dump 這個(gè)工具看著樹(shù)的結(jié)構(gòu)一點(diǎn)點(diǎn)調(diào)整很快就能建立直覺(jué)。比如插入觸發(fā)叔父紅變色時(shí)你會(huì)發(fā)現(xiàn)中間幾層同時(shí)變色走旋轉(zhuǎn)分支時(shí)子樹(shù)根節(jié)點(diǎn)會(huì)整體遷移。這種直覺(jué)一旦建立再回頭讀內(nèi)核源碼理解速度會(huì)有質(zhì)變。另外一個(gè)常見(jiàn)的面試加分項(xiàng)是思考紅黑樹(shù)和 B 樹(shù)、跳表的取舍。我會(huì)在面試中這樣回答內(nèi)存中數(shù)據(jù)量中等、插入刪除頻繁、需要穩(wěn)定迭代器時(shí)選紅黑樹(shù)磁盤(pán)存儲(chǔ)、范圍查詢、塊式讀寫(xiě)時(shí)選 B 樹(shù)并發(fā)讀寫(xiě)頻繁且想要簡(jiǎn)單的概率平衡時(shí)選跳表。這個(gè)回答比單純背性質(zhì)有用得多因?yàn)槟惆褦?shù)據(jù)結(jié)構(gòu)和應(yīng)用場(chǎng)景掛上了鉤。我對(duì)紅黑樹(shù)源碼最深的體會(huì)是讀十遍都不如自己動(dòng)手寫(xiě)一遍有用。我建議大家拿到一份源碼不管是我上面這份迷你版還是內(nèi)核原版先自己在本地跑通插入和查找再一步一步加刪除和修復(fù)。寫(xiě)刪除修復(fù)時(shí)卡住很正?？ㄗ≌f(shuō)明你對(duì)“雙黑”這個(gè)概念還沒(méi)有真正吃透這時(shí)候回看 ____rb_erase_color 的四種情況會(huì)有一種豁然開(kāi)朗的感覺(jué)。紅黑樹(shù)源碼這份功課早晚要補(bǔ)趁早補(bǔ)完后面看任何平衡樹(shù)相關(guān)的代碼都會(huì)輕松許多。