lib/rbtree.c

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