11da177e4SLinus Torvalds /* 21da177e4SLinus Torvalds Red Black Trees 31da177e4SLinus Torvalds (C) 1999 Andrea Arcangeli <andrea@suse.de> 41da177e4SLinus Torvalds (C) 2002 David Woodhouse <dwmw2@infradead.org> 546b6135aSMichel Lespinasse (C) 2012 Michel Lespinasse <walken@google.com> 61da177e4SLinus Torvalds 71da177e4SLinus Torvalds This program is free software; you can redistribute it and/or modify 81da177e4SLinus Torvalds it under the terms of the GNU General Public License as published by 91da177e4SLinus Torvalds the Free Software Foundation; either version 2 of the License, or 101da177e4SLinus Torvalds (at your option) any later version. 111da177e4SLinus Torvalds 121da177e4SLinus Torvalds This program is distributed in the hope that it will be useful, 131da177e4SLinus Torvalds but WITHOUT ANY WARRANTY; without even the implied warranty of 141da177e4SLinus Torvalds MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 151da177e4SLinus Torvalds GNU General Public License for more details. 161da177e4SLinus Torvalds 171da177e4SLinus Torvalds You should have received a copy of the GNU General Public License 181da177e4SLinus Torvalds along with this program; if not, write to the Free Software 191da177e4SLinus Torvalds Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA 02111-1307 USA 201da177e4SLinus Torvalds 211da177e4SLinus Torvalds linux/lib/rbtree.c 221da177e4SLinus Torvalds */ 231da177e4SLinus Torvalds 249c079addSMichel Lespinasse #include <linux/rbtree_augmented.h> 258bc3bcc9SPaul Gortmaker #include <linux/export.h> 261da177e4SLinus Torvalds 275bc9188aSMichel Lespinasse /* 285bc9188aSMichel Lespinasse * red-black trees properties: http://en.wikipedia.org/wiki/Rbtree 295bc9188aSMichel Lespinasse * 305bc9188aSMichel Lespinasse * 1) A node is either red or black 315bc9188aSMichel Lespinasse * 2) The root is black 325bc9188aSMichel Lespinasse * 3) All leaves (NULL) are black 335bc9188aSMichel Lespinasse * 4) Both children of every red node are black 345bc9188aSMichel Lespinasse * 5) Every simple path from root to leaves contains the same number 355bc9188aSMichel Lespinasse * of black nodes. 365bc9188aSMichel Lespinasse * 375bc9188aSMichel Lespinasse * 4 and 5 give the O(log n) guarantee, since 4 implies you cannot have two 385bc9188aSMichel Lespinasse * consecutive red nodes in a path and every red node is therefore followed by 395bc9188aSMichel Lespinasse * a black. So if B is the number of black nodes on every simple path (as per 405bc9188aSMichel Lespinasse * 5), then the longest possible path due to 4 is 2B. 415bc9188aSMichel Lespinasse * 425bc9188aSMichel Lespinasse * We shall indicate color with case, where black nodes are uppercase and red 436280d235SMichel Lespinasse * nodes will be lowercase. Unknown color nodes shall be drawn as red within 446280d235SMichel Lespinasse * parentheses and have some accompanying text comment. 455bc9188aSMichel Lespinasse */ 465bc9188aSMichel Lespinasse 47d72da4a4SPeter Zijlstra /* 48d72da4a4SPeter Zijlstra * Notes on lockless lookups: 49d72da4a4SPeter Zijlstra * 50d72da4a4SPeter Zijlstra * All stores to the tree structure (rb_left and rb_right) must be done using 51d72da4a4SPeter Zijlstra * WRITE_ONCE(). And we must not inadvertently cause (temporary) loops in the 52d72da4a4SPeter Zijlstra * tree structure as seen in program order. 53d72da4a4SPeter Zijlstra * 54d72da4a4SPeter Zijlstra * These two requirements will allow lockless iteration of the tree -- not 55d72da4a4SPeter Zijlstra * correct iteration mind you, tree rotations are not atomic so a lookup might 56d72da4a4SPeter Zijlstra * miss entire subtrees. 57d72da4a4SPeter Zijlstra * 58d72da4a4SPeter Zijlstra * But they do guarantee that any such traversal will only see valid elements 59d72da4a4SPeter Zijlstra * and that it will indeed complete -- does not get stuck in a loop. 60d72da4a4SPeter Zijlstra * 61d72da4a4SPeter Zijlstra * It also guarantees that if the lookup returns an element it is the 'correct' 62d72da4a4SPeter Zijlstra * one. But not returning an element does _NOT_ mean it's not present. 63d72da4a4SPeter Zijlstra * 64d72da4a4SPeter Zijlstra * NOTE: 65d72da4a4SPeter Zijlstra * 66d72da4a4SPeter Zijlstra * Stores to __rb_parent_color are not important for simple lookups so those 67d72da4a4SPeter Zijlstra * are left undone as of now. Nor did I check for loops involving parent 68d72da4a4SPeter Zijlstra * pointers. 69d72da4a4SPeter Zijlstra */ 70d72da4a4SPeter Zijlstra 7146b6135aSMichel Lespinasse static inline void rb_set_black(struct rb_node *rb) 7246b6135aSMichel Lespinasse { 7346b6135aSMichel Lespinasse rb->__rb_parent_color |= RB_BLACK; 7446b6135aSMichel Lespinasse } 7546b6135aSMichel Lespinasse 765bc9188aSMichel Lespinasse static inline struct rb_node *rb_red_parent(struct rb_node *red) 775bc9188aSMichel Lespinasse { 785bc9188aSMichel Lespinasse return (struct rb_node *)red->__rb_parent_color; 795bc9188aSMichel Lespinasse } 805bc9188aSMichel Lespinasse 815bc9188aSMichel Lespinasse /* 825bc9188aSMichel Lespinasse * Helper function for rotations: 835bc9188aSMichel Lespinasse * - old's parent and color get assigned to new 845bc9188aSMichel Lespinasse * - old gets assigned new as a parent and 'color' as a color. 855bc9188aSMichel Lespinasse */ 865bc9188aSMichel Lespinasse static inline void 875bc9188aSMichel Lespinasse __rb_rotate_set_parents(struct rb_node *old, struct rb_node *new, 885bc9188aSMichel Lespinasse struct rb_root *root, int color) 895bc9188aSMichel Lespinasse { 905bc9188aSMichel Lespinasse struct rb_node *parent = rb_parent(old); 915bc9188aSMichel Lespinasse new->__rb_parent_color = old->__rb_parent_color; 925bc9188aSMichel Lespinasse rb_set_parent_color(old, new, color); 937abc704aSMichel Lespinasse __rb_change_child(old, new, parent, root); 945bc9188aSMichel Lespinasse } 955bc9188aSMichel Lespinasse 9614b94af0SMichel Lespinasse static __always_inline void 9714b94af0SMichel Lespinasse __rb_insert(struct rb_node *node, struct rb_root *root, 98*cd9e61edSDavidlohr Bueso bool newleft, struct rb_node **leftmost, 9914b94af0SMichel Lespinasse void (*augment_rotate)(struct rb_node *old, struct rb_node *new)) 1001da177e4SLinus Torvalds { 1015bc9188aSMichel Lespinasse struct rb_node *parent = rb_red_parent(node), *gparent, *tmp; 1021da177e4SLinus Torvalds 103*cd9e61edSDavidlohr Bueso if (newleft) 104*cd9e61edSDavidlohr Bueso *leftmost = node; 105*cd9e61edSDavidlohr Bueso 1066d58452dSMichel Lespinasse while (true) { 1076d58452dSMichel Lespinasse /* 1086d58452dSMichel Lespinasse * Loop invariant: node is red 1096d58452dSMichel Lespinasse * 1106d58452dSMichel Lespinasse * If there is a black parent, we are done. 1116d58452dSMichel Lespinasse * Otherwise, take some corrective action as we don't 1126d58452dSMichel Lespinasse * want a red root or two consecutive red nodes. 1136d58452dSMichel Lespinasse */ 1146d58452dSMichel Lespinasse if (!parent) { 1155bc9188aSMichel Lespinasse rb_set_parent_color(node, NULL, RB_BLACK); 1166d58452dSMichel Lespinasse break; 1176d58452dSMichel Lespinasse } else if (rb_is_black(parent)) 1186d58452dSMichel Lespinasse break; 1196d58452dSMichel Lespinasse 1205bc9188aSMichel Lespinasse gparent = rb_red_parent(parent); 1211da177e4SLinus Torvalds 1225bc9188aSMichel Lespinasse tmp = gparent->rb_right; 12359633abfSMichel Lespinasse if (parent != tmp) { /* parent == gparent->rb_left */ 1245bc9188aSMichel Lespinasse if (tmp && rb_is_red(tmp)) { 1255bc9188aSMichel Lespinasse /* 1265bc9188aSMichel Lespinasse * Case 1 - color flips 1275bc9188aSMichel Lespinasse * 1285bc9188aSMichel Lespinasse * G g 1295bc9188aSMichel Lespinasse * / \ / \ 1305bc9188aSMichel Lespinasse * p u --> P U 1315bc9188aSMichel Lespinasse * / / 1321b9c53e8SWei Yang * n n 1335bc9188aSMichel Lespinasse * 1345bc9188aSMichel Lespinasse * However, since g's parent might be red, and 1355bc9188aSMichel Lespinasse * 4) does not allow this, we need to recurse 1365bc9188aSMichel Lespinasse * at g. 1375bc9188aSMichel Lespinasse */ 1385bc9188aSMichel Lespinasse rb_set_parent_color(tmp, gparent, RB_BLACK); 1395bc9188aSMichel Lespinasse rb_set_parent_color(parent, gparent, RB_BLACK); 1401da177e4SLinus Torvalds node = gparent; 1415bc9188aSMichel Lespinasse parent = rb_parent(node); 1425bc9188aSMichel Lespinasse rb_set_parent_color(node, parent, RB_RED); 1431da177e4SLinus Torvalds continue; 1441da177e4SLinus Torvalds } 1451da177e4SLinus Torvalds 14659633abfSMichel Lespinasse tmp = parent->rb_right; 14759633abfSMichel Lespinasse if (node == tmp) { 1485bc9188aSMichel Lespinasse /* 1495bc9188aSMichel Lespinasse * Case 2 - left rotate at parent 1505bc9188aSMichel Lespinasse * 1515bc9188aSMichel Lespinasse * G G 1525bc9188aSMichel Lespinasse * / \ / \ 1535bc9188aSMichel Lespinasse * p U --> n U 1545bc9188aSMichel Lespinasse * \ / 1555bc9188aSMichel Lespinasse * n p 1565bc9188aSMichel Lespinasse * 1575bc9188aSMichel Lespinasse * This still leaves us in violation of 4), the 1585bc9188aSMichel Lespinasse * continuation into Case 3 will fix that. 1595bc9188aSMichel Lespinasse */ 160d72da4a4SPeter Zijlstra tmp = node->rb_left; 161d72da4a4SPeter Zijlstra WRITE_ONCE(parent->rb_right, tmp); 162d72da4a4SPeter Zijlstra WRITE_ONCE(node->rb_left, parent); 1635bc9188aSMichel Lespinasse if (tmp) 1645bc9188aSMichel Lespinasse rb_set_parent_color(tmp, parent, 1655bc9188aSMichel Lespinasse RB_BLACK); 1665bc9188aSMichel Lespinasse rb_set_parent_color(parent, node, RB_RED); 16714b94af0SMichel Lespinasse augment_rotate(parent, node); 1681da177e4SLinus Torvalds parent = node; 16959633abfSMichel Lespinasse tmp = node->rb_right; 1701da177e4SLinus Torvalds } 1711da177e4SLinus Torvalds 1725bc9188aSMichel Lespinasse /* 1735bc9188aSMichel Lespinasse * Case 3 - right rotate at gparent 1745bc9188aSMichel Lespinasse * 1755bc9188aSMichel Lespinasse * G P 1765bc9188aSMichel Lespinasse * / \ / \ 1775bc9188aSMichel Lespinasse * p U --> n g 1785bc9188aSMichel Lespinasse * / \ 1795bc9188aSMichel Lespinasse * n U 1805bc9188aSMichel Lespinasse */ 181d72da4a4SPeter Zijlstra WRITE_ONCE(gparent->rb_left, tmp); /* == parent->rb_right */ 182d72da4a4SPeter Zijlstra WRITE_ONCE(parent->rb_right, gparent); 1835bc9188aSMichel Lespinasse if (tmp) 1845bc9188aSMichel Lespinasse rb_set_parent_color(tmp, gparent, RB_BLACK); 1855bc9188aSMichel Lespinasse __rb_rotate_set_parents(gparent, parent, root, RB_RED); 18614b94af0SMichel Lespinasse augment_rotate(gparent, parent); 1871f052865SMichel Lespinasse break; 1881da177e4SLinus Torvalds } else { 1895bc9188aSMichel Lespinasse tmp = gparent->rb_left; 1905bc9188aSMichel Lespinasse if (tmp && rb_is_red(tmp)) { 1915bc9188aSMichel Lespinasse /* Case 1 - color flips */ 1925bc9188aSMichel Lespinasse rb_set_parent_color(tmp, gparent, RB_BLACK); 1935bc9188aSMichel Lespinasse rb_set_parent_color(parent, gparent, RB_BLACK); 1941da177e4SLinus Torvalds node = gparent; 1955bc9188aSMichel Lespinasse parent = rb_parent(node); 1965bc9188aSMichel Lespinasse rb_set_parent_color(node, parent, RB_RED); 1971da177e4SLinus Torvalds continue; 1981da177e4SLinus Torvalds } 1991da177e4SLinus Torvalds 20059633abfSMichel Lespinasse tmp = parent->rb_left; 20159633abfSMichel Lespinasse if (node == tmp) { 2025bc9188aSMichel Lespinasse /* Case 2 - right rotate at parent */ 203d72da4a4SPeter Zijlstra tmp = node->rb_right; 204d72da4a4SPeter Zijlstra WRITE_ONCE(parent->rb_left, tmp); 205d72da4a4SPeter Zijlstra WRITE_ONCE(node->rb_right, parent); 2065bc9188aSMichel Lespinasse if (tmp) 2075bc9188aSMichel Lespinasse rb_set_parent_color(tmp, parent, 2085bc9188aSMichel Lespinasse RB_BLACK); 2095bc9188aSMichel Lespinasse rb_set_parent_color(parent, node, RB_RED); 21014b94af0SMichel Lespinasse augment_rotate(parent, node); 2111da177e4SLinus Torvalds parent = node; 21259633abfSMichel Lespinasse tmp = node->rb_left; 2131da177e4SLinus Torvalds } 2141da177e4SLinus Torvalds 2155bc9188aSMichel Lespinasse /* Case 3 - left rotate at gparent */ 216d72da4a4SPeter Zijlstra WRITE_ONCE(gparent->rb_right, tmp); /* == parent->rb_left */ 217d72da4a4SPeter Zijlstra WRITE_ONCE(parent->rb_left, gparent); 2185bc9188aSMichel Lespinasse if (tmp) 2195bc9188aSMichel Lespinasse rb_set_parent_color(tmp, gparent, RB_BLACK); 2205bc9188aSMichel Lespinasse __rb_rotate_set_parents(gparent, parent, root, RB_RED); 22114b94af0SMichel Lespinasse augment_rotate(gparent, parent); 2221f052865SMichel Lespinasse break; 2231da177e4SLinus Torvalds } 2241da177e4SLinus Torvalds } 2251da177e4SLinus Torvalds } 2261da177e4SLinus Torvalds 2273cb7a563SMichel Lespinasse /* 2283cb7a563SMichel Lespinasse * Inline version for rb_erase() use - we want to be able to inline 2293cb7a563SMichel Lespinasse * and eliminate the dummy_rotate callback there 2303cb7a563SMichel Lespinasse */ 2313cb7a563SMichel Lespinasse static __always_inline void 2323cb7a563SMichel Lespinasse ____rb_erase_color(struct rb_node *parent, struct rb_root *root, 2339c079addSMichel Lespinasse void (*augment_rotate)(struct rb_node *old, struct rb_node *new)) 2341da177e4SLinus Torvalds { 23546b6135aSMichel Lespinasse struct rb_node *node = NULL, *sibling, *tmp1, *tmp2; 2361da177e4SLinus Torvalds 237d6ff1273SMichel Lespinasse while (true) { 238d6ff1273SMichel Lespinasse /* 23946b6135aSMichel Lespinasse * Loop invariants: 24046b6135aSMichel Lespinasse * - node is black (or NULL on first iteration) 24146b6135aSMichel Lespinasse * - node is not the root (parent is not NULL) 24246b6135aSMichel Lespinasse * - All leaf paths going through parent and node have a 243d6ff1273SMichel Lespinasse * black node count that is 1 lower than other leaf paths. 244d6ff1273SMichel Lespinasse */ 2456280d235SMichel Lespinasse sibling = parent->rb_right; 24659633abfSMichel Lespinasse if (node != sibling) { /* node == parent->rb_left */ 2476280d235SMichel Lespinasse if (rb_is_red(sibling)) { 2486280d235SMichel Lespinasse /* 2496280d235SMichel Lespinasse * Case 1 - left rotate at parent 2506280d235SMichel Lespinasse * 2516280d235SMichel Lespinasse * P S 2526280d235SMichel Lespinasse * / \ / \ 2536280d235SMichel Lespinasse * N s --> p Sr 2546280d235SMichel Lespinasse * / \ / \ 2556280d235SMichel Lespinasse * Sl Sr N Sl 2566280d235SMichel Lespinasse */ 257d72da4a4SPeter Zijlstra tmp1 = sibling->rb_left; 258d72da4a4SPeter Zijlstra WRITE_ONCE(parent->rb_right, tmp1); 259d72da4a4SPeter Zijlstra WRITE_ONCE(sibling->rb_left, parent); 2606280d235SMichel Lespinasse rb_set_parent_color(tmp1, parent, RB_BLACK); 2616280d235SMichel Lespinasse __rb_rotate_set_parents(parent, sibling, root, 2626280d235SMichel Lespinasse RB_RED); 2639c079addSMichel Lespinasse augment_rotate(parent, sibling); 2646280d235SMichel Lespinasse sibling = tmp1; 2651da177e4SLinus Torvalds } 2666280d235SMichel Lespinasse tmp1 = sibling->rb_right; 2676280d235SMichel Lespinasse if (!tmp1 || rb_is_black(tmp1)) { 2686280d235SMichel Lespinasse tmp2 = sibling->rb_left; 2696280d235SMichel Lespinasse if (!tmp2 || rb_is_black(tmp2)) { 2706280d235SMichel Lespinasse /* 2716280d235SMichel Lespinasse * Case 2 - sibling color flip 2726280d235SMichel Lespinasse * (p could be either color here) 2736280d235SMichel Lespinasse * 2746280d235SMichel Lespinasse * (p) (p) 2756280d235SMichel Lespinasse * / \ / \ 2766280d235SMichel Lespinasse * N S --> N s 2776280d235SMichel Lespinasse * / \ / \ 2786280d235SMichel Lespinasse * Sl Sr Sl Sr 2796280d235SMichel Lespinasse * 28046b6135aSMichel Lespinasse * This leaves us violating 5) which 28146b6135aSMichel Lespinasse * can be fixed by flipping p to black 28246b6135aSMichel Lespinasse * if it was red, or by recursing at p. 28346b6135aSMichel Lespinasse * p is red when coming from Case 1. 2846280d235SMichel Lespinasse */ 2856280d235SMichel Lespinasse rb_set_parent_color(sibling, parent, 2866280d235SMichel Lespinasse RB_RED); 28746b6135aSMichel Lespinasse if (rb_is_red(parent)) 28846b6135aSMichel Lespinasse rb_set_black(parent); 28946b6135aSMichel Lespinasse else { 2901da177e4SLinus Torvalds node = parent; 29155a98102SDavid Woodhouse parent = rb_parent(node); 29246b6135aSMichel Lespinasse if (parent) 293e125d147SMichel Lespinasse continue; 2941da177e4SLinus Torvalds } 29546b6135aSMichel Lespinasse break; 29646b6135aSMichel Lespinasse } 2976280d235SMichel Lespinasse /* 2986280d235SMichel Lespinasse * Case 3 - right rotate at sibling 2996280d235SMichel Lespinasse * (p could be either color here) 3006280d235SMichel Lespinasse * 3016280d235SMichel Lespinasse * (p) (p) 3026280d235SMichel Lespinasse * / \ / \ 303ce093a04SJie Chen * N S --> N sl 3046280d235SMichel Lespinasse * / \ \ 305ce093a04SJie Chen * sl Sr S 306ce093a04SJie Chen * \ 307ce093a04SJie Chen * Sr 308ce093a04SJie Chen * 309ce093a04SJie Chen * Note: p might be red, and then both 310ce093a04SJie Chen * p and sl are red after rotation(which 311ce093a04SJie Chen * breaks property 4). This is fixed in 312ce093a04SJie Chen * Case 4 (in __rb_rotate_set_parents() 313ce093a04SJie Chen * which set sl the color of p 314ce093a04SJie Chen * and set p RB_BLACK) 315ce093a04SJie Chen * 316ce093a04SJie Chen * (p) (sl) 317ce093a04SJie Chen * / \ / \ 318ce093a04SJie Chen * N sl --> P S 319ce093a04SJie Chen * \ / \ 320ce093a04SJie Chen * S N Sr 3216280d235SMichel Lespinasse * \ 3226280d235SMichel Lespinasse * Sr 3236280d235SMichel Lespinasse */ 324d72da4a4SPeter Zijlstra tmp1 = tmp2->rb_right; 325d72da4a4SPeter Zijlstra WRITE_ONCE(sibling->rb_left, tmp1); 326d72da4a4SPeter Zijlstra WRITE_ONCE(tmp2->rb_right, sibling); 327d72da4a4SPeter Zijlstra WRITE_ONCE(parent->rb_right, tmp2); 3286280d235SMichel Lespinasse if (tmp1) 3296280d235SMichel Lespinasse rb_set_parent_color(tmp1, sibling, 3306280d235SMichel Lespinasse RB_BLACK); 3319c079addSMichel Lespinasse augment_rotate(sibling, tmp2); 3326280d235SMichel Lespinasse tmp1 = sibling; 3336280d235SMichel Lespinasse sibling = tmp2; 3341da177e4SLinus Torvalds } 3356280d235SMichel Lespinasse /* 3366280d235SMichel Lespinasse * Case 4 - left rotate at parent + color flips 3376280d235SMichel Lespinasse * (p and sl could be either color here. 3386280d235SMichel Lespinasse * After rotation, p becomes black, s acquires 3396280d235SMichel Lespinasse * p's color, and sl keeps its color) 3406280d235SMichel Lespinasse * 3416280d235SMichel Lespinasse * (p) (s) 3426280d235SMichel Lespinasse * / \ / \ 3436280d235SMichel Lespinasse * N S --> P Sr 3446280d235SMichel Lespinasse * / \ / \ 3456280d235SMichel Lespinasse * (sl) sr N (sl) 3466280d235SMichel Lespinasse */ 347d72da4a4SPeter Zijlstra tmp2 = sibling->rb_left; 348d72da4a4SPeter Zijlstra WRITE_ONCE(parent->rb_right, tmp2); 349d72da4a4SPeter Zijlstra WRITE_ONCE(sibling->rb_left, parent); 3506280d235SMichel Lespinasse rb_set_parent_color(tmp1, sibling, RB_BLACK); 3516280d235SMichel Lespinasse if (tmp2) 3526280d235SMichel Lespinasse rb_set_parent(tmp2, parent); 3536280d235SMichel Lespinasse __rb_rotate_set_parents(parent, sibling, root, 3546280d235SMichel Lespinasse RB_BLACK); 3559c079addSMichel Lespinasse augment_rotate(parent, sibling); 3561da177e4SLinus Torvalds break; 357d6ff1273SMichel Lespinasse } else { 3586280d235SMichel Lespinasse sibling = parent->rb_left; 3596280d235SMichel Lespinasse if (rb_is_red(sibling)) { 3606280d235SMichel Lespinasse /* Case 1 - right rotate at parent */ 361d72da4a4SPeter Zijlstra tmp1 = sibling->rb_right; 362d72da4a4SPeter Zijlstra WRITE_ONCE(parent->rb_left, tmp1); 363d72da4a4SPeter Zijlstra WRITE_ONCE(sibling->rb_right, parent); 3646280d235SMichel Lespinasse rb_set_parent_color(tmp1, parent, RB_BLACK); 3656280d235SMichel Lespinasse __rb_rotate_set_parents(parent, sibling, root, 3666280d235SMichel Lespinasse RB_RED); 3679c079addSMichel Lespinasse augment_rotate(parent, sibling); 3686280d235SMichel Lespinasse sibling = tmp1; 3691da177e4SLinus Torvalds } 3706280d235SMichel Lespinasse tmp1 = sibling->rb_left; 3716280d235SMichel Lespinasse if (!tmp1 || rb_is_black(tmp1)) { 3726280d235SMichel Lespinasse tmp2 = sibling->rb_right; 3736280d235SMichel Lespinasse if (!tmp2 || rb_is_black(tmp2)) { 3746280d235SMichel Lespinasse /* Case 2 - sibling color flip */ 3756280d235SMichel Lespinasse rb_set_parent_color(sibling, parent, 3766280d235SMichel Lespinasse RB_RED); 37746b6135aSMichel Lespinasse if (rb_is_red(parent)) 37846b6135aSMichel Lespinasse rb_set_black(parent); 37946b6135aSMichel Lespinasse else { 3801da177e4SLinus Torvalds node = parent; 38155a98102SDavid Woodhouse parent = rb_parent(node); 38246b6135aSMichel Lespinasse if (parent) 383e125d147SMichel Lespinasse continue; 3841da177e4SLinus Torvalds } 38546b6135aSMichel Lespinasse break; 38646b6135aSMichel Lespinasse } 387ce093a04SJie Chen /* Case 3 - left rotate at sibling */ 388d72da4a4SPeter Zijlstra tmp1 = tmp2->rb_left; 389d72da4a4SPeter Zijlstra WRITE_ONCE(sibling->rb_right, tmp1); 390d72da4a4SPeter Zijlstra WRITE_ONCE(tmp2->rb_left, sibling); 391d72da4a4SPeter Zijlstra WRITE_ONCE(parent->rb_left, tmp2); 3926280d235SMichel Lespinasse if (tmp1) 3936280d235SMichel Lespinasse rb_set_parent_color(tmp1, sibling, 3946280d235SMichel Lespinasse RB_BLACK); 3959c079addSMichel Lespinasse augment_rotate(sibling, tmp2); 3966280d235SMichel Lespinasse tmp1 = sibling; 3976280d235SMichel Lespinasse sibling = tmp2; 3981da177e4SLinus Torvalds } 399ce093a04SJie Chen /* Case 4 - right rotate at parent + color flips */ 400d72da4a4SPeter Zijlstra tmp2 = sibling->rb_right; 401d72da4a4SPeter Zijlstra WRITE_ONCE(parent->rb_left, tmp2); 402d72da4a4SPeter Zijlstra WRITE_ONCE(sibling->rb_right, parent); 4036280d235SMichel Lespinasse rb_set_parent_color(tmp1, sibling, RB_BLACK); 4046280d235SMichel Lespinasse if (tmp2) 4056280d235SMichel Lespinasse rb_set_parent(tmp2, parent); 4066280d235SMichel Lespinasse __rb_rotate_set_parents(parent, sibling, root, 4076280d235SMichel Lespinasse RB_BLACK); 4089c079addSMichel Lespinasse augment_rotate(parent, sibling); 4091da177e4SLinus Torvalds break; 4101da177e4SLinus Torvalds } 4111da177e4SLinus Torvalds } 4121da177e4SLinus Torvalds } 4133cb7a563SMichel Lespinasse 4143cb7a563SMichel Lespinasse /* Non-inline version for rb_erase_augmented() use */ 4153cb7a563SMichel Lespinasse void __rb_erase_color(struct rb_node *parent, struct rb_root *root, 4163cb7a563SMichel Lespinasse void (*augment_rotate)(struct rb_node *old, struct rb_node *new)) 4173cb7a563SMichel Lespinasse { 4183cb7a563SMichel Lespinasse ____rb_erase_color(parent, root, augment_rotate); 4193cb7a563SMichel Lespinasse } 4209c079addSMichel Lespinasse EXPORT_SYMBOL(__rb_erase_color); 42114b94af0SMichel Lespinasse 42214b94af0SMichel Lespinasse /* 42314b94af0SMichel Lespinasse * Non-augmented rbtree manipulation functions. 42414b94af0SMichel Lespinasse * 42514b94af0SMichel Lespinasse * We use dummy augmented callbacks here, and have the compiler optimize them 42614b94af0SMichel Lespinasse * out of the rb_insert_color() and rb_erase() function definitions. 42714b94af0SMichel Lespinasse */ 42814b94af0SMichel Lespinasse 42914b94af0SMichel Lespinasse static inline void dummy_propagate(struct rb_node *node, struct rb_node *stop) {} 43014b94af0SMichel Lespinasse static inline void dummy_copy(struct rb_node *old, struct rb_node *new) {} 43114b94af0SMichel Lespinasse static inline void dummy_rotate(struct rb_node *old, struct rb_node *new) {} 43214b94af0SMichel Lespinasse 43314b94af0SMichel Lespinasse static const struct rb_augment_callbacks dummy_callbacks = { 434f231aebfSKees Cook .propagate = dummy_propagate, 435f231aebfSKees Cook .copy = dummy_copy, 436f231aebfSKees Cook .rotate = dummy_rotate 43714b94af0SMichel Lespinasse }; 43814b94af0SMichel Lespinasse 43914b94af0SMichel Lespinasse void rb_insert_color(struct rb_node *node, struct rb_root *root) 44014b94af0SMichel Lespinasse { 441*cd9e61edSDavidlohr Bueso __rb_insert(node, root, false, NULL, dummy_rotate); 44214b94af0SMichel Lespinasse } 44314b94af0SMichel Lespinasse EXPORT_SYMBOL(rb_insert_color); 44414b94af0SMichel Lespinasse 44514b94af0SMichel Lespinasse void rb_erase(struct rb_node *node, struct rb_root *root) 44614b94af0SMichel Lespinasse { 4473cb7a563SMichel Lespinasse struct rb_node *rebalance; 448*cd9e61edSDavidlohr Bueso rebalance = __rb_erase_augmented(node, root, 449*cd9e61edSDavidlohr Bueso NULL, &dummy_callbacks); 4503cb7a563SMichel Lespinasse if (rebalance) 4513cb7a563SMichel Lespinasse ____rb_erase_color(rebalance, root, dummy_rotate); 4521da177e4SLinus Torvalds } 4531da177e4SLinus Torvalds EXPORT_SYMBOL(rb_erase); 4541da177e4SLinus Torvalds 455*cd9e61edSDavidlohr Bueso void rb_insert_color_cached(struct rb_node *node, 456*cd9e61edSDavidlohr Bueso struct rb_root_cached *root, bool leftmost) 457*cd9e61edSDavidlohr Bueso { 458*cd9e61edSDavidlohr Bueso __rb_insert(node, &root->rb_root, leftmost, 459*cd9e61edSDavidlohr Bueso &root->rb_leftmost, dummy_rotate); 460*cd9e61edSDavidlohr Bueso } 461*cd9e61edSDavidlohr Bueso EXPORT_SYMBOL(rb_insert_color_cached); 462*cd9e61edSDavidlohr Bueso 463*cd9e61edSDavidlohr Bueso void rb_erase_cached(struct rb_node *node, struct rb_root_cached *root) 464*cd9e61edSDavidlohr Bueso { 465*cd9e61edSDavidlohr Bueso struct rb_node *rebalance; 466*cd9e61edSDavidlohr Bueso rebalance = __rb_erase_augmented(node, &root->rb_root, 467*cd9e61edSDavidlohr Bueso &root->rb_leftmost, &dummy_callbacks); 468*cd9e61edSDavidlohr Bueso if (rebalance) 469*cd9e61edSDavidlohr Bueso ____rb_erase_color(rebalance, &root->rb_root, dummy_rotate); 470*cd9e61edSDavidlohr Bueso } 471*cd9e61edSDavidlohr Bueso EXPORT_SYMBOL(rb_erase_cached); 472*cd9e61edSDavidlohr Bueso 47314b94af0SMichel Lespinasse /* 47414b94af0SMichel Lespinasse * Augmented rbtree manipulation functions. 47514b94af0SMichel Lespinasse * 47614b94af0SMichel Lespinasse * This instantiates the same __always_inline functions as in the non-augmented 47714b94af0SMichel Lespinasse * case, but this time with user-defined callbacks. 47814b94af0SMichel Lespinasse */ 47914b94af0SMichel Lespinasse 48014b94af0SMichel Lespinasse void __rb_insert_augmented(struct rb_node *node, struct rb_root *root, 481*cd9e61edSDavidlohr Bueso bool newleft, struct rb_node **leftmost, 48214b94af0SMichel Lespinasse void (*augment_rotate)(struct rb_node *old, struct rb_node *new)) 48314b94af0SMichel Lespinasse { 484*cd9e61edSDavidlohr Bueso __rb_insert(node, root, newleft, leftmost, augment_rotate); 48514b94af0SMichel Lespinasse } 48614b94af0SMichel Lespinasse EXPORT_SYMBOL(__rb_insert_augmented); 48714b94af0SMichel Lespinasse 4881da177e4SLinus Torvalds /* 4891da177e4SLinus Torvalds * This function returns the first node (in sort order) of the tree. 4901da177e4SLinus Torvalds */ 491f4b477c4SArtem Bityutskiy struct rb_node *rb_first(const struct rb_root *root) 4921da177e4SLinus Torvalds { 4931da177e4SLinus Torvalds struct rb_node *n; 4941da177e4SLinus Torvalds 4951da177e4SLinus Torvalds n = root->rb_node; 4961da177e4SLinus Torvalds if (!n) 4971da177e4SLinus Torvalds return NULL; 4981da177e4SLinus Torvalds while (n->rb_left) 4991da177e4SLinus Torvalds n = n->rb_left; 5001da177e4SLinus Torvalds return n; 5011da177e4SLinus Torvalds } 5021da177e4SLinus Torvalds EXPORT_SYMBOL(rb_first); 5031da177e4SLinus Torvalds 504f4b477c4SArtem Bityutskiy struct rb_node *rb_last(const struct rb_root *root) 5051da177e4SLinus Torvalds { 5061da177e4SLinus Torvalds struct rb_node *n; 5071da177e4SLinus Torvalds 5081da177e4SLinus Torvalds n = root->rb_node; 5091da177e4SLinus Torvalds if (!n) 5101da177e4SLinus Torvalds return NULL; 5111da177e4SLinus Torvalds while (n->rb_right) 5121da177e4SLinus Torvalds n = n->rb_right; 5131da177e4SLinus Torvalds return n; 5141da177e4SLinus Torvalds } 5151da177e4SLinus Torvalds EXPORT_SYMBOL(rb_last); 5161da177e4SLinus Torvalds 517f4b477c4SArtem Bityutskiy struct rb_node *rb_next(const struct rb_node *node) 5181da177e4SLinus Torvalds { 51955a98102SDavid Woodhouse struct rb_node *parent; 52055a98102SDavid Woodhouse 5214c199a93SMichel Lespinasse if (RB_EMPTY_NODE(node)) 52210fd48f2SJens Axboe return NULL; 52310fd48f2SJens Axboe 5247ce6ff9eSMichel Lespinasse /* 5257ce6ff9eSMichel Lespinasse * If we have a right-hand child, go down and then left as far 5267ce6ff9eSMichel Lespinasse * as we can. 5277ce6ff9eSMichel Lespinasse */ 5281da177e4SLinus Torvalds if (node->rb_right) { 5291da177e4SLinus Torvalds node = node->rb_right; 5301da177e4SLinus Torvalds while (node->rb_left) 5311da177e4SLinus Torvalds node=node->rb_left; 532f4b477c4SArtem Bityutskiy return (struct rb_node *)node; 5331da177e4SLinus Torvalds } 5341da177e4SLinus Torvalds 5357ce6ff9eSMichel Lespinasse /* 5367ce6ff9eSMichel Lespinasse * No right-hand children. Everything down and left is smaller than us, 5377ce6ff9eSMichel Lespinasse * so any 'next' node must be in the general direction of our parent. 5387ce6ff9eSMichel Lespinasse * Go up the tree; any time the ancestor is a right-hand child of its 5397ce6ff9eSMichel Lespinasse * parent, keep going up. First time it's a left-hand child of its 5407ce6ff9eSMichel Lespinasse * parent, said parent is our 'next' node. 5417ce6ff9eSMichel Lespinasse */ 54255a98102SDavid Woodhouse while ((parent = rb_parent(node)) && node == parent->rb_right) 54355a98102SDavid Woodhouse node = parent; 5441da177e4SLinus Torvalds 54555a98102SDavid Woodhouse return parent; 5461da177e4SLinus Torvalds } 5471da177e4SLinus Torvalds EXPORT_SYMBOL(rb_next); 5481da177e4SLinus Torvalds 549f4b477c4SArtem Bityutskiy struct rb_node *rb_prev(const struct rb_node *node) 5501da177e4SLinus Torvalds { 55155a98102SDavid Woodhouse struct rb_node *parent; 55255a98102SDavid Woodhouse 5534c199a93SMichel Lespinasse if (RB_EMPTY_NODE(node)) 55410fd48f2SJens Axboe return NULL; 55510fd48f2SJens Axboe 5567ce6ff9eSMichel Lespinasse /* 5577ce6ff9eSMichel Lespinasse * If we have a left-hand child, go down and then right as far 5587ce6ff9eSMichel Lespinasse * as we can. 5597ce6ff9eSMichel Lespinasse */ 5601da177e4SLinus Torvalds if (node->rb_left) { 5611da177e4SLinus Torvalds node = node->rb_left; 5621da177e4SLinus Torvalds while (node->rb_right) 5631da177e4SLinus Torvalds node=node->rb_right; 564f4b477c4SArtem Bityutskiy return (struct rb_node *)node; 5651da177e4SLinus Torvalds } 5661da177e4SLinus Torvalds 5677ce6ff9eSMichel Lespinasse /* 5687ce6ff9eSMichel Lespinasse * No left-hand children. Go up till we find an ancestor which 5697ce6ff9eSMichel Lespinasse * is a right-hand child of its parent. 5707ce6ff9eSMichel Lespinasse */ 57155a98102SDavid Woodhouse while ((parent = rb_parent(node)) && node == parent->rb_left) 57255a98102SDavid Woodhouse node = parent; 5731da177e4SLinus Torvalds 57455a98102SDavid Woodhouse return parent; 5751da177e4SLinus Torvalds } 5761da177e4SLinus Torvalds EXPORT_SYMBOL(rb_prev); 5771da177e4SLinus Torvalds 5781da177e4SLinus Torvalds void rb_replace_node(struct rb_node *victim, struct rb_node *new, 5791da177e4SLinus Torvalds struct rb_root *root) 5801da177e4SLinus Torvalds { 58155a98102SDavid Woodhouse struct rb_node *parent = rb_parent(victim); 5821da177e4SLinus Torvalds 583c1adf200SDavid Howells /* Copy the pointers/colour from the victim to the replacement */ 584c1adf200SDavid Howells *new = *victim; 585c1adf200SDavid Howells 5861da177e4SLinus Torvalds /* Set the surrounding nodes to point to the replacement */ 587c1adf200SDavid Howells if (victim->rb_left) 588c1adf200SDavid Howells rb_set_parent(victim->rb_left, new); 589c1adf200SDavid Howells if (victim->rb_right) 590c1adf200SDavid Howells rb_set_parent(victim->rb_right, new); 5917abc704aSMichel Lespinasse __rb_change_child(victim, new, parent, root); 592c1adf200SDavid Howells } 593c1adf200SDavid Howells EXPORT_SYMBOL(rb_replace_node); 594c1adf200SDavid Howells 595c1adf200SDavid Howells void rb_replace_node_rcu(struct rb_node *victim, struct rb_node *new, 596c1adf200SDavid Howells struct rb_root *root) 597c1adf200SDavid Howells { 598c1adf200SDavid Howells struct rb_node *parent = rb_parent(victim); 599c1adf200SDavid Howells 600c1adf200SDavid Howells /* Copy the pointers/colour from the victim to the replacement */ 601c1adf200SDavid Howells *new = *victim; 602c1adf200SDavid Howells 603c1adf200SDavid Howells /* Set the surrounding nodes to point to the replacement */ 6041da177e4SLinus Torvalds if (victim->rb_left) 60555a98102SDavid Woodhouse rb_set_parent(victim->rb_left, new); 6061da177e4SLinus Torvalds if (victim->rb_right) 60755a98102SDavid Woodhouse rb_set_parent(victim->rb_right, new); 6081da177e4SLinus Torvalds 609c1adf200SDavid Howells /* Set the parent's pointer to the new node last after an RCU barrier 610c1adf200SDavid Howells * so that the pointers onwards are seen to be set correctly when doing 611c1adf200SDavid Howells * an RCU walk over the tree. 612c1adf200SDavid Howells */ 613c1adf200SDavid Howells __rb_change_child_rcu(victim, new, parent, root); 6141da177e4SLinus Torvalds } 615c1adf200SDavid Howells EXPORT_SYMBOL(rb_replace_node_rcu); 6169dee5c51SCody P Schafer 6179dee5c51SCody P Schafer static struct rb_node *rb_left_deepest_node(const struct rb_node *node) 6189dee5c51SCody P Schafer { 6199dee5c51SCody P Schafer for (;;) { 6209dee5c51SCody P Schafer if (node->rb_left) 6219dee5c51SCody P Schafer node = node->rb_left; 6229dee5c51SCody P Schafer else if (node->rb_right) 6239dee5c51SCody P Schafer node = node->rb_right; 6249dee5c51SCody P Schafer else 6259dee5c51SCody P Schafer return (struct rb_node *)node; 6269dee5c51SCody P Schafer } 6279dee5c51SCody P Schafer } 6289dee5c51SCody P Schafer 6299dee5c51SCody P Schafer struct rb_node *rb_next_postorder(const struct rb_node *node) 6309dee5c51SCody P Schafer { 6319dee5c51SCody P Schafer const struct rb_node *parent; 6329dee5c51SCody P Schafer if (!node) 6339dee5c51SCody P Schafer return NULL; 6349dee5c51SCody P Schafer parent = rb_parent(node); 6359dee5c51SCody P Schafer 6369dee5c51SCody P Schafer /* If we're sitting on node, we've already seen our children */ 6379dee5c51SCody P Schafer if (parent && node == parent->rb_left && parent->rb_right) { 6389dee5c51SCody P Schafer /* If we are the parent's left node, go to the parent's right 6399dee5c51SCody P Schafer * node then all the way down to the left */ 6409dee5c51SCody P Schafer return rb_left_deepest_node(parent->rb_right); 6419dee5c51SCody P Schafer } else 6429dee5c51SCody P Schafer /* Otherwise we are the parent's right node, and the parent 6439dee5c51SCody P Schafer * should be next */ 6449dee5c51SCody P Schafer return (struct rb_node *)parent; 6459dee5c51SCody P Schafer } 6469dee5c51SCody P Schafer EXPORT_SYMBOL(rb_next_postorder); 6479dee5c51SCody P Schafer 6489dee5c51SCody P Schafer struct rb_node *rb_first_postorder(const struct rb_root *root) 6499dee5c51SCody P Schafer { 6509dee5c51SCody P Schafer if (!root->rb_node) 6519dee5c51SCody P Schafer return NULL; 6529dee5c51SCody P Schafer 6539dee5c51SCody P Schafer return rb_left_deepest_node(root->rb_node); 6549dee5c51SCody P Schafer } 6559dee5c51SCody P Schafer EXPORT_SYMBOL(rb_first_postorder); 656