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, 9814b94af0SMichel Lespinasse void (*augment_rotate)(struct rb_node *old, struct rb_node *new)) 991da177e4SLinus Torvalds { 1005bc9188aSMichel Lespinasse struct rb_node *parent = rb_red_parent(node), *gparent, *tmp; 1011da177e4SLinus Torvalds 1026d58452dSMichel Lespinasse while (true) { 1036d58452dSMichel Lespinasse /* 1046d58452dSMichel Lespinasse * Loop invariant: node is red 1056d58452dSMichel Lespinasse * 1066d58452dSMichel Lespinasse * If there is a black parent, we are done. 1076d58452dSMichel Lespinasse * Otherwise, take some corrective action as we don't 1086d58452dSMichel Lespinasse * want a red root or two consecutive red nodes. 1096d58452dSMichel Lespinasse */ 1106d58452dSMichel Lespinasse if (!parent) { 1115bc9188aSMichel Lespinasse rb_set_parent_color(node, NULL, RB_BLACK); 1126d58452dSMichel Lespinasse break; 1136d58452dSMichel Lespinasse } else if (rb_is_black(parent)) 1146d58452dSMichel Lespinasse break; 1156d58452dSMichel Lespinasse 1165bc9188aSMichel Lespinasse gparent = rb_red_parent(parent); 1171da177e4SLinus Torvalds 1185bc9188aSMichel Lespinasse tmp = gparent->rb_right; 11959633abfSMichel Lespinasse if (parent != tmp) { /* parent == gparent->rb_left */ 1205bc9188aSMichel Lespinasse if (tmp && rb_is_red(tmp)) { 1215bc9188aSMichel Lespinasse /* 1225bc9188aSMichel Lespinasse * Case 1 - color flips 1235bc9188aSMichel Lespinasse * 1245bc9188aSMichel Lespinasse * G g 1255bc9188aSMichel Lespinasse * / \ / \ 1265bc9188aSMichel Lespinasse * p u --> P U 1275bc9188aSMichel Lespinasse * / / 1281b9c53e8SWei Yang * n n 1295bc9188aSMichel Lespinasse * 1305bc9188aSMichel Lespinasse * However, since g's parent might be red, and 1315bc9188aSMichel Lespinasse * 4) does not allow this, we need to recurse 1325bc9188aSMichel Lespinasse * at g. 1335bc9188aSMichel Lespinasse */ 1345bc9188aSMichel Lespinasse rb_set_parent_color(tmp, gparent, RB_BLACK); 1355bc9188aSMichel Lespinasse rb_set_parent_color(parent, gparent, RB_BLACK); 1361da177e4SLinus Torvalds node = gparent; 1375bc9188aSMichel Lespinasse parent = rb_parent(node); 1385bc9188aSMichel Lespinasse rb_set_parent_color(node, parent, RB_RED); 1391da177e4SLinus Torvalds continue; 1401da177e4SLinus Torvalds } 1411da177e4SLinus Torvalds 14259633abfSMichel Lespinasse tmp = parent->rb_right; 14359633abfSMichel Lespinasse if (node == tmp) { 1445bc9188aSMichel Lespinasse /* 1455bc9188aSMichel Lespinasse * Case 2 - left rotate at parent 1465bc9188aSMichel Lespinasse * 1475bc9188aSMichel Lespinasse * G G 1485bc9188aSMichel Lespinasse * / \ / \ 1495bc9188aSMichel Lespinasse * p U --> n U 1505bc9188aSMichel Lespinasse * \ / 1515bc9188aSMichel Lespinasse * n p 1525bc9188aSMichel Lespinasse * 1535bc9188aSMichel Lespinasse * This still leaves us in violation of 4), the 1545bc9188aSMichel Lespinasse * continuation into Case 3 will fix that. 1555bc9188aSMichel Lespinasse */ 156d72da4a4SPeter Zijlstra tmp = node->rb_left; 157d72da4a4SPeter Zijlstra WRITE_ONCE(parent->rb_right, tmp); 158d72da4a4SPeter Zijlstra WRITE_ONCE(node->rb_left, parent); 1595bc9188aSMichel Lespinasse if (tmp) 1605bc9188aSMichel Lespinasse rb_set_parent_color(tmp, parent, 1615bc9188aSMichel Lespinasse RB_BLACK); 1625bc9188aSMichel Lespinasse rb_set_parent_color(parent, node, RB_RED); 16314b94af0SMichel Lespinasse augment_rotate(parent, node); 1641da177e4SLinus Torvalds parent = node; 16559633abfSMichel Lespinasse tmp = node->rb_right; 1661da177e4SLinus Torvalds } 1671da177e4SLinus Torvalds 1685bc9188aSMichel Lespinasse /* 1695bc9188aSMichel Lespinasse * Case 3 - right rotate at gparent 1705bc9188aSMichel Lespinasse * 1715bc9188aSMichel Lespinasse * G P 1725bc9188aSMichel Lespinasse * / \ / \ 1735bc9188aSMichel Lespinasse * p U --> n g 1745bc9188aSMichel Lespinasse * / \ 1755bc9188aSMichel Lespinasse * n U 1765bc9188aSMichel Lespinasse */ 177d72da4a4SPeter Zijlstra WRITE_ONCE(gparent->rb_left, tmp); /* == parent->rb_right */ 178d72da4a4SPeter Zijlstra WRITE_ONCE(parent->rb_right, gparent); 1795bc9188aSMichel Lespinasse if (tmp) 1805bc9188aSMichel Lespinasse rb_set_parent_color(tmp, gparent, RB_BLACK); 1815bc9188aSMichel Lespinasse __rb_rotate_set_parents(gparent, parent, root, RB_RED); 18214b94af0SMichel Lespinasse augment_rotate(gparent, parent); 1831f052865SMichel Lespinasse break; 1841da177e4SLinus Torvalds } else { 1855bc9188aSMichel Lespinasse tmp = gparent->rb_left; 1865bc9188aSMichel Lespinasse if (tmp && rb_is_red(tmp)) { 1875bc9188aSMichel Lespinasse /* Case 1 - color flips */ 1885bc9188aSMichel Lespinasse rb_set_parent_color(tmp, gparent, RB_BLACK); 1895bc9188aSMichel Lespinasse rb_set_parent_color(parent, gparent, RB_BLACK); 1901da177e4SLinus Torvalds node = gparent; 1915bc9188aSMichel Lespinasse parent = rb_parent(node); 1925bc9188aSMichel Lespinasse rb_set_parent_color(node, parent, RB_RED); 1931da177e4SLinus Torvalds continue; 1941da177e4SLinus Torvalds } 1951da177e4SLinus Torvalds 19659633abfSMichel Lespinasse tmp = parent->rb_left; 19759633abfSMichel Lespinasse if (node == tmp) { 1985bc9188aSMichel Lespinasse /* Case 2 - right rotate at parent */ 199d72da4a4SPeter Zijlstra tmp = node->rb_right; 200d72da4a4SPeter Zijlstra WRITE_ONCE(parent->rb_left, tmp); 201d72da4a4SPeter Zijlstra WRITE_ONCE(node->rb_right, parent); 2025bc9188aSMichel Lespinasse if (tmp) 2035bc9188aSMichel Lespinasse rb_set_parent_color(tmp, parent, 2045bc9188aSMichel Lespinasse RB_BLACK); 2055bc9188aSMichel Lespinasse rb_set_parent_color(parent, node, RB_RED); 20614b94af0SMichel Lespinasse augment_rotate(parent, node); 2071da177e4SLinus Torvalds parent = node; 20859633abfSMichel Lespinasse tmp = node->rb_left; 2091da177e4SLinus Torvalds } 2101da177e4SLinus Torvalds 2115bc9188aSMichel Lespinasse /* Case 3 - left rotate at gparent */ 212d72da4a4SPeter Zijlstra WRITE_ONCE(gparent->rb_right, tmp); /* == parent->rb_left */ 213d72da4a4SPeter Zijlstra WRITE_ONCE(parent->rb_left, gparent); 2145bc9188aSMichel Lespinasse if (tmp) 2155bc9188aSMichel Lespinasse rb_set_parent_color(tmp, gparent, RB_BLACK); 2165bc9188aSMichel Lespinasse __rb_rotate_set_parents(gparent, parent, root, RB_RED); 21714b94af0SMichel Lespinasse augment_rotate(gparent, parent); 2181f052865SMichel Lespinasse break; 2191da177e4SLinus Torvalds } 2201da177e4SLinus Torvalds } 2211da177e4SLinus Torvalds } 2221da177e4SLinus Torvalds 2233cb7a563SMichel Lespinasse /* 2243cb7a563SMichel Lespinasse * Inline version for rb_erase() use - we want to be able to inline 2253cb7a563SMichel Lespinasse * and eliminate the dummy_rotate callback there 2263cb7a563SMichel Lespinasse */ 2273cb7a563SMichel Lespinasse static __always_inline void 2283cb7a563SMichel Lespinasse ____rb_erase_color(struct rb_node *parent, struct rb_root *root, 2299c079addSMichel Lespinasse void (*augment_rotate)(struct rb_node *old, struct rb_node *new)) 2301da177e4SLinus Torvalds { 23146b6135aSMichel Lespinasse struct rb_node *node = NULL, *sibling, *tmp1, *tmp2; 2321da177e4SLinus Torvalds 233d6ff1273SMichel Lespinasse while (true) { 234d6ff1273SMichel Lespinasse /* 23546b6135aSMichel Lespinasse * Loop invariants: 23646b6135aSMichel Lespinasse * - node is black (or NULL on first iteration) 23746b6135aSMichel Lespinasse * - node is not the root (parent is not NULL) 23846b6135aSMichel Lespinasse * - All leaf paths going through parent and node have a 239d6ff1273SMichel Lespinasse * black node count that is 1 lower than other leaf paths. 240d6ff1273SMichel Lespinasse */ 2416280d235SMichel Lespinasse sibling = parent->rb_right; 24259633abfSMichel Lespinasse if (node != sibling) { /* node == parent->rb_left */ 2436280d235SMichel Lespinasse if (rb_is_red(sibling)) { 2446280d235SMichel Lespinasse /* 2456280d235SMichel Lespinasse * Case 1 - left rotate at parent 2466280d235SMichel Lespinasse * 2476280d235SMichel Lespinasse * P S 2486280d235SMichel Lespinasse * / \ / \ 2496280d235SMichel Lespinasse * N s --> p Sr 2506280d235SMichel Lespinasse * / \ / \ 2516280d235SMichel Lespinasse * Sl Sr N Sl 2526280d235SMichel Lespinasse */ 253d72da4a4SPeter Zijlstra tmp1 = sibling->rb_left; 254d72da4a4SPeter Zijlstra WRITE_ONCE(parent->rb_right, tmp1); 255d72da4a4SPeter Zijlstra WRITE_ONCE(sibling->rb_left, parent); 2566280d235SMichel Lespinasse rb_set_parent_color(tmp1, parent, RB_BLACK); 2576280d235SMichel Lespinasse __rb_rotate_set_parents(parent, sibling, root, 2586280d235SMichel Lespinasse RB_RED); 2599c079addSMichel Lespinasse augment_rotate(parent, sibling); 2606280d235SMichel Lespinasse sibling = tmp1; 2611da177e4SLinus Torvalds } 2626280d235SMichel Lespinasse tmp1 = sibling->rb_right; 2636280d235SMichel Lespinasse if (!tmp1 || rb_is_black(tmp1)) { 2646280d235SMichel Lespinasse tmp2 = sibling->rb_left; 2656280d235SMichel Lespinasse if (!tmp2 || rb_is_black(tmp2)) { 2666280d235SMichel Lespinasse /* 2676280d235SMichel Lespinasse * Case 2 - sibling color flip 2686280d235SMichel Lespinasse * (p could be either color here) 2696280d235SMichel Lespinasse * 2706280d235SMichel Lespinasse * (p) (p) 2716280d235SMichel Lespinasse * / \ / \ 2726280d235SMichel Lespinasse * N S --> N s 2736280d235SMichel Lespinasse * / \ / \ 2746280d235SMichel Lespinasse * Sl Sr Sl Sr 2756280d235SMichel Lespinasse * 27646b6135aSMichel Lespinasse * This leaves us violating 5) which 27746b6135aSMichel Lespinasse * can be fixed by flipping p to black 27846b6135aSMichel Lespinasse * if it was red, or by recursing at p. 27946b6135aSMichel Lespinasse * p is red when coming from Case 1. 2806280d235SMichel Lespinasse */ 2816280d235SMichel Lespinasse rb_set_parent_color(sibling, parent, 2826280d235SMichel Lespinasse RB_RED); 28346b6135aSMichel Lespinasse if (rb_is_red(parent)) 28446b6135aSMichel Lespinasse rb_set_black(parent); 28546b6135aSMichel Lespinasse else { 2861da177e4SLinus Torvalds node = parent; 28755a98102SDavid Woodhouse parent = rb_parent(node); 28846b6135aSMichel Lespinasse if (parent) 289e125d147SMichel Lespinasse continue; 2901da177e4SLinus Torvalds } 29146b6135aSMichel Lespinasse break; 29246b6135aSMichel Lespinasse } 2936280d235SMichel Lespinasse /* 2946280d235SMichel Lespinasse * Case 3 - right rotate at sibling 2956280d235SMichel Lespinasse * (p could be either color here) 2966280d235SMichel Lespinasse * 2976280d235SMichel Lespinasse * (p) (p) 2986280d235SMichel Lespinasse * / \ / \ 299ce093a04SJie Chen * N S --> N sl 3006280d235SMichel Lespinasse * / \ \ 301ce093a04SJie Chen * sl Sr S 302ce093a04SJie Chen * \ 303ce093a04SJie Chen * Sr 304ce093a04SJie Chen * 305ce093a04SJie Chen * Note: p might be red, and then both 306ce093a04SJie Chen * p and sl are red after rotation(which 307ce093a04SJie Chen * breaks property 4). This is fixed in 308ce093a04SJie Chen * Case 4 (in __rb_rotate_set_parents() 309ce093a04SJie Chen * which set sl the color of p 310ce093a04SJie Chen * and set p RB_BLACK) 311ce093a04SJie Chen * 312ce093a04SJie Chen * (p) (sl) 313ce093a04SJie Chen * / \ / \ 314ce093a04SJie Chen * N sl --> P S 315ce093a04SJie Chen * \ / \ 316ce093a04SJie Chen * S N Sr 3176280d235SMichel Lespinasse * \ 3186280d235SMichel Lespinasse * Sr 3196280d235SMichel Lespinasse */ 320d72da4a4SPeter Zijlstra tmp1 = tmp2->rb_right; 321d72da4a4SPeter Zijlstra WRITE_ONCE(sibling->rb_left, tmp1); 322d72da4a4SPeter Zijlstra WRITE_ONCE(tmp2->rb_right, sibling); 323d72da4a4SPeter Zijlstra WRITE_ONCE(parent->rb_right, tmp2); 3246280d235SMichel Lespinasse if (tmp1) 3256280d235SMichel Lespinasse rb_set_parent_color(tmp1, sibling, 3266280d235SMichel Lespinasse RB_BLACK); 3279c079addSMichel Lespinasse augment_rotate(sibling, tmp2); 3286280d235SMichel Lespinasse tmp1 = sibling; 3296280d235SMichel Lespinasse sibling = tmp2; 3301da177e4SLinus Torvalds } 3316280d235SMichel Lespinasse /* 3326280d235SMichel Lespinasse * Case 4 - left rotate at parent + color flips 3336280d235SMichel Lespinasse * (p and sl could be either color here. 3346280d235SMichel Lespinasse * After rotation, p becomes black, s acquires 3356280d235SMichel Lespinasse * p's color, and sl keeps its color) 3366280d235SMichel Lespinasse * 3376280d235SMichel Lespinasse * (p) (s) 3386280d235SMichel Lespinasse * / \ / \ 3396280d235SMichel Lespinasse * N S --> P Sr 3406280d235SMichel Lespinasse * / \ / \ 3416280d235SMichel Lespinasse * (sl) sr N (sl) 3426280d235SMichel Lespinasse */ 343d72da4a4SPeter Zijlstra tmp2 = sibling->rb_left; 344d72da4a4SPeter Zijlstra WRITE_ONCE(parent->rb_right, tmp2); 345d72da4a4SPeter Zijlstra WRITE_ONCE(sibling->rb_left, parent); 3466280d235SMichel Lespinasse rb_set_parent_color(tmp1, sibling, RB_BLACK); 3476280d235SMichel Lespinasse if (tmp2) 3486280d235SMichel Lespinasse rb_set_parent(tmp2, parent); 3496280d235SMichel Lespinasse __rb_rotate_set_parents(parent, sibling, root, 3506280d235SMichel Lespinasse RB_BLACK); 3519c079addSMichel Lespinasse augment_rotate(parent, sibling); 3521da177e4SLinus Torvalds break; 353d6ff1273SMichel Lespinasse } else { 3546280d235SMichel Lespinasse sibling = parent->rb_left; 3556280d235SMichel Lespinasse if (rb_is_red(sibling)) { 3566280d235SMichel Lespinasse /* Case 1 - right rotate at parent */ 357d72da4a4SPeter Zijlstra tmp1 = sibling->rb_right; 358d72da4a4SPeter Zijlstra WRITE_ONCE(parent->rb_left, tmp1); 359d72da4a4SPeter Zijlstra WRITE_ONCE(sibling->rb_right, parent); 3606280d235SMichel Lespinasse rb_set_parent_color(tmp1, parent, RB_BLACK); 3616280d235SMichel Lespinasse __rb_rotate_set_parents(parent, sibling, root, 3626280d235SMichel Lespinasse RB_RED); 3639c079addSMichel Lespinasse augment_rotate(parent, sibling); 3646280d235SMichel Lespinasse sibling = tmp1; 3651da177e4SLinus Torvalds } 3666280d235SMichel Lespinasse tmp1 = sibling->rb_left; 3676280d235SMichel Lespinasse if (!tmp1 || rb_is_black(tmp1)) { 3686280d235SMichel Lespinasse tmp2 = sibling->rb_right; 3696280d235SMichel Lespinasse if (!tmp2 || rb_is_black(tmp2)) { 3706280d235SMichel Lespinasse /* Case 2 - sibling color flip */ 3716280d235SMichel Lespinasse rb_set_parent_color(sibling, parent, 3726280d235SMichel Lespinasse RB_RED); 37346b6135aSMichel Lespinasse if (rb_is_red(parent)) 37446b6135aSMichel Lespinasse rb_set_black(parent); 37546b6135aSMichel Lespinasse else { 3761da177e4SLinus Torvalds node = parent; 37755a98102SDavid Woodhouse parent = rb_parent(node); 37846b6135aSMichel Lespinasse if (parent) 379e125d147SMichel Lespinasse continue; 3801da177e4SLinus Torvalds } 38146b6135aSMichel Lespinasse break; 38246b6135aSMichel Lespinasse } 383ce093a04SJie Chen /* Case 3 - left rotate at sibling */ 384d72da4a4SPeter Zijlstra tmp1 = tmp2->rb_left; 385d72da4a4SPeter Zijlstra WRITE_ONCE(sibling->rb_right, tmp1); 386d72da4a4SPeter Zijlstra WRITE_ONCE(tmp2->rb_left, sibling); 387d72da4a4SPeter Zijlstra WRITE_ONCE(parent->rb_left, tmp2); 3886280d235SMichel Lespinasse if (tmp1) 3896280d235SMichel Lespinasse rb_set_parent_color(tmp1, sibling, 3906280d235SMichel Lespinasse RB_BLACK); 3919c079addSMichel Lespinasse augment_rotate(sibling, tmp2); 3926280d235SMichel Lespinasse tmp1 = sibling; 3936280d235SMichel Lespinasse sibling = tmp2; 3941da177e4SLinus Torvalds } 395ce093a04SJie Chen /* Case 4 - right rotate at parent + color flips */ 396d72da4a4SPeter Zijlstra tmp2 = sibling->rb_right; 397d72da4a4SPeter Zijlstra WRITE_ONCE(parent->rb_left, tmp2); 398d72da4a4SPeter Zijlstra WRITE_ONCE(sibling->rb_right, parent); 3996280d235SMichel Lespinasse rb_set_parent_color(tmp1, sibling, RB_BLACK); 4006280d235SMichel Lespinasse if (tmp2) 4016280d235SMichel Lespinasse rb_set_parent(tmp2, parent); 4026280d235SMichel Lespinasse __rb_rotate_set_parents(parent, sibling, root, 4036280d235SMichel Lespinasse RB_BLACK); 4049c079addSMichel Lespinasse augment_rotate(parent, sibling); 4051da177e4SLinus Torvalds break; 4061da177e4SLinus Torvalds } 4071da177e4SLinus Torvalds } 4081da177e4SLinus Torvalds } 4093cb7a563SMichel Lespinasse 4103cb7a563SMichel Lespinasse /* Non-inline version for rb_erase_augmented() use */ 4113cb7a563SMichel Lespinasse void __rb_erase_color(struct rb_node *parent, struct rb_root *root, 4123cb7a563SMichel Lespinasse void (*augment_rotate)(struct rb_node *old, struct rb_node *new)) 4133cb7a563SMichel Lespinasse { 4143cb7a563SMichel Lespinasse ____rb_erase_color(parent, root, augment_rotate); 4153cb7a563SMichel Lespinasse } 4169c079addSMichel Lespinasse EXPORT_SYMBOL(__rb_erase_color); 41714b94af0SMichel Lespinasse 41814b94af0SMichel Lespinasse /* 41914b94af0SMichel Lespinasse * Non-augmented rbtree manipulation functions. 42014b94af0SMichel Lespinasse * 42114b94af0SMichel Lespinasse * We use dummy augmented callbacks here, and have the compiler optimize them 42214b94af0SMichel Lespinasse * out of the rb_insert_color() and rb_erase() function definitions. 42314b94af0SMichel Lespinasse */ 42414b94af0SMichel Lespinasse 42514b94af0SMichel Lespinasse static inline void dummy_propagate(struct rb_node *node, struct rb_node *stop) {} 42614b94af0SMichel Lespinasse static inline void dummy_copy(struct rb_node *old, struct rb_node *new) {} 42714b94af0SMichel Lespinasse static inline void dummy_rotate(struct rb_node *old, struct rb_node *new) {} 42814b94af0SMichel Lespinasse 42914b94af0SMichel Lespinasse static const struct rb_augment_callbacks dummy_callbacks = { 430*f231aebfSKees Cook .propagate = dummy_propagate, 431*f231aebfSKees Cook .copy = dummy_copy, 432*f231aebfSKees Cook .rotate = dummy_rotate 43314b94af0SMichel Lespinasse }; 43414b94af0SMichel Lespinasse 43514b94af0SMichel Lespinasse void rb_insert_color(struct rb_node *node, struct rb_root *root) 43614b94af0SMichel Lespinasse { 43714b94af0SMichel Lespinasse __rb_insert(node, root, dummy_rotate); 43814b94af0SMichel Lespinasse } 43914b94af0SMichel Lespinasse EXPORT_SYMBOL(rb_insert_color); 44014b94af0SMichel Lespinasse 44114b94af0SMichel Lespinasse void rb_erase(struct rb_node *node, struct rb_root *root) 44214b94af0SMichel Lespinasse { 4433cb7a563SMichel Lespinasse struct rb_node *rebalance; 4443cb7a563SMichel Lespinasse rebalance = __rb_erase_augmented(node, root, &dummy_callbacks); 4453cb7a563SMichel Lespinasse if (rebalance) 4463cb7a563SMichel Lespinasse ____rb_erase_color(rebalance, root, dummy_rotate); 4471da177e4SLinus Torvalds } 4481da177e4SLinus Torvalds EXPORT_SYMBOL(rb_erase); 4491da177e4SLinus Torvalds 45014b94af0SMichel Lespinasse /* 45114b94af0SMichel Lespinasse * Augmented rbtree manipulation functions. 45214b94af0SMichel Lespinasse * 45314b94af0SMichel Lespinasse * This instantiates the same __always_inline functions as in the non-augmented 45414b94af0SMichel Lespinasse * case, but this time with user-defined callbacks. 45514b94af0SMichel Lespinasse */ 45614b94af0SMichel Lespinasse 45714b94af0SMichel Lespinasse void __rb_insert_augmented(struct rb_node *node, struct rb_root *root, 45814b94af0SMichel Lespinasse void (*augment_rotate)(struct rb_node *old, struct rb_node *new)) 45914b94af0SMichel Lespinasse { 46014b94af0SMichel Lespinasse __rb_insert(node, root, augment_rotate); 46114b94af0SMichel Lespinasse } 46214b94af0SMichel Lespinasse EXPORT_SYMBOL(__rb_insert_augmented); 46314b94af0SMichel Lespinasse 4641da177e4SLinus Torvalds /* 4651da177e4SLinus Torvalds * This function returns the first node (in sort order) of the tree. 4661da177e4SLinus Torvalds */ 467f4b477c4SArtem Bityutskiy struct rb_node *rb_first(const struct rb_root *root) 4681da177e4SLinus Torvalds { 4691da177e4SLinus Torvalds struct rb_node *n; 4701da177e4SLinus Torvalds 4711da177e4SLinus Torvalds n = root->rb_node; 4721da177e4SLinus Torvalds if (!n) 4731da177e4SLinus Torvalds return NULL; 4741da177e4SLinus Torvalds while (n->rb_left) 4751da177e4SLinus Torvalds n = n->rb_left; 4761da177e4SLinus Torvalds return n; 4771da177e4SLinus Torvalds } 4781da177e4SLinus Torvalds EXPORT_SYMBOL(rb_first); 4791da177e4SLinus Torvalds 480f4b477c4SArtem Bityutskiy struct rb_node *rb_last(const struct rb_root *root) 4811da177e4SLinus Torvalds { 4821da177e4SLinus Torvalds struct rb_node *n; 4831da177e4SLinus Torvalds 4841da177e4SLinus Torvalds n = root->rb_node; 4851da177e4SLinus Torvalds if (!n) 4861da177e4SLinus Torvalds return NULL; 4871da177e4SLinus Torvalds while (n->rb_right) 4881da177e4SLinus Torvalds n = n->rb_right; 4891da177e4SLinus Torvalds return n; 4901da177e4SLinus Torvalds } 4911da177e4SLinus Torvalds EXPORT_SYMBOL(rb_last); 4921da177e4SLinus Torvalds 493f4b477c4SArtem Bityutskiy struct rb_node *rb_next(const struct rb_node *node) 4941da177e4SLinus Torvalds { 49555a98102SDavid Woodhouse struct rb_node *parent; 49655a98102SDavid Woodhouse 4974c199a93SMichel Lespinasse if (RB_EMPTY_NODE(node)) 49810fd48f2SJens Axboe return NULL; 49910fd48f2SJens Axboe 5007ce6ff9eSMichel Lespinasse /* 5017ce6ff9eSMichel Lespinasse * If we have a right-hand child, go down and then left as far 5027ce6ff9eSMichel Lespinasse * as we can. 5037ce6ff9eSMichel Lespinasse */ 5041da177e4SLinus Torvalds if (node->rb_right) { 5051da177e4SLinus Torvalds node = node->rb_right; 5061da177e4SLinus Torvalds while (node->rb_left) 5071da177e4SLinus Torvalds node=node->rb_left; 508f4b477c4SArtem Bityutskiy return (struct rb_node *)node; 5091da177e4SLinus Torvalds } 5101da177e4SLinus Torvalds 5117ce6ff9eSMichel Lespinasse /* 5127ce6ff9eSMichel Lespinasse * No right-hand children. Everything down and left is smaller than us, 5137ce6ff9eSMichel Lespinasse * so any 'next' node must be in the general direction of our parent. 5147ce6ff9eSMichel Lespinasse * Go up the tree; any time the ancestor is a right-hand child of its 5157ce6ff9eSMichel Lespinasse * parent, keep going up. First time it's a left-hand child of its 5167ce6ff9eSMichel Lespinasse * parent, said parent is our 'next' node. 5177ce6ff9eSMichel Lespinasse */ 51855a98102SDavid Woodhouse while ((parent = rb_parent(node)) && node == parent->rb_right) 51955a98102SDavid Woodhouse node = parent; 5201da177e4SLinus Torvalds 52155a98102SDavid Woodhouse return parent; 5221da177e4SLinus Torvalds } 5231da177e4SLinus Torvalds EXPORT_SYMBOL(rb_next); 5241da177e4SLinus Torvalds 525f4b477c4SArtem Bityutskiy struct rb_node *rb_prev(const struct rb_node *node) 5261da177e4SLinus Torvalds { 52755a98102SDavid Woodhouse struct rb_node *parent; 52855a98102SDavid Woodhouse 5294c199a93SMichel Lespinasse if (RB_EMPTY_NODE(node)) 53010fd48f2SJens Axboe return NULL; 53110fd48f2SJens Axboe 5327ce6ff9eSMichel Lespinasse /* 5337ce6ff9eSMichel Lespinasse * If we have a left-hand child, go down and then right as far 5347ce6ff9eSMichel Lespinasse * as we can. 5357ce6ff9eSMichel Lespinasse */ 5361da177e4SLinus Torvalds if (node->rb_left) { 5371da177e4SLinus Torvalds node = node->rb_left; 5381da177e4SLinus Torvalds while (node->rb_right) 5391da177e4SLinus Torvalds node=node->rb_right; 540f4b477c4SArtem Bityutskiy return (struct rb_node *)node; 5411da177e4SLinus Torvalds } 5421da177e4SLinus Torvalds 5437ce6ff9eSMichel Lespinasse /* 5447ce6ff9eSMichel Lespinasse * No left-hand children. Go up till we find an ancestor which 5457ce6ff9eSMichel Lespinasse * is a right-hand child of its parent. 5467ce6ff9eSMichel Lespinasse */ 54755a98102SDavid Woodhouse while ((parent = rb_parent(node)) && node == parent->rb_left) 54855a98102SDavid Woodhouse node = parent; 5491da177e4SLinus Torvalds 55055a98102SDavid Woodhouse return parent; 5511da177e4SLinus Torvalds } 5521da177e4SLinus Torvalds EXPORT_SYMBOL(rb_prev); 5531da177e4SLinus Torvalds 5541da177e4SLinus Torvalds void rb_replace_node(struct rb_node *victim, struct rb_node *new, 5551da177e4SLinus Torvalds struct rb_root *root) 5561da177e4SLinus Torvalds { 55755a98102SDavid Woodhouse struct rb_node *parent = rb_parent(victim); 5581da177e4SLinus Torvalds 559c1adf200SDavid Howells /* Copy the pointers/colour from the victim to the replacement */ 560c1adf200SDavid Howells *new = *victim; 561c1adf200SDavid Howells 5621da177e4SLinus Torvalds /* Set the surrounding nodes to point to the replacement */ 563c1adf200SDavid Howells if (victim->rb_left) 564c1adf200SDavid Howells rb_set_parent(victim->rb_left, new); 565c1adf200SDavid Howells if (victim->rb_right) 566c1adf200SDavid Howells rb_set_parent(victim->rb_right, new); 5677abc704aSMichel Lespinasse __rb_change_child(victim, new, parent, root); 568c1adf200SDavid Howells } 569c1adf200SDavid Howells EXPORT_SYMBOL(rb_replace_node); 570c1adf200SDavid Howells 571c1adf200SDavid Howells void rb_replace_node_rcu(struct rb_node *victim, struct rb_node *new, 572c1adf200SDavid Howells struct rb_root *root) 573c1adf200SDavid Howells { 574c1adf200SDavid Howells struct rb_node *parent = rb_parent(victim); 575c1adf200SDavid Howells 576c1adf200SDavid Howells /* Copy the pointers/colour from the victim to the replacement */ 577c1adf200SDavid Howells *new = *victim; 578c1adf200SDavid Howells 579c1adf200SDavid Howells /* Set the surrounding nodes to point to the replacement */ 5801da177e4SLinus Torvalds if (victim->rb_left) 58155a98102SDavid Woodhouse rb_set_parent(victim->rb_left, new); 5821da177e4SLinus Torvalds if (victim->rb_right) 58355a98102SDavid Woodhouse rb_set_parent(victim->rb_right, new); 5841da177e4SLinus Torvalds 585c1adf200SDavid Howells /* Set the parent's pointer to the new node last after an RCU barrier 586c1adf200SDavid Howells * so that the pointers onwards are seen to be set correctly when doing 587c1adf200SDavid Howells * an RCU walk over the tree. 588c1adf200SDavid Howells */ 589c1adf200SDavid Howells __rb_change_child_rcu(victim, new, parent, root); 5901da177e4SLinus Torvalds } 591c1adf200SDavid Howells EXPORT_SYMBOL(rb_replace_node_rcu); 5929dee5c51SCody P Schafer 5939dee5c51SCody P Schafer static struct rb_node *rb_left_deepest_node(const struct rb_node *node) 5949dee5c51SCody P Schafer { 5959dee5c51SCody P Schafer for (;;) { 5969dee5c51SCody P Schafer if (node->rb_left) 5979dee5c51SCody P Schafer node = node->rb_left; 5989dee5c51SCody P Schafer else if (node->rb_right) 5999dee5c51SCody P Schafer node = node->rb_right; 6009dee5c51SCody P Schafer else 6019dee5c51SCody P Schafer return (struct rb_node *)node; 6029dee5c51SCody P Schafer } 6039dee5c51SCody P Schafer } 6049dee5c51SCody P Schafer 6059dee5c51SCody P Schafer struct rb_node *rb_next_postorder(const struct rb_node *node) 6069dee5c51SCody P Schafer { 6079dee5c51SCody P Schafer const struct rb_node *parent; 6089dee5c51SCody P Schafer if (!node) 6099dee5c51SCody P Schafer return NULL; 6109dee5c51SCody P Schafer parent = rb_parent(node); 6119dee5c51SCody P Schafer 6129dee5c51SCody P Schafer /* If we're sitting on node, we've already seen our children */ 6139dee5c51SCody P Schafer if (parent && node == parent->rb_left && parent->rb_right) { 6149dee5c51SCody P Schafer /* If we are the parent's left node, go to the parent's right 6159dee5c51SCody P Schafer * node then all the way down to the left */ 6169dee5c51SCody P Schafer return rb_left_deepest_node(parent->rb_right); 6179dee5c51SCody P Schafer } else 6189dee5c51SCody P Schafer /* Otherwise we are the parent's right node, and the parent 6199dee5c51SCody P Schafer * should be next */ 6209dee5c51SCody P Schafer return (struct rb_node *)parent; 6219dee5c51SCody P Schafer } 6229dee5c51SCody P Schafer EXPORT_SYMBOL(rb_next_postorder); 6239dee5c51SCody P Schafer 6249dee5c51SCody P Schafer struct rb_node *rb_first_postorder(const struct rb_root *root) 6259dee5c51SCody P Schafer { 6269dee5c51SCody P Schafer if (!root->rb_node) 6279dee5c51SCody P Schafer return NULL; 6289dee5c51SCody P Schafer 6299dee5c51SCody P Schafer return rb_left_deepest_node(root->rb_node); 6309dee5c51SCody P Schafer } 6319dee5c51SCody P Schafer EXPORT_SYMBOL(rb_first_postorder); 632