lib/rbtree.c

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