11a59d1b8SThomas Gleixner // SPDX-License-Identifier: GPL-2.0-or-later 21da177e4SLinus Torvalds /* 31da177e4SLinus Torvalds Red Black Trees 41da177e4SLinus Torvalds (C) 1999 Andrea Arcangeli <andrea@suse.de> 51da177e4SLinus Torvalds (C) 2002 David Woodhouse <dwmw2@infradead.org> 646b6135aSMichel Lespinasse (C) 2012 Michel Lespinasse <walken@google.com> 71da177e4SLinus Torvalds 81da177e4SLinus Torvalds 91da177e4SLinus Torvalds linux/lib/rbtree.c 101da177e4SLinus Torvalds */ 111da177e4SLinus Torvalds 129c079addSMichel Lespinasse #include <linux/rbtree_augmented.h> 138bc3bcc9SPaul Gortmaker #include <linux/export.h> 141da177e4SLinus Torvalds 155bc9188aSMichel Lespinasse /* 16*d89775fcSAlexander A. Klimov * red-black trees properties: https://en.wikipedia.org/wiki/Rbtree 175bc9188aSMichel Lespinasse * 185bc9188aSMichel Lespinasse * 1) A node is either red or black 195bc9188aSMichel Lespinasse * 2) The root is black 205bc9188aSMichel Lespinasse * 3) All leaves (NULL) are black 215bc9188aSMichel Lespinasse * 4) Both children of every red node are black 225bc9188aSMichel Lespinasse * 5) Every simple path from root to leaves contains the same number 235bc9188aSMichel Lespinasse * of black nodes. 245bc9188aSMichel Lespinasse * 255bc9188aSMichel Lespinasse * 4 and 5 give the O(log n) guarantee, since 4 implies you cannot have two 265bc9188aSMichel Lespinasse * consecutive red nodes in a path and every red node is therefore followed by 275bc9188aSMichel Lespinasse * a black. So if B is the number of black nodes on every simple path (as per 285bc9188aSMichel Lespinasse * 5), then the longest possible path due to 4 is 2B. 295bc9188aSMichel Lespinasse * 305bc9188aSMichel Lespinasse * We shall indicate color with case, where black nodes are uppercase and red 316280d235SMichel Lespinasse * nodes will be lowercase. Unknown color nodes shall be drawn as red within 326280d235SMichel Lespinasse * parentheses and have some accompanying text comment. 335bc9188aSMichel Lespinasse */ 345bc9188aSMichel Lespinasse 35d72da4a4SPeter Zijlstra /* 36d72da4a4SPeter Zijlstra * Notes on lockless lookups: 37d72da4a4SPeter Zijlstra * 38d72da4a4SPeter Zijlstra * All stores to the tree structure (rb_left and rb_right) must be done using 39d72da4a4SPeter Zijlstra * WRITE_ONCE(). And we must not inadvertently cause (temporary) loops in the 40d72da4a4SPeter Zijlstra * tree structure as seen in program order. 41d72da4a4SPeter Zijlstra * 42d72da4a4SPeter Zijlstra * These two requirements will allow lockless iteration of the tree -- not 43d72da4a4SPeter Zijlstra * correct iteration mind you, tree rotations are not atomic so a lookup might 44d72da4a4SPeter Zijlstra * miss entire subtrees. 45d72da4a4SPeter Zijlstra * 46d72da4a4SPeter Zijlstra * But they do guarantee that any such traversal will only see valid elements 47d72da4a4SPeter Zijlstra * and that it will indeed complete -- does not get stuck in a loop. 48d72da4a4SPeter Zijlstra * 49d72da4a4SPeter Zijlstra * It also guarantees that if the lookup returns an element it is the 'correct' 50d72da4a4SPeter Zijlstra * one. But not returning an element does _NOT_ mean it's not present. 51d72da4a4SPeter Zijlstra * 52d72da4a4SPeter Zijlstra * NOTE: 53d72da4a4SPeter Zijlstra * 54d72da4a4SPeter Zijlstra * Stores to __rb_parent_color are not important for simple lookups so those 55d72da4a4SPeter Zijlstra * are left undone as of now. Nor did I check for loops involving parent 56d72da4a4SPeter Zijlstra * pointers. 57d72da4a4SPeter Zijlstra */ 58d72da4a4SPeter Zijlstra 5946b6135aSMichel Lespinasse static inline void rb_set_black(struct rb_node *rb) 6046b6135aSMichel Lespinasse { 6146b6135aSMichel Lespinasse rb->__rb_parent_color |= RB_BLACK; 6246b6135aSMichel Lespinasse } 6346b6135aSMichel Lespinasse 645bc9188aSMichel Lespinasse static inline struct rb_node *rb_red_parent(struct rb_node *red) 655bc9188aSMichel Lespinasse { 665bc9188aSMichel Lespinasse return (struct rb_node *)red->__rb_parent_color; 675bc9188aSMichel Lespinasse } 685bc9188aSMichel Lespinasse 695bc9188aSMichel Lespinasse /* 705bc9188aSMichel Lespinasse * Helper function for rotations: 715bc9188aSMichel Lespinasse * - old's parent and color get assigned to new 725bc9188aSMichel Lespinasse * - old gets assigned new as a parent and 'color' as a color. 735bc9188aSMichel Lespinasse */ 745bc9188aSMichel Lespinasse static inline void 755bc9188aSMichel Lespinasse __rb_rotate_set_parents(struct rb_node *old, struct rb_node *new, 765bc9188aSMichel Lespinasse struct rb_root *root, int color) 775bc9188aSMichel Lespinasse { 785bc9188aSMichel Lespinasse struct rb_node *parent = rb_parent(old); 795bc9188aSMichel Lespinasse new->__rb_parent_color = old->__rb_parent_color; 805bc9188aSMichel Lespinasse rb_set_parent_color(old, new, color); 817abc704aSMichel Lespinasse __rb_change_child(old, new, parent, root); 825bc9188aSMichel Lespinasse } 835bc9188aSMichel Lespinasse 8414b94af0SMichel Lespinasse static __always_inline void 8514b94af0SMichel Lespinasse __rb_insert(struct rb_node *node, struct rb_root *root, 8614b94af0SMichel Lespinasse void (*augment_rotate)(struct rb_node *old, struct rb_node *new)) 871da177e4SLinus Torvalds { 885bc9188aSMichel Lespinasse struct rb_node *parent = rb_red_parent(node), *gparent, *tmp; 891da177e4SLinus Torvalds 906d58452dSMichel Lespinasse while (true) { 916d58452dSMichel Lespinasse /* 922aadf7fcSDavidlohr Bueso * Loop invariant: node is red. 936d58452dSMichel Lespinasse */ 942aadf7fcSDavidlohr Bueso if (unlikely(!parent)) { 952aadf7fcSDavidlohr Bueso /* 962aadf7fcSDavidlohr Bueso * The inserted node is root. Either this is the 972aadf7fcSDavidlohr Bueso * first node, or we recursed at Case 1 below and 982aadf7fcSDavidlohr Bueso * are no longer violating 4). 992aadf7fcSDavidlohr Bueso */ 1005bc9188aSMichel Lespinasse rb_set_parent_color(node, NULL, RB_BLACK); 1016d58452dSMichel Lespinasse break; 1022aadf7fcSDavidlohr Bueso } 1032aadf7fcSDavidlohr Bueso 1042aadf7fcSDavidlohr Bueso /* 1052aadf7fcSDavidlohr Bueso * If there is a black parent, we are done. 1062aadf7fcSDavidlohr Bueso * Otherwise, take some corrective action as, 1072aadf7fcSDavidlohr Bueso * per 4), we don't want a red root or two 1082aadf7fcSDavidlohr Bueso * consecutive red nodes. 1092aadf7fcSDavidlohr Bueso */ 1102aadf7fcSDavidlohr Bueso if(rb_is_black(parent)) 1116d58452dSMichel Lespinasse break; 1126d58452dSMichel Lespinasse 1135bc9188aSMichel Lespinasse gparent = rb_red_parent(parent); 1141da177e4SLinus Torvalds 1155bc9188aSMichel Lespinasse tmp = gparent->rb_right; 11659633abfSMichel Lespinasse if (parent != tmp) { /* parent == gparent->rb_left */ 1175bc9188aSMichel Lespinasse if (tmp && rb_is_red(tmp)) { 1185bc9188aSMichel Lespinasse /* 11935dc67d7SDavidlohr Bueso * Case 1 - node's uncle is red (color flips). 1205bc9188aSMichel Lespinasse * 1215bc9188aSMichel Lespinasse * G g 1225bc9188aSMichel Lespinasse * / \ / \ 1235bc9188aSMichel Lespinasse * p u --> P U 1245bc9188aSMichel Lespinasse * / / 1251b9c53e8SWei Yang * n n 1265bc9188aSMichel Lespinasse * 1275bc9188aSMichel Lespinasse * However, since g's parent might be red, and 1285bc9188aSMichel Lespinasse * 4) does not allow this, we need to recurse 1295bc9188aSMichel Lespinasse * at g. 1305bc9188aSMichel Lespinasse */ 1315bc9188aSMichel Lespinasse rb_set_parent_color(tmp, gparent, RB_BLACK); 1325bc9188aSMichel Lespinasse rb_set_parent_color(parent, gparent, RB_BLACK); 1331da177e4SLinus Torvalds node = gparent; 1345bc9188aSMichel Lespinasse parent = rb_parent(node); 1355bc9188aSMichel Lespinasse rb_set_parent_color(node, parent, RB_RED); 1361da177e4SLinus Torvalds continue; 1371da177e4SLinus Torvalds } 1381da177e4SLinus Torvalds 13959633abfSMichel Lespinasse tmp = parent->rb_right; 14059633abfSMichel Lespinasse if (node == tmp) { 1415bc9188aSMichel Lespinasse /* 14235dc67d7SDavidlohr Bueso * Case 2 - node's uncle is black and node is 14335dc67d7SDavidlohr Bueso * the parent's right child (left rotate at parent). 1445bc9188aSMichel Lespinasse * 1455bc9188aSMichel Lespinasse * G G 1465bc9188aSMichel Lespinasse * / \ / \ 1475bc9188aSMichel Lespinasse * p U --> n U 1485bc9188aSMichel Lespinasse * \ / 1495bc9188aSMichel Lespinasse * n p 1505bc9188aSMichel Lespinasse * 1515bc9188aSMichel Lespinasse * This still leaves us in violation of 4), the 1525bc9188aSMichel Lespinasse * continuation into Case 3 will fix that. 1535bc9188aSMichel Lespinasse */ 154d72da4a4SPeter Zijlstra tmp = node->rb_left; 155d72da4a4SPeter Zijlstra WRITE_ONCE(parent->rb_right, tmp); 156d72da4a4SPeter Zijlstra WRITE_ONCE(node->rb_left, parent); 1575bc9188aSMichel Lespinasse if (tmp) 1585bc9188aSMichel Lespinasse rb_set_parent_color(tmp, parent, 1595bc9188aSMichel Lespinasse RB_BLACK); 1605bc9188aSMichel Lespinasse rb_set_parent_color(parent, node, RB_RED); 16114b94af0SMichel Lespinasse augment_rotate(parent, node); 1621da177e4SLinus Torvalds parent = node; 16359633abfSMichel Lespinasse tmp = node->rb_right; 1641da177e4SLinus Torvalds } 1651da177e4SLinus Torvalds 1665bc9188aSMichel Lespinasse /* 16735dc67d7SDavidlohr Bueso * Case 3 - node's uncle is black and node is 16835dc67d7SDavidlohr Bueso * the parent's left child (right rotate at gparent). 1695bc9188aSMichel Lespinasse * 1705bc9188aSMichel Lespinasse * G P 1715bc9188aSMichel Lespinasse * / \ / \ 1725bc9188aSMichel Lespinasse * p U --> n g 1735bc9188aSMichel Lespinasse * / \ 1745bc9188aSMichel Lespinasse * n U 1755bc9188aSMichel Lespinasse */ 176d72da4a4SPeter Zijlstra WRITE_ONCE(gparent->rb_left, tmp); /* == parent->rb_right */ 177d72da4a4SPeter Zijlstra WRITE_ONCE(parent->rb_right, gparent); 1785bc9188aSMichel Lespinasse if (tmp) 1795bc9188aSMichel Lespinasse rb_set_parent_color(tmp, gparent, RB_BLACK); 1805bc9188aSMichel Lespinasse __rb_rotate_set_parents(gparent, parent, root, RB_RED); 18114b94af0SMichel Lespinasse augment_rotate(gparent, parent); 1821f052865SMichel Lespinasse break; 1831da177e4SLinus Torvalds } else { 1845bc9188aSMichel Lespinasse tmp = gparent->rb_left; 1855bc9188aSMichel Lespinasse if (tmp && rb_is_red(tmp)) { 1865bc9188aSMichel Lespinasse /* Case 1 - color flips */ 1875bc9188aSMichel Lespinasse rb_set_parent_color(tmp, gparent, RB_BLACK); 1885bc9188aSMichel Lespinasse rb_set_parent_color(parent, gparent, RB_BLACK); 1891da177e4SLinus Torvalds node = gparent; 1905bc9188aSMichel Lespinasse parent = rb_parent(node); 1915bc9188aSMichel Lespinasse rb_set_parent_color(node, parent, RB_RED); 1921da177e4SLinus Torvalds continue; 1931da177e4SLinus Torvalds } 1941da177e4SLinus Torvalds 19559633abfSMichel Lespinasse tmp = parent->rb_left; 19659633abfSMichel Lespinasse if (node == tmp) { 1975bc9188aSMichel Lespinasse /* Case 2 - right rotate at parent */ 198d72da4a4SPeter Zijlstra tmp = node->rb_right; 199d72da4a4SPeter Zijlstra WRITE_ONCE(parent->rb_left, tmp); 200d72da4a4SPeter Zijlstra WRITE_ONCE(node->rb_right, parent); 2015bc9188aSMichel Lespinasse if (tmp) 2025bc9188aSMichel Lespinasse rb_set_parent_color(tmp, parent, 2035bc9188aSMichel Lespinasse RB_BLACK); 2045bc9188aSMichel Lespinasse rb_set_parent_color(parent, node, RB_RED); 20514b94af0SMichel Lespinasse augment_rotate(parent, node); 2061da177e4SLinus Torvalds parent = node; 20759633abfSMichel Lespinasse tmp = node->rb_left; 2081da177e4SLinus Torvalds } 2091da177e4SLinus Torvalds 2105bc9188aSMichel Lespinasse /* Case 3 - left rotate at gparent */ 211d72da4a4SPeter Zijlstra WRITE_ONCE(gparent->rb_right, tmp); /* == parent->rb_left */ 212d72da4a4SPeter Zijlstra WRITE_ONCE(parent->rb_left, gparent); 2135bc9188aSMichel Lespinasse if (tmp) 2145bc9188aSMichel Lespinasse rb_set_parent_color(tmp, gparent, RB_BLACK); 2155bc9188aSMichel Lespinasse __rb_rotate_set_parents(gparent, parent, root, RB_RED); 21614b94af0SMichel Lespinasse augment_rotate(gparent, parent); 2171f052865SMichel Lespinasse break; 2181da177e4SLinus Torvalds } 2191da177e4SLinus Torvalds } 2201da177e4SLinus Torvalds } 2211da177e4SLinus Torvalds 2223cb7a563SMichel Lespinasse /* 2233cb7a563SMichel Lespinasse * Inline version for rb_erase() use - we want to be able to inline 2243cb7a563SMichel Lespinasse * and eliminate the dummy_rotate callback there 2253cb7a563SMichel Lespinasse */ 2263cb7a563SMichel Lespinasse static __always_inline void 2273cb7a563SMichel Lespinasse ____rb_erase_color(struct rb_node *parent, struct rb_root *root, 2289c079addSMichel Lespinasse void (*augment_rotate)(struct rb_node *old, struct rb_node *new)) 2291da177e4SLinus Torvalds { 23046b6135aSMichel Lespinasse struct rb_node *node = NULL, *sibling, *tmp1, *tmp2; 2311da177e4SLinus Torvalds 232d6ff1273SMichel Lespinasse while (true) { 233d6ff1273SMichel Lespinasse /* 23446b6135aSMichel Lespinasse * Loop invariants: 23546b6135aSMichel Lespinasse * - node is black (or NULL on first iteration) 23646b6135aSMichel Lespinasse * - node is not the root (parent is not NULL) 23746b6135aSMichel Lespinasse * - All leaf paths going through parent and node have a 238d6ff1273SMichel Lespinasse * black node count that is 1 lower than other leaf paths. 239d6ff1273SMichel Lespinasse */ 2406280d235SMichel Lespinasse sibling = parent->rb_right; 24159633abfSMichel Lespinasse if (node != sibling) { /* node == parent->rb_left */ 2426280d235SMichel Lespinasse if (rb_is_red(sibling)) { 2436280d235SMichel Lespinasse /* 2446280d235SMichel Lespinasse * Case 1 - left rotate at parent 2456280d235SMichel Lespinasse * 2466280d235SMichel Lespinasse * P S 2476280d235SMichel Lespinasse * / \ / \ 2486280d235SMichel Lespinasse * N s --> p Sr 2496280d235SMichel Lespinasse * / \ / \ 2506280d235SMichel Lespinasse * Sl Sr N Sl 2516280d235SMichel Lespinasse */ 252d72da4a4SPeter Zijlstra tmp1 = sibling->rb_left; 253d72da4a4SPeter Zijlstra WRITE_ONCE(parent->rb_right, tmp1); 254d72da4a4SPeter Zijlstra WRITE_ONCE(sibling->rb_left, parent); 2556280d235SMichel Lespinasse rb_set_parent_color(tmp1, parent, RB_BLACK); 2566280d235SMichel Lespinasse __rb_rotate_set_parents(parent, sibling, root, 2576280d235SMichel Lespinasse RB_RED); 2589c079addSMichel Lespinasse augment_rotate(parent, sibling); 2596280d235SMichel Lespinasse sibling = tmp1; 2601da177e4SLinus Torvalds } 2616280d235SMichel Lespinasse tmp1 = sibling->rb_right; 2626280d235SMichel Lespinasse if (!tmp1 || rb_is_black(tmp1)) { 2636280d235SMichel Lespinasse tmp2 = sibling->rb_left; 2646280d235SMichel Lespinasse if (!tmp2 || rb_is_black(tmp2)) { 2656280d235SMichel Lespinasse /* 2666280d235SMichel Lespinasse * Case 2 - sibling color flip 2676280d235SMichel Lespinasse * (p could be either color here) 2686280d235SMichel Lespinasse * 2696280d235SMichel Lespinasse * (p) (p) 2706280d235SMichel Lespinasse * / \ / \ 2716280d235SMichel Lespinasse * N S --> N s 2726280d235SMichel Lespinasse * / \ / \ 2736280d235SMichel Lespinasse * Sl Sr Sl Sr 2746280d235SMichel Lespinasse * 27546b6135aSMichel Lespinasse * This leaves us violating 5) which 27646b6135aSMichel Lespinasse * can be fixed by flipping p to black 27746b6135aSMichel Lespinasse * if it was red, or by recursing at p. 27846b6135aSMichel Lespinasse * p is red when coming from Case 1. 2796280d235SMichel Lespinasse */ 2806280d235SMichel Lespinasse rb_set_parent_color(sibling, parent, 2816280d235SMichel Lespinasse RB_RED); 28246b6135aSMichel Lespinasse if (rb_is_red(parent)) 28346b6135aSMichel Lespinasse rb_set_black(parent); 28446b6135aSMichel Lespinasse else { 2851da177e4SLinus Torvalds node = parent; 28655a98102SDavid Woodhouse parent = rb_parent(node); 28746b6135aSMichel Lespinasse if (parent) 288e125d147SMichel Lespinasse continue; 2891da177e4SLinus Torvalds } 29046b6135aSMichel Lespinasse break; 29146b6135aSMichel Lespinasse } 2926280d235SMichel Lespinasse /* 2936280d235SMichel Lespinasse * Case 3 - right rotate at sibling 2946280d235SMichel Lespinasse * (p could be either color here) 2956280d235SMichel Lespinasse * 2966280d235SMichel Lespinasse * (p) (p) 2976280d235SMichel Lespinasse * / \ / \ 298ce093a04SJie Chen * N S --> N sl 2996280d235SMichel Lespinasse * / \ \ 300ce093a04SJie Chen * sl Sr S 301ce093a04SJie Chen * \ 302ce093a04SJie Chen * Sr 303ce093a04SJie Chen * 304ce093a04SJie Chen * Note: p might be red, and then both 305ce093a04SJie Chen * p and sl are red after rotation(which 306ce093a04SJie Chen * breaks property 4). This is fixed in 307ce093a04SJie Chen * Case 4 (in __rb_rotate_set_parents() 308ce093a04SJie Chen * which set sl the color of p 309ce093a04SJie Chen * and set p RB_BLACK) 310ce093a04SJie Chen * 311ce093a04SJie Chen * (p) (sl) 312ce093a04SJie Chen * / \ / \ 313ce093a04SJie Chen * N sl --> P S 314ce093a04SJie Chen * \ / \ 315ce093a04SJie Chen * S N Sr 3166280d235SMichel Lespinasse * \ 3176280d235SMichel Lespinasse * Sr 3186280d235SMichel Lespinasse */ 319d72da4a4SPeter Zijlstra tmp1 = tmp2->rb_right; 320d72da4a4SPeter Zijlstra WRITE_ONCE(sibling->rb_left, tmp1); 321d72da4a4SPeter Zijlstra WRITE_ONCE(tmp2->rb_right, sibling); 322d72da4a4SPeter Zijlstra WRITE_ONCE(parent->rb_right, tmp2); 3236280d235SMichel Lespinasse if (tmp1) 3246280d235SMichel Lespinasse rb_set_parent_color(tmp1, sibling, 3256280d235SMichel Lespinasse RB_BLACK); 3269c079addSMichel Lespinasse augment_rotate(sibling, tmp2); 3276280d235SMichel Lespinasse tmp1 = sibling; 3286280d235SMichel Lespinasse sibling = tmp2; 3291da177e4SLinus Torvalds } 3306280d235SMichel Lespinasse /* 3316280d235SMichel Lespinasse * Case 4 - left rotate at parent + color flips 3326280d235SMichel Lespinasse * (p and sl could be either color here. 3336280d235SMichel Lespinasse * After rotation, p becomes black, s acquires 3346280d235SMichel Lespinasse * p's color, and sl keeps its color) 3356280d235SMichel Lespinasse * 3366280d235SMichel Lespinasse * (p) (s) 3376280d235SMichel Lespinasse * / \ / \ 3386280d235SMichel Lespinasse * N S --> P Sr 3396280d235SMichel Lespinasse * / \ / \ 3406280d235SMichel Lespinasse * (sl) sr N (sl) 3416280d235SMichel Lespinasse */ 342d72da4a4SPeter Zijlstra tmp2 = sibling->rb_left; 343d72da4a4SPeter Zijlstra WRITE_ONCE(parent->rb_right, tmp2); 344d72da4a4SPeter Zijlstra WRITE_ONCE(sibling->rb_left, parent); 3456280d235SMichel Lespinasse rb_set_parent_color(tmp1, sibling, RB_BLACK); 3466280d235SMichel Lespinasse if (tmp2) 3476280d235SMichel Lespinasse rb_set_parent(tmp2, parent); 3486280d235SMichel Lespinasse __rb_rotate_set_parents(parent, sibling, root, 3496280d235SMichel Lespinasse RB_BLACK); 3509c079addSMichel Lespinasse augment_rotate(parent, sibling); 3511da177e4SLinus Torvalds break; 352d6ff1273SMichel Lespinasse } else { 3536280d235SMichel Lespinasse sibling = parent->rb_left; 3546280d235SMichel Lespinasse if (rb_is_red(sibling)) { 3556280d235SMichel Lespinasse /* Case 1 - right rotate at parent */ 356d72da4a4SPeter Zijlstra tmp1 = sibling->rb_right; 357d72da4a4SPeter Zijlstra WRITE_ONCE(parent->rb_left, tmp1); 358d72da4a4SPeter Zijlstra WRITE_ONCE(sibling->rb_right, parent); 3596280d235SMichel Lespinasse rb_set_parent_color(tmp1, parent, RB_BLACK); 3606280d235SMichel Lespinasse __rb_rotate_set_parents(parent, sibling, root, 3616280d235SMichel Lespinasse RB_RED); 3629c079addSMichel Lespinasse augment_rotate(parent, sibling); 3636280d235SMichel Lespinasse sibling = tmp1; 3641da177e4SLinus Torvalds } 3656280d235SMichel Lespinasse tmp1 = sibling->rb_left; 3666280d235SMichel Lespinasse if (!tmp1 || rb_is_black(tmp1)) { 3676280d235SMichel Lespinasse tmp2 = sibling->rb_right; 3686280d235SMichel Lespinasse if (!tmp2 || rb_is_black(tmp2)) { 3696280d235SMichel Lespinasse /* Case 2 - sibling color flip */ 3706280d235SMichel Lespinasse rb_set_parent_color(sibling, parent, 3716280d235SMichel Lespinasse RB_RED); 37246b6135aSMichel Lespinasse if (rb_is_red(parent)) 37346b6135aSMichel Lespinasse rb_set_black(parent); 37446b6135aSMichel Lespinasse else { 3751da177e4SLinus Torvalds node = parent; 37655a98102SDavid Woodhouse parent = rb_parent(node); 37746b6135aSMichel Lespinasse if (parent) 378e125d147SMichel Lespinasse continue; 3791da177e4SLinus Torvalds } 38046b6135aSMichel Lespinasse break; 38146b6135aSMichel Lespinasse } 382ce093a04SJie Chen /* Case 3 - left rotate at sibling */ 383d72da4a4SPeter Zijlstra tmp1 = tmp2->rb_left; 384d72da4a4SPeter Zijlstra WRITE_ONCE(sibling->rb_right, tmp1); 385d72da4a4SPeter Zijlstra WRITE_ONCE(tmp2->rb_left, sibling); 386d72da4a4SPeter Zijlstra WRITE_ONCE(parent->rb_left, tmp2); 3876280d235SMichel Lespinasse if (tmp1) 3886280d235SMichel Lespinasse rb_set_parent_color(tmp1, sibling, 3896280d235SMichel Lespinasse RB_BLACK); 3909c079addSMichel Lespinasse augment_rotate(sibling, tmp2); 3916280d235SMichel Lespinasse tmp1 = sibling; 3926280d235SMichel Lespinasse sibling = tmp2; 3931da177e4SLinus Torvalds } 394ce093a04SJie Chen /* Case 4 - right rotate at parent + color flips */ 395d72da4a4SPeter Zijlstra tmp2 = sibling->rb_right; 396d72da4a4SPeter Zijlstra WRITE_ONCE(parent->rb_left, tmp2); 397d72da4a4SPeter Zijlstra WRITE_ONCE(sibling->rb_right, parent); 3986280d235SMichel Lespinasse rb_set_parent_color(tmp1, sibling, RB_BLACK); 3996280d235SMichel Lespinasse if (tmp2) 4006280d235SMichel Lespinasse rb_set_parent(tmp2, parent); 4016280d235SMichel Lespinasse __rb_rotate_set_parents(parent, sibling, root, 4026280d235SMichel Lespinasse RB_BLACK); 4039c079addSMichel Lespinasse augment_rotate(parent, sibling); 4041da177e4SLinus Torvalds break; 4051da177e4SLinus Torvalds } 4061da177e4SLinus Torvalds } 4071da177e4SLinus Torvalds } 4083cb7a563SMichel Lespinasse 4093cb7a563SMichel Lespinasse /* Non-inline version for rb_erase_augmented() use */ 4103cb7a563SMichel Lespinasse void __rb_erase_color(struct rb_node *parent, struct rb_root *root, 4113cb7a563SMichel Lespinasse void (*augment_rotate)(struct rb_node *old, struct rb_node *new)) 4123cb7a563SMichel Lespinasse { 4133cb7a563SMichel Lespinasse ____rb_erase_color(parent, root, augment_rotate); 4143cb7a563SMichel Lespinasse } 4159c079addSMichel Lespinasse EXPORT_SYMBOL(__rb_erase_color); 41614b94af0SMichel Lespinasse 41714b94af0SMichel Lespinasse /* 41814b94af0SMichel Lespinasse * Non-augmented rbtree manipulation functions. 41914b94af0SMichel Lespinasse * 42014b94af0SMichel Lespinasse * We use dummy augmented callbacks here, and have the compiler optimize them 42114b94af0SMichel Lespinasse * out of the rb_insert_color() and rb_erase() function definitions. 42214b94af0SMichel Lespinasse */ 42314b94af0SMichel Lespinasse 42414b94af0SMichel Lespinasse static inline void dummy_propagate(struct rb_node *node, struct rb_node *stop) {} 42514b94af0SMichel Lespinasse static inline void dummy_copy(struct rb_node *old, struct rb_node *new) {} 42614b94af0SMichel Lespinasse static inline void dummy_rotate(struct rb_node *old, struct rb_node *new) {} 42714b94af0SMichel Lespinasse 42814b94af0SMichel Lespinasse static const struct rb_augment_callbacks dummy_callbacks = { 429f231aebfSKees Cook .propagate = dummy_propagate, 430f231aebfSKees Cook .copy = dummy_copy, 431f231aebfSKees Cook .rotate = dummy_rotate 43214b94af0SMichel Lespinasse }; 43314b94af0SMichel Lespinasse 43414b94af0SMichel Lespinasse void rb_insert_color(struct rb_node *node, struct rb_root *root) 43514b94af0SMichel Lespinasse { 4369f973cb3SMichel Lespinasse __rb_insert(node, root, dummy_rotate); 43714b94af0SMichel Lespinasse } 43814b94af0SMichel Lespinasse EXPORT_SYMBOL(rb_insert_color); 43914b94af0SMichel Lespinasse 44014b94af0SMichel Lespinasse void rb_erase(struct rb_node *node, struct rb_root *root) 44114b94af0SMichel Lespinasse { 4423cb7a563SMichel Lespinasse struct rb_node *rebalance; 4439f973cb3SMichel Lespinasse rebalance = __rb_erase_augmented(node, root, &dummy_callbacks); 4443cb7a563SMichel Lespinasse if (rebalance) 4453cb7a563SMichel Lespinasse ____rb_erase_color(rebalance, root, dummy_rotate); 4461da177e4SLinus Torvalds } 4471da177e4SLinus Torvalds EXPORT_SYMBOL(rb_erase); 4481da177e4SLinus Torvalds 44914b94af0SMichel Lespinasse /* 45014b94af0SMichel Lespinasse * Augmented rbtree manipulation functions. 45114b94af0SMichel Lespinasse * 45214b94af0SMichel Lespinasse * This instantiates the same __always_inline functions as in the non-augmented 45314b94af0SMichel Lespinasse * case, but this time with user-defined callbacks. 45414b94af0SMichel Lespinasse */ 45514b94af0SMichel Lespinasse 45614b94af0SMichel Lespinasse void __rb_insert_augmented(struct rb_node *node, struct rb_root *root, 45714b94af0SMichel Lespinasse void (*augment_rotate)(struct rb_node *old, struct rb_node *new)) 45814b94af0SMichel Lespinasse { 4599f973cb3SMichel Lespinasse __rb_insert(node, root, augment_rotate); 46014b94af0SMichel Lespinasse } 46114b94af0SMichel Lespinasse EXPORT_SYMBOL(__rb_insert_augmented); 46214b94af0SMichel Lespinasse 4631da177e4SLinus Torvalds /* 4641da177e4SLinus Torvalds * This function returns the first node (in sort order) of the tree. 4651da177e4SLinus Torvalds */ 466f4b477c4SArtem Bityutskiy struct rb_node *rb_first(const struct rb_root *root) 4671da177e4SLinus Torvalds { 4681da177e4SLinus Torvalds struct rb_node *n; 4691da177e4SLinus Torvalds 4701da177e4SLinus Torvalds n = root->rb_node; 4711da177e4SLinus Torvalds if (!n) 4721da177e4SLinus Torvalds return NULL; 4731da177e4SLinus Torvalds while (n->rb_left) 4741da177e4SLinus Torvalds n = n->rb_left; 4751da177e4SLinus Torvalds return n; 4761da177e4SLinus Torvalds } 4771da177e4SLinus Torvalds EXPORT_SYMBOL(rb_first); 4781da177e4SLinus Torvalds 479f4b477c4SArtem Bityutskiy struct rb_node *rb_last(const struct rb_root *root) 4801da177e4SLinus Torvalds { 4811da177e4SLinus Torvalds struct rb_node *n; 4821da177e4SLinus Torvalds 4831da177e4SLinus Torvalds n = root->rb_node; 4841da177e4SLinus Torvalds if (!n) 4851da177e4SLinus Torvalds return NULL; 4861da177e4SLinus Torvalds while (n->rb_right) 4871da177e4SLinus Torvalds n = n->rb_right; 4881da177e4SLinus Torvalds return n; 4891da177e4SLinus Torvalds } 4901da177e4SLinus Torvalds EXPORT_SYMBOL(rb_last); 4911da177e4SLinus Torvalds 492f4b477c4SArtem Bityutskiy struct rb_node *rb_next(const struct rb_node *node) 4931da177e4SLinus Torvalds { 49455a98102SDavid Woodhouse struct rb_node *parent; 49555a98102SDavid Woodhouse 4964c199a93SMichel Lespinasse if (RB_EMPTY_NODE(node)) 49710fd48f2SJens Axboe return NULL; 49810fd48f2SJens Axboe 4997ce6ff9eSMichel Lespinasse /* 5007ce6ff9eSMichel Lespinasse * If we have a right-hand child, go down and then left as far 5017ce6ff9eSMichel Lespinasse * as we can. 5027ce6ff9eSMichel Lespinasse */ 5031da177e4SLinus Torvalds if (node->rb_right) { 5041da177e4SLinus Torvalds node = node->rb_right; 5051da177e4SLinus Torvalds while (node->rb_left) 5061da177e4SLinus Torvalds node = node->rb_left; 507f4b477c4SArtem Bityutskiy return (struct rb_node *)node; 5081da177e4SLinus Torvalds } 5091da177e4SLinus Torvalds 5107ce6ff9eSMichel Lespinasse /* 5117ce6ff9eSMichel Lespinasse * No right-hand children. Everything down and left is smaller than us, 5127ce6ff9eSMichel Lespinasse * so any 'next' node must be in the general direction of our parent. 5137ce6ff9eSMichel Lespinasse * Go up the tree; any time the ancestor is a right-hand child of its 5147ce6ff9eSMichel Lespinasse * parent, keep going up. First time it's a left-hand child of its 5157ce6ff9eSMichel Lespinasse * parent, said parent is our 'next' node. 5167ce6ff9eSMichel Lespinasse */ 51755a98102SDavid Woodhouse while ((parent = rb_parent(node)) && node == parent->rb_right) 51855a98102SDavid Woodhouse node = parent; 5191da177e4SLinus Torvalds 52055a98102SDavid Woodhouse return parent; 5211da177e4SLinus Torvalds } 5221da177e4SLinus Torvalds EXPORT_SYMBOL(rb_next); 5231da177e4SLinus Torvalds 524f4b477c4SArtem Bityutskiy struct rb_node *rb_prev(const struct rb_node *node) 5251da177e4SLinus Torvalds { 52655a98102SDavid Woodhouse struct rb_node *parent; 52755a98102SDavid Woodhouse 5284c199a93SMichel Lespinasse if (RB_EMPTY_NODE(node)) 52910fd48f2SJens Axboe return NULL; 53010fd48f2SJens Axboe 5317ce6ff9eSMichel Lespinasse /* 5327ce6ff9eSMichel Lespinasse * If we have a left-hand child, go down and then right as far 5337ce6ff9eSMichel Lespinasse * as we can. 5347ce6ff9eSMichel Lespinasse */ 5351da177e4SLinus Torvalds if (node->rb_left) { 5361da177e4SLinus Torvalds node = node->rb_left; 5371da177e4SLinus Torvalds while (node->rb_right) 5381da177e4SLinus Torvalds node = node->rb_right; 539f4b477c4SArtem Bityutskiy return (struct rb_node *)node; 5401da177e4SLinus Torvalds } 5411da177e4SLinus Torvalds 5427ce6ff9eSMichel Lespinasse /* 5437ce6ff9eSMichel Lespinasse * No left-hand children. Go up till we find an ancestor which 5447ce6ff9eSMichel Lespinasse * is a right-hand child of its parent. 5457ce6ff9eSMichel Lespinasse */ 54655a98102SDavid Woodhouse while ((parent = rb_parent(node)) && node == parent->rb_left) 54755a98102SDavid Woodhouse node = parent; 5481da177e4SLinus Torvalds 54955a98102SDavid Woodhouse return parent; 5501da177e4SLinus Torvalds } 5511da177e4SLinus Torvalds EXPORT_SYMBOL(rb_prev); 5521da177e4SLinus Torvalds 5531da177e4SLinus Torvalds void rb_replace_node(struct rb_node *victim, struct rb_node *new, 5541da177e4SLinus Torvalds struct rb_root *root) 5551da177e4SLinus Torvalds { 55655a98102SDavid Woodhouse struct rb_node *parent = rb_parent(victim); 5571da177e4SLinus Torvalds 558c1adf200SDavid Howells /* Copy the pointers/colour from the victim to the replacement */ 559c1adf200SDavid Howells *new = *victim; 560c1adf200SDavid Howells 5611da177e4SLinus Torvalds /* Set the surrounding nodes to point to the replacement */ 562c1adf200SDavid Howells if (victim->rb_left) 563c1adf200SDavid Howells rb_set_parent(victim->rb_left, new); 564c1adf200SDavid Howells if (victim->rb_right) 565c1adf200SDavid Howells rb_set_parent(victim->rb_right, new); 5667abc704aSMichel Lespinasse __rb_change_child(victim, new, parent, root); 567c1adf200SDavid Howells } 568c1adf200SDavid Howells EXPORT_SYMBOL(rb_replace_node); 569c1adf200SDavid Howells 570c1adf200SDavid Howells void rb_replace_node_rcu(struct rb_node *victim, struct rb_node *new, 571c1adf200SDavid Howells struct rb_root *root) 572c1adf200SDavid Howells { 573c1adf200SDavid Howells struct rb_node *parent = rb_parent(victim); 574c1adf200SDavid Howells 575c1adf200SDavid Howells /* Copy the pointers/colour from the victim to the replacement */ 576c1adf200SDavid Howells *new = *victim; 577c1adf200SDavid Howells 578c1adf200SDavid Howells /* Set the surrounding nodes to point to the replacement */ 5791da177e4SLinus Torvalds if (victim->rb_left) 58055a98102SDavid Woodhouse rb_set_parent(victim->rb_left, new); 5811da177e4SLinus Torvalds if (victim->rb_right) 58255a98102SDavid Woodhouse rb_set_parent(victim->rb_right, new); 5831da177e4SLinus Torvalds 584c1adf200SDavid Howells /* Set the parent's pointer to the new node last after an RCU barrier 585c1adf200SDavid Howells * so that the pointers onwards are seen to be set correctly when doing 586c1adf200SDavid Howells * an RCU walk over the tree. 587c1adf200SDavid Howells */ 588c1adf200SDavid Howells __rb_change_child_rcu(victim, new, parent, root); 5891da177e4SLinus Torvalds } 590c1adf200SDavid Howells EXPORT_SYMBOL(rb_replace_node_rcu); 5919dee5c51SCody P Schafer 5929dee5c51SCody P Schafer static struct rb_node *rb_left_deepest_node(const struct rb_node *node) 5939dee5c51SCody P Schafer { 5949dee5c51SCody P Schafer for (;;) { 5959dee5c51SCody P Schafer if (node->rb_left) 5969dee5c51SCody P Schafer node = node->rb_left; 5979dee5c51SCody P Schafer else if (node->rb_right) 5989dee5c51SCody P Schafer node = node->rb_right; 5999dee5c51SCody P Schafer else 6009dee5c51SCody P Schafer return (struct rb_node *)node; 6019dee5c51SCody P Schafer } 6029dee5c51SCody P Schafer } 6039dee5c51SCody P Schafer 6049dee5c51SCody P Schafer struct rb_node *rb_next_postorder(const struct rb_node *node) 6059dee5c51SCody P Schafer { 6069dee5c51SCody P Schafer const struct rb_node *parent; 6079dee5c51SCody P Schafer if (!node) 6089dee5c51SCody P Schafer return NULL; 6099dee5c51SCody P Schafer parent = rb_parent(node); 6109dee5c51SCody P Schafer 6119dee5c51SCody P Schafer /* If we're sitting on node, we've already seen our children */ 6129dee5c51SCody P Schafer if (parent && node == parent->rb_left && parent->rb_right) { 6139dee5c51SCody P Schafer /* If we are the parent's left node, go to the parent's right 6149dee5c51SCody P Schafer * node then all the way down to the left */ 6159dee5c51SCody P Schafer return rb_left_deepest_node(parent->rb_right); 6169dee5c51SCody P Schafer } else 6179dee5c51SCody P Schafer /* Otherwise we are the parent's right node, and the parent 6189dee5c51SCody P Schafer * should be next */ 6199dee5c51SCody P Schafer return (struct rb_node *)parent; 6209dee5c51SCody P Schafer } 6219dee5c51SCody P Schafer EXPORT_SYMBOL(rb_next_postorder); 6229dee5c51SCody P Schafer 6239dee5c51SCody P Schafer struct rb_node *rb_first_postorder(const struct rb_root *root) 6249dee5c51SCody P Schafer { 6259dee5c51SCody P Schafer if (!root->rb_node) 6269dee5c51SCody P Schafer return NULL; 6279dee5c51SCody P Schafer 6289dee5c51SCody P Schafer return rb_left_deepest_node(root->rb_node); 6299dee5c51SCody P Schafer } 6309dee5c51SCody P Schafer EXPORT_SYMBOL(rb_first_postorder); 631