lib/rbtree.c

   1 // SPDX-License-Identifier: GPL-2.0-or-later
   2 /*
   3   Red Black Trees
   4   (C) 1999  Andrea Arcangeli <andrea@suse.de>
   5   (C) 2002  David Woodhouse <dwmw2@infradead.org>
   6   (C) 2012  Michel Lespinasse <walken@google.com>
   7
   8
   9   linux/lib/rbtree.c
  10 */
  11
  12 #include <linux/rbtree_augmented.h>
  13 #include <linux/export.h>
  14
  15 /*
  16  * red-black trees properties:  https://en.wikipedia.org/wiki/Rbtree
  17  *
  18  *  1) A node is either red or black
  19  *  2) The root is black
  20  *  3) All leaves (NULL) are black
  21  *  4) Both children of every red node are black
  22  *  5) Every simple path from root to leaves contains the same number
  23  *     of black nodes.
  24  *
  25  *  4 and 5 give the O(log n) guarantee, since 4 implies you cannot have two
  26  *  consecutive red nodes in a path and every red node is therefore followed by
  27  *  a black. So if B is the number of black nodes on every simple path (as per
  28  *  5), then the longest possible path due to 4 is 2B.
  29  *
  30  *  We shall indicate color with case, where black nodes are uppercase and red
  31  *  nodes will be lowercase. Unknown color nodes shall be drawn as red within
  32  *  parentheses and have some accompanying text comment.
  33  */
  34
  35 /*
  36  * Notes on lockless lookups:
  37  *
  38  * All stores to the tree structure (rb_left and rb_right) must be done using
  39  * WRITE_ONCE(). And we must not inadvertently cause (temporary) loops in the
  40  * tree structure as seen in program order.
  41  *
  42  * These two requirements will allow lockless iteration of the tree -- not
  43  * correct iteration mind you, tree rotations are not atomic so a lookup might
  44  * miss entire subtrees.
  45  *
  46  * But they do guarantee that any such traversal will only see valid elements
  47  * and that it will indeed complete -- does not get stuck in a loop.
  48  *
  49  * It also guarantees that if the lookup returns an element it is the 'correct'
  50  * one. But not returning an element does _NOT_ mean it's not present.
  51  *
  52  * NOTE:
  53  *
  54  * Stores to __rb_parent_color are not important for simple lookups so those
  55  * are left undone as of now. Nor did I check for loops involving parent
  56  * pointers.
  57  */
  58
  59 static inline void rb_set_black(struct rb_node *rb)
  60 {
  61         rb->__rb_parent_color += RB_BLACK;
  62 }
  63
  64 static inline struct rb_node *rb_red_parent(struct rb_node *red)
  65 {
  66         return (struct rb_node *)red->__rb_parent_color;
  67 }
  68
  69 /*
  70  * Helper function for rotations:
  71  * - old's parent and color get assigned to new
  72  * - old gets assigned new as a parent and 'color' as a color.
  73  */
  74 static inline void
  75 __rb_rotate_set_parents(struct rb_node *old, struct rb_node *new,
  76                         struct rb_root *root, int color)
  77 {
  78         struct rb_node *parent = rb_parent(old);
  79         new->__rb_parent_color = old->__rb_parent_color;
  80         rb_set_parent_color(old, new, color);
  81         __rb_change_child(old, new, parent, root);
  82 }
  83
  84 static __always_inline void
  85 __rb_insert(struct rb_node *node, struct rb_root *root,
  86             void (*augment_rotate)(struct rb_node *old, struct rb_node *new))
  87 {
  88         struct rb_node *parent = rb_red_parent(node), *gparent, *tmp;
  89
  90         while (true) {
  91                 /*
  92                  * Loop invariant: node is red.
  93                  */
  94                 if (unlikely(!parent)) {
  95                         /*
  96                          * The inserted node is root. Either this is the
  97                          * first node, or we recursed at Case 1 below and
  98                          * are no longer violating 4).
  99                          */
 100                         rb_set_parent_color(node, NULL, RB_BLACK);
 101                         break;
 102                 }
 103
 104                 /*
 105                  * If there is a black parent, we are done.
 106                  * Otherwise, take some corrective action as,
 107                  * per 4), we don't want a red root or two
 108                  * consecutive red nodes.
 109                  */
 110                 if(rb_is_black(parent))
 111                         break;
 112
 113                 gparent = rb_red_parent(parent);
 114
 115                 tmp = gparent->rb_right;
 116                 if (parent != tmp) {    /* parent == gparent->rb_left */
 117                         if (tmp && rb_is_red(tmp)) {
 118                                 /*
 119                                  * Case 1 - node's uncle is red (color flips).
 120                                  *
 121                                  *       G            g
 122                                  *      / \          / \
 123                                  *     p   u  -->   P   U
 124                                  *    /            /
 125                                  *   n            n
 126                                  *
 127                                  * However, since g's parent might be red, and
 128                                  * 4) does not allow this, we need to recurse
 129                                  * at g.
 130                                  */
 131                                 rb_set_parent_color(tmp, gparent, RB_BLACK);
 132                                 rb_set_parent_color(parent, gparent, RB_BLACK);
 133                                 node = gparent;
 134                                 parent = rb_parent(node);
 135                                 rb_set_parent_color(node, parent, RB_RED);
 136                                 continue;
 137                         }
 138
 139                         tmp = parent->rb_right;
 140                         if (node == tmp) {
 141                                 /*
 142                                  * Case 2 - node's uncle is black and node is
 143                                  * the parent's right child (left rotate at parent).
 144                                  *
 145                                  *      G             G
 146                                  *     / \           / \
 147                                  *    p   U  -->    n   U
 148                                  *     \           /
 149                                  *      n         p
 150                                  *
 151                                  * This still leaves us in violation of 4), the
 152                                  * continuation into Case 3 will fix that.
 153                                  */
 154                                 tmp = node->rb_left;
 155                                 WRITE_ONCE(parent->rb_right, tmp);
 156                                 WRITE_ONCE(node->rb_left, parent);
 157                                 if (tmp)
 158                                         rb_set_parent_color(tmp, parent,
 159                                                             RB_BLACK);
 160                                 rb_set_parent_color(parent, node, RB_RED);
 161                                 augment_rotate(parent, node);
 162                                 parent = node;
 163                                 tmp = node->rb_right;
 164                         }
 165
 166                         /*
 167                          * Case 3 - node's uncle is black and node is
 168                          * the parent's left child (right rotate at gparent).
 169                          *
 170                          *        G           P
 171                          *       / \         / \
 172                          *      p   U  -->  n   g
 173                          *     /                 \
 174                          *    n                   U
 175                          */
 176                         WRITE_ONCE(gparent->rb_left, tmp); /* == parent->rb_right */
 177                         WRITE_ONCE(parent->rb_right, gparent);
 178                         if (tmp)
 179                                 rb_set_parent_color(tmp, gparent, RB_BLACK);
 180                         __rb_rotate_set_parents(gparent, parent, root, RB_RED);
 181                         augment_rotate(gparent, parent);
 182                         break;
 183                 } else {
 184                         tmp = gparent->rb_left;
 185                         if (tmp && rb_is_red(tmp)) {
 186                                 /* Case 1 - color flips */
 187                                 rb_set_parent_color(tmp, gparent, RB_BLACK);
 188                                 rb_set_parent_color(parent, gparent, RB_BLACK);
 189                                 node = gparent;
 190                                 parent = rb_parent(node);
 191                                 rb_set_parent_color(node, parent, RB_RED);
 192                                 continue;
 193                         }
 194
 195                         tmp = parent->rb_left;
 196                         if (node == tmp) {
 197                                 /* Case 2 - right rotate at parent */
 198                                 tmp = node->rb_right;
 199                                 WRITE_ONCE(parent->rb_left, tmp);
 200                                 WRITE_ONCE(node->rb_right, parent);
 201                                 if (tmp)
 202                                         rb_set_parent_color(tmp, parent,
 203                                                             RB_BLACK);
 204                                 rb_set_parent_color(parent, node, RB_RED);
 205                                 augment_rotate(parent, node);
 206                                 parent = node;
 207                                 tmp = node->rb_left;
 208                         }
 209
 210                         /* Case 3 - left rotate at gparent */
 211                         WRITE_ONCE(gparent->rb_right, tmp); /* == parent->rb_left */
 212                         WRITE_ONCE(parent->rb_left, gparent);
 213                         if (tmp)
 214                                 rb_set_parent_color(tmp, gparent, RB_BLACK);
 215                         __rb_rotate_set_parents(gparent, parent, root, RB_RED);
 216                         augment_rotate(gparent, parent);
 217                         break;
 218                 }
 219         }
 220 }
 221
 222 /*
 223  * Inline version for rb_erase() use - we want to be able to inline
 224  * and eliminate the dummy_rotate callback there
 225  */
 226 static __always_inline void
 227 ____rb_erase_color(struct rb_node *parent, struct rb_root *root,
 228         void (*augment_rotate)(struct rb_node *old, struct rb_node *new))
 229 {
 230         struct rb_node *node = NULL, *sibling, *tmp1, *tmp2;
 231
 232         while (true) {
 233                 /*
 234                  * Loop invariants:
 235                  * - node is black (or NULL on first iteration)
 236                  * - node is not the root (parent is not NULL)
 237                  * - All leaf paths going through parent and node have a
 238                  *   black node count that is 1 lower than other leaf paths.
 239                  */
 240                 sibling = parent->rb_right;
 241                 if (node != sibling) {  /* node == parent->rb_left */
 242                         if (rb_is_red(sibling)) {
 243                                 /*
 244                                  * Case 1 - left rotate at parent
 245                                  *
 246                                  *     P               S
 247                                  *    / \             / \
 248                                  *   N   s    -->    p   Sr
 249                                  *      / \         / \
 250                                  *     Sl  Sr      N   Sl
 251                                  */
 252                                 tmp1 = sibling->rb_left;
 253                                 WRITE_ONCE(parent->rb_right, tmp1);
 254                                 WRITE_ONCE(sibling->rb_left, parent);
 255                                 rb_set_parent_color(tmp1, parent, RB_BLACK);
 256                                 __rb_rotate_set_parents(parent, sibling, root,
 257                                                         RB_RED);
 258                                 augment_rotate(parent, sibling);
 259                                 sibling = tmp1;
 260                         }
 261                         tmp1 = sibling->rb_right;
 262                         if (!tmp1 || rb_is_black(tmp1)) {
 263                                 tmp2 = sibling->rb_left;
 264                                 if (!tmp2 || rb_is_black(tmp2)) {
 265                                         /*
 266                                          * Case 2 - sibling color flip
 267                                          * (p could be either color here)
 268                                          *
 269                                          *    (p)           (p)
 270                                          *    / \           / \
 271                                          *   N   S    -->  N   s
 272                                          *      / \           / \
 273                                          *     Sl  Sr        Sl  Sr
 274                                          *
 275                                          * This leaves us violating 5) which
 276                                          * can be fixed by flipping p to black
 277                                          * if it was red, or by recursing at p.
 278                                          * p is red when coming from Case 1.
 279                                          */
 280                                         rb_set_parent_color(sibling, parent,
 281                                                             RB_RED);
 282                                         if (rb_is_red(parent))
 283                                                 rb_set_black(parent);
 284                                         else {
 285                                                 node = parent;
 286                                                 parent = rb_parent(node);
 287                                                 if (parent)
 288                                                         continue;
 289                                         }
 290                                         break;
 291                                 }
 292                                 /*
 293                                  * Case 3 - right rotate at sibling
 294                                  * (p could be either color here)
 295                                  *
 296                                  *   (p)           (p)
 297                                  *   / \           / \
 298                                  *  N   S    -->  N   sl
 299                                  *     / \             \
 300                                  *    sl  Sr            S
 301                                  *                       \
 302                                  *                        Sr
 303                                  *
 304                                  * Note: p might be red, and then both
 305                                  * p and sl are red after rotation(which
 306                                  * breaks property 4). This is fixed in
 307                                  * Case 4 (in __rb_rotate_set_parents()
 308                                  *         which set sl the color of p
 309                                  *         and set p RB_BLACK)
 310                                  *
 311                                  *   (p)            (sl)
 312                                  *   / \            /  \
 313                                  *  N   sl   -->   P    S
 314                                  *       \        /      \
 315                                  *        S      N        Sr
 316                                  *         \
 317                                  *          Sr
 318                                  */
 319                                 tmp1 = tmp2->rb_right;
 320                                 WRITE_ONCE(sibling->rb_left, tmp1);
 321                                 WRITE_ONCE(tmp2->rb_right, sibling);
 322                                 WRITE_ONCE(parent->rb_right, tmp2);
 323                                 if (tmp1)
 324                                         rb_set_parent_color(tmp1, sibling,
 325                                                             RB_BLACK);
 326                                 augment_rotate(sibling, tmp2);
 327                                 tmp1 = sibling;
 328                                 sibling = tmp2;
 329                         }
 330                         /*
 331                          * Case 4 - left rotate at parent + color flips
 332                          * (p and sl could be either color here.
 333                          *  After rotation, p becomes black, s acquires
 334                          *  p's color, and sl keeps its color)
 335                          *
 336                          *      (p)             (s)
 337                          *      / \             / \
 338                          *     N   S     -->   P   Sr
 339                          *        / \         / \
 340                          *      (sl) sr      N  (sl)
 341                          */
 342                         tmp2 = sibling->rb_left;
 343                         WRITE_ONCE(parent->rb_right, tmp2);
 344                         WRITE_ONCE(sibling->rb_left, parent);
 345                         rb_set_parent_color(tmp1, sibling, RB_BLACK);
 346                         if (tmp2)
 347                                 rb_set_parent(tmp2, parent);
 348                         __rb_rotate_set_parents(parent, sibling, root,
 349                                                 RB_BLACK);
 350                         augment_rotate(parent, sibling);
 351                         break;
 352                 } else {
 353                         sibling = parent->rb_left;
 354                         if (rb_is_red(sibling)) {
 355                                 /* Case 1 - right rotate at parent */
 356                                 tmp1 = sibling->rb_right;
 357                                 WRITE_ONCE(parent->rb_left, tmp1);
 358                                 WRITE_ONCE(sibling->rb_right, parent);
 359                                 rb_set_parent_color(tmp1, parent, RB_BLACK);
 360                                 __rb_rotate_set_parents(parent, sibling, root,
 361                                                         RB_RED);
 362                                 augment_rotate(parent, sibling);
 363                                 sibling = tmp1;
 364                         }
 365                         tmp1 = sibling->rb_left;
 366                         if (!tmp1 || rb_is_black(tmp1)) {
 367                                 tmp2 = sibling->rb_right;
 368                                 if (!tmp2 || rb_is_black(tmp2)) {
 369                                         /* Case 2 - sibling color flip */
 370                                         rb_set_parent_color(sibling, parent,
 371                                                             RB_RED);
 372                                         if (rb_is_red(parent))
 373                                                 rb_set_black(parent);
 374                                         else {
 375                                                 node = parent;
 376                                                 parent = rb_parent(node);
 377                                                 if (parent)
 378                                                         continue;
 379                                         }
 380                                         break;
 381                                 }
 382                                 /* Case 3 - left rotate at sibling */
 383                                 tmp1 = tmp2->rb_left;
 384                                 WRITE_ONCE(sibling->rb_right, tmp1);
 385                                 WRITE_ONCE(tmp2->rb_left, sibling);
 386                                 WRITE_ONCE(parent->rb_left, tmp2);
 387                                 if (tmp1)
 388                                         rb_set_parent_color(tmp1, sibling,
 389                                                             RB_BLACK);
 390                                 augment_rotate(sibling, tmp2);
 391                                 tmp1 = sibling;
 392                                 sibling = tmp2;
 393                         }
 394                         /* Case 4 - right rotate at parent + color flips */
 395                         tmp2 = sibling->rb_right;
 396                         WRITE_ONCE(parent->rb_left, tmp2);
 397                         WRITE_ONCE(sibling->rb_right, parent);
 398                         rb_set_parent_color(tmp1, sibling, RB_BLACK);
 399                         if (tmp2)
 400                                 rb_set_parent(tmp2, parent);
 401                         __rb_rotate_set_parents(parent, sibling, root,
 402                                                 RB_BLACK);
 403                         augment_rotate(parent, sibling);
 404                         break;
 405                 }
 406         }
 407 }
 408
 409 /* Non-inline version for rb_erase_augmented() use */
 410 void __rb_erase_color(struct rb_node *parent, struct rb_root *root,
 411         void (*augment_rotate)(struct rb_node *old, struct rb_node *new))
 412 {
 413         ____rb_erase_color(parent, root, augment_rotate);
 414 }
 415 EXPORT_SYMBOL(__rb_erase_color);
 416
 417 /*
 418  * Non-augmented rbtree manipulation functions.
 419  *
 420  * We use dummy augmented callbacks here, and have the compiler optimize them
 421  * out of the rb_insert_color() and rb_erase() function definitions.
 422  */
 423
 424 static inline void dummy_propagate(struct rb_node *node, struct rb_node *stop) {}
 425 static inline void dummy_copy(struct rb_node *old, struct rb_node *new) {}
 426 static inline void dummy_rotate(struct rb_node *old, struct rb_node *new) {}
 427
 428 static const struct rb_augment_callbacks dummy_callbacks = {
 429         .propagate = dummy_propagate,
 430         .copy = dummy_copy,
 431         .rotate = dummy_rotate
 432 };
 433
 434 void rb_insert_color(struct rb_node *node, struct rb_root *root)
 435 {
 436         __rb_insert(node, root, dummy_rotate);
 437 }
 438 EXPORT_SYMBOL(rb_insert_color);
 439
 440 void rb_erase(struct rb_node *node, struct rb_root *root)
 441 {
 442         struct rb_node *rebalance;
 443         rebalance = __rb_erase_augmented(node, root, &dummy_callbacks);
 444         if (rebalance)
 445                 ____rb_erase_color(rebalance, root, dummy_rotate);
 446 }
 447 EXPORT_SYMBOL(rb_erase);
 448
 449 /*
 450  * Augmented rbtree manipulation functions.
 451  *
 452  * This instantiates the same __always_inline functions as in the non-augmented
 453  * case, but this time with user-defined callbacks.
 454  */
 455
 456 void __rb_insert_augmented(struct rb_node *node, struct rb_root *root,
 457         void (*augment_rotate)(struct rb_node *old, struct rb_node *new))
 458 {
 459         __rb_insert(node, root, augment_rotate);
 460 }
 461 EXPORT_SYMBOL(__rb_insert_augmented);
 462
 463 /*
 464  * This function returns the first node (in sort order) of the tree.
 465  */
 466 struct rb_node *rb_first(const struct rb_root *root)
 467 {
 468         struct rb_node  *n;
 469
 470         n = root->rb_node;
 471         if (!n)
 472                 return NULL;
 473         while (n->rb_left)
 474                 n = n->rb_left;
 475         return n;
 476 }
 477 EXPORT_SYMBOL(rb_first);
 478
 479 struct rb_node *rb_last(const struct rb_root *root)
 480 {
 481         struct rb_node  *n;
 482
 483         n = root->rb_node;
 484         if (!n)
 485                 return NULL;
 486         while (n->rb_right)
 487                 n = n->rb_right;
 488         return n;
 489 }
 490 EXPORT_SYMBOL(rb_last);
 491
 492 struct rb_node *rb_next(const struct rb_node *node)
 493 {
 494         struct rb_node *parent;
 495
 496         if (RB_EMPTY_NODE(node))
 497                 return NULL;
 498
 499         /*
 500          * If we have a right-hand child, go down and then left as far
 501          * as we can.
 502          */
 503         if (node->rb_right) {
 504                 node = node->rb_right;
 505                 while (node->rb_left)
 506                         node = node->rb_left;
 507                 return (struct rb_node *)node;
 508         }
 509
 510         /*
 511          * No right-hand children. Everything down and left is smaller than us,
 512          * so any 'next' node must be in the general direction of our parent.
 513          * Go up the tree; any time the ancestor is a right-hand child of its
 514          * parent, keep going up. First time it's a left-hand child of its
 515          * parent, said parent is our 'next' node.
 516          */
 517         while ((parent = rb_parent(node)) && node == parent->rb_right)
 518                 node = parent;
 519
 520         return parent;
 521 }
 522 EXPORT_SYMBOL(rb_next);
 523
 524 struct rb_node *rb_prev(const struct rb_node *node)
 525 {
 526         struct rb_node *parent;
 527
 528         if (RB_EMPTY_NODE(node))
 529                 return NULL;
 530
 531         /*
 532          * If we have a left-hand child, go down and then right as far
 533          * as we can.
 534          */
 535         if (node->rb_left) {
 536                 node = node->rb_left;
 537                 while (node->rb_right)
 538                         node = node->rb_right;
 539                 return (struct rb_node *)node;
 540         }
 541
 542         /*
 543          * No left-hand children. Go up till we find an ancestor which
 544          * is a right-hand child of its parent.
 545          */
 546         while ((parent = rb_parent(node)) && node == parent->rb_left)
 547                 node = parent;
 548
 549         return parent;
 550 }
 551 EXPORT_SYMBOL(rb_prev);
 552
 553 void rb_replace_node(struct rb_node *victim, struct rb_node *new,
 554                      struct rb_root *root)
 555 {
 556         struct rb_node *parent = rb_parent(victim);
 557
 558         /* Copy the pointers/colour from the victim to the replacement */
 559         *new = *victim;
 560
 561         /* Set the surrounding nodes to point to the replacement */
 562         if (victim->rb_left)
 563                 rb_set_parent(victim->rb_left, new);
 564         if (victim->rb_right)
 565                 rb_set_parent(victim->rb_right, new);
 566         __rb_change_child(victim, new, parent, root);
 567 }
 568 EXPORT_SYMBOL(rb_replace_node);
 569
 570 void rb_replace_node_rcu(struct rb_node *victim, struct rb_node *new,
 571                          struct rb_root *root)
 572 {
 573         struct rb_node *parent = rb_parent(victim);
 574
 575         /* Copy the pointers/colour from the victim to the replacement */
 576         *new = *victim;
 577
 578         /* Set the surrounding nodes to point to the replacement */
 579         if (victim->rb_left)
 580                 rb_set_parent(victim->rb_left, new);
 581         if (victim->rb_right)
 582                 rb_set_parent(victim->rb_right, new);
 583
 584         /* Set the parent's pointer to the new node last after an RCU barrier
 585          * so that the pointers onwards are seen to be set correctly when doing
 586          * an RCU walk over the tree.
 587          */
 588         __rb_change_child_rcu(victim, new, parent, root);
 589 }
 590 EXPORT_SYMBOL(rb_replace_node_rcu);
 591
 592 static struct rb_node *rb_left_deepest_node(const struct rb_node *node)
 593 {
 594         for (;;) {
 595                 if (node->rb_left)
 596                         node = node->rb_left;
 597                 else if (node->rb_right)
 598                         node = node->rb_right;
 599                 else
 600                         return (struct rb_node *)node;
 601         }
 602 }
 603
 604 struct rb_node *rb_next_postorder(const struct rb_node *node)
 605 {
 606         const struct rb_node *parent;
 607         if (!node)
 608                 return NULL;
 609         parent = rb_parent(node);
 610
 611         /* If we're sitting on node, we've already seen our children */
 612         if (parent && node == parent->rb_left && parent->rb_right) {
 613                 /* If we are the parent's left node, go to the parent's right
 614                  * node then all the way down to the left */
 615                 return rb_left_deepest_node(parent->rb_right);
 616         } else
 617                 /* Otherwise we are the parent's right node, and the parent
 618                  * should be next */
 619                 return (struct rb_node *)parent;
 620 }
 621 EXPORT_SYMBOL(rb_next_postorder);
 622
 623 struct rb_node *rb_first_postorder(const struct rb_root *root)
 624 {
 625         if (!root->rb_node)
 626                 return NULL;
 627
 628         return rb_left_deepest_node(root->rb_node);
 629 }
 630 EXPORT_SYMBOL(rb_first_postorder);