tools/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
 416 /*
 417  * Non-augmented rbtree manipulation functions.
 418  *
 419  * We use dummy augmented callbacks here, and have the compiler optimize them
 420  * out of the rb_insert_color() and rb_erase() function definitions.
 421  */
 422
 423 static inline void dummy_propagate(struct rb_node *node, struct rb_node *stop) {}
 424 static inline void dummy_copy(struct rb_node *old, struct rb_node *new) {}
 425 static inline void dummy_rotate(struct rb_node *old, struct rb_node *new) {}
 426
 427 static const struct rb_augment_callbacks dummy_callbacks = {
 428         .propagate = dummy_propagate,
 429         .copy = dummy_copy,
 430         .rotate = dummy_rotate
 431 };
 432
 433 void rb_insert_color(struct rb_node *node, struct rb_root *root)
 434 {
 435         __rb_insert(node, root, dummy_rotate);
 436 }
 437
 438 void rb_erase(struct rb_node *node, struct rb_root *root)
 439 {
 440         struct rb_node *rebalance;
 441         rebalance = __rb_erase_augmented(node, root, &dummy_callbacks);
 442         if (rebalance)
 443                 ____rb_erase_color(rebalance, root, dummy_rotate);
 444 }
 445
 446 /*
 447  * Augmented rbtree manipulation functions.
 448  *
 449  * This instantiates the same __always_inline functions as in the non-augmented
 450  * case, but this time with user-defined callbacks.
 451  */
 452
 453 void __rb_insert_augmented(struct rb_node *node, struct rb_root *root,
 454         void (*augment_rotate)(struct rb_node *old, struct rb_node *new))
 455 {
 456         __rb_insert(node, root, augment_rotate);
 457 }
 458
 459 /*
 460  * This function returns the first node (in sort order) of the tree.
 461  */
 462 struct rb_node *rb_first(const struct rb_root *root)
 463 {
 464         struct rb_node  *n;
 465
 466         n = root->rb_node;
 467         if (!n)
 468                 return NULL;
 469         while (n->rb_left)
 470                 n = n->rb_left;
 471         return n;
 472 }
 473
 474 struct rb_node *rb_last(const struct rb_root *root)
 475 {
 476         struct rb_node  *n;
 477
 478         n = root->rb_node;
 479         if (!n)
 480                 return NULL;
 481         while (n->rb_right)
 482                 n = n->rb_right;
 483         return n;
 484 }
 485
 486 struct rb_node *rb_next(const struct rb_node *node)
 487 {
 488         struct rb_node *parent;
 489
 490         if (RB_EMPTY_NODE(node))
 491                 return NULL;
 492
 493         /*
 494          * If we have a right-hand child, go down and then left as far
 495          * as we can.
 496          */
 497         if (node->rb_right) {
 498                 node = node->rb_right;
 499                 while (node->rb_left)
 500                         node = node->rb_left;
 501                 return (struct rb_node *)node;
 502         }
 503
 504         /*
 505          * No right-hand children. Everything down and left is smaller than us,
 506          * so any 'next' node must be in the general direction of our parent.
 507          * Go up the tree; any time the ancestor is a right-hand child of its
 508          * parent, keep going up. First time it's a left-hand child of its
 509          * parent, said parent is our 'next' node.
 510          */
 511         while ((parent = rb_parent(node)) && node == parent->rb_right)
 512                 node = parent;
 513
 514         return parent;
 515 }
 516
 517 struct rb_node *rb_prev(const struct rb_node *node)
 518 {
 519         struct rb_node *parent;
 520
 521         if (RB_EMPTY_NODE(node))
 522                 return NULL;
 523
 524         /*
 525          * If we have a left-hand child, go down and then right as far
 526          * as we can.
 527          */
 528         if (node->rb_left) {
 529                 node = node->rb_left;
 530                 while (node->rb_right)
 531                         node = node->rb_right;
 532                 return (struct rb_node *)node;
 533         }
 534
 535         /*
 536          * No left-hand children. Go up till we find an ancestor which
 537          * is a right-hand child of its parent.
 538          */
 539         while ((parent = rb_parent(node)) && node == parent->rb_left)
 540                 node = parent;
 541
 542         return parent;
 543 }
 544
 545 void rb_replace_node(struct rb_node *victim, struct rb_node *new,
 546                      struct rb_root *root)
 547 {
 548         struct rb_node *parent = rb_parent(victim);
 549
 550         /* Copy the pointers/colour from the victim to the replacement */
 551         *new = *victim;
 552
 553         /* Set the surrounding nodes to point to the replacement */
 554         if (victim->rb_left)
 555                 rb_set_parent(victim->rb_left, new);
 556         if (victim->rb_right)
 557                 rb_set_parent(victim->rb_right, new);
 558         __rb_change_child(victim, new, parent, root);
 559 }
 560
 561 static struct rb_node *rb_left_deepest_node(const struct rb_node *node)
 562 {
 563         for (;;) {
 564                 if (node->rb_left)
 565                         node = node->rb_left;
 566                 else if (node->rb_right)
 567                         node = node->rb_right;
 568                 else
 569                         return (struct rb_node *)node;
 570         }
 571 }
 572
 573 struct rb_node *rb_next_postorder(const struct rb_node *node)
 574 {
 575         const struct rb_node *parent;
 576         if (!node)
 577                 return NULL;
 578         parent = rb_parent(node);
 579
 580         /* If we're sitting on node, we've already seen our children */
 581         if (parent && node == parent->rb_left && parent->rb_right) {
 582                 /* If we are the parent's left node, go to the parent's right
 583                  * node then all the way down to the left */
 584                 return rb_left_deepest_node(parent->rb_right);
 585         } else
 586                 /* Otherwise we are the parent's right node, and the parent
 587                  * should be next */
 588                 return (struct rb_node *)parent;
 589 }
 590
 591 struct rb_node *rb_first_postorder(const struct rb_root *root)
 592 {
 593         if (!root->rb_node)
 594                 return NULL;
 595
 596         return rb_left_deepest_node(root->rb_node);
 597 }