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
   6   This program is free software; you can redistribute it and/or modify
   7   it under the terms of the GNU General Public License as published by
   8   the Free Software Foundation; either version 2 of the License, or
   9   (at your option) any later version.
  10
  11   This program is distributed in the hope that it will be useful,
  12   but WITHOUT ANY WARRANTY; without even the implied warranty of
  13   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  14   GNU General Public License for more details.
  15
  16   You should have received a copy of the GNU General Public License
  17   along with this program; if not, write to the Free Software
  18   Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
  19
  20   linux/lib/rbtree.c
  21 */
  22
  23 #include <linux/rbtree.h>
  24 #include <linux/export.h>
  25
  26 /*
  27  * red-black trees properties:  http://en.wikipedia.org/wiki/Rbtree
  28  *
  29  *  1) A node is either red or black
  30  *  2) The root is black
  31  *  3) All leaves (NULL) are black
  32  *  4) Both children of every red node are black
  33  *  5) Every simple path from root to leaves contains the same number
  34  *     of black nodes.
  35  *
  36  *  4 and 5 give the O(log n) guarantee, since 4 implies you cannot have two
  37  *  consecutive red nodes in a path and every red node is therefore followed by
  38  *  a black. So if B is the number of black nodes on every simple path (as per
  39  *  5), then the longest possible path due to 4 is 2B.
  40  *
  41  *  We shall indicate color with case, where black nodes are uppercase and red
  42  *  nodes will be lowercase. Unknown color nodes shall be drawn as red within
  43  *  parentheses and have some accompanying text comment.
  44  */
  45
  46 #define RB_RED          0
  47 #define RB_BLACK        1
  48
  49 #define rb_color(r)   ((r)->__rb_parent_color & 1)
  50 #define rb_is_red(r)   (!rb_color(r))
  51 #define rb_is_black(r) rb_color(r)
  52
  53 static inline void rb_set_parent(struct rb_node *rb, struct rb_node *p)
  54 {
  55         rb->__rb_parent_color = rb_color(rb) | (unsigned long)p;
  56 }
  57
  58 static inline void rb_set_parent_color(struct rb_node *rb,
  59                                        struct rb_node *p, int color)
  60 {
  61         rb->__rb_parent_color = (unsigned long)p | color;
  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         if (parent) {
  82                 if (parent->rb_left == old)
  83                         parent->rb_left = new;
  84                 else
  85                         parent->rb_right = new;
  86         } else
  87                 root->rb_node = new;
  88 }
  89
  90 void rb_insert_color(struct rb_node *node, struct rb_root *root)
  91 {
  92         struct rb_node *parent = rb_red_parent(node), *gparent, *tmp;
  93
  94         while (true) {
  95                 /*
  96                  * Loop invariant: node is red
  97                  *
  98                  * If there is a black parent, we are done.
  99                  * Otherwise, take some corrective action as we don't
 100                  * want a red root or two consecutive red nodes.
 101                  */
 102                 if (!parent) {
 103                         rb_set_parent_color(node, NULL, RB_BLACK);
 104                         break;
 105                 } else if (rb_is_black(parent))
 106                         break;
 107
 108                 gparent = rb_red_parent(parent);
 109
 110                 tmp = gparent->rb_right;
 111                 if (parent != tmp) {    /* parent == gparent->rb_left */
 112                         if (tmp && rb_is_red(tmp)) {
 113                                 /*
 114                                  * Case 1 - color flips
 115                                  *
 116                                  *       G            g
 117                                  *      / \          / \
 118                                  *     p   u  -->   P   U
 119                                  *    /            /
 120                                  *   n            N
 121                                  *
 122                                  * However, since g's parent might be red, and
 123                                  * 4) does not allow this, we need to recurse
 124                                  * at g.
 125                                  */
 126                                 rb_set_parent_color(tmp, gparent, RB_BLACK);
 127                                 rb_set_parent_color(parent, gparent, RB_BLACK);
 128                                 node = gparent;
 129                                 parent = rb_parent(node);
 130                                 rb_set_parent_color(node, parent, RB_RED);
 131                                 continue;
 132                         }
 133
 134                         tmp = parent->rb_right;
 135                         if (node == tmp) {
 136                                 /*
 137                                  * Case 2 - left rotate at parent
 138                                  *
 139                                  *      G             G
 140                                  *     / \           / \
 141                                  *    p   U  -->    n   U
 142                                  *     \           /
 143                                  *      n         p
 144                                  *
 145                                  * This still leaves us in violation of 4), the
 146                                  * continuation into Case 3 will fix that.
 147                                  */
 148                                 parent->rb_right = tmp = node->rb_left;
 149                                 node->rb_left = parent;
 150                                 if (tmp)
 151                                         rb_set_parent_color(tmp, parent,
 152                                                             RB_BLACK);
 153                                 rb_set_parent_color(parent, node, RB_RED);
 154                                 parent = node;
 155                                 tmp = node->rb_right;
 156                         }
 157
 158                         /*
 159                          * Case 3 - right rotate at gparent
 160                          *
 161                          *        G           P
 162                          *       / \         / \
 163                          *      p   U  -->  n   g
 164                          *     /                 \
 165                          *    n                   U
 166                          */
 167                         gparent->rb_left = tmp;  /* == parent->rb_right */
 168                         parent->rb_right = gparent;
 169                         if (tmp)
 170                                 rb_set_parent_color(tmp, gparent, RB_BLACK);
 171                         __rb_rotate_set_parents(gparent, parent, root, RB_RED);
 172                         break;
 173                 } else {
 174                         tmp = gparent->rb_left;
 175                         if (tmp && rb_is_red(tmp)) {
 176                                 /* Case 1 - color flips */
 177                                 rb_set_parent_color(tmp, gparent, RB_BLACK);
 178                                 rb_set_parent_color(parent, gparent, RB_BLACK);
 179                                 node = gparent;
 180                                 parent = rb_parent(node);
 181                                 rb_set_parent_color(node, parent, RB_RED);
 182                                 continue;
 183                         }
 184
 185                         tmp = parent->rb_left;
 186                         if (node == tmp) {
 187                                 /* Case 2 - right rotate at parent */
 188                                 parent->rb_left = tmp = node->rb_right;
 189                                 node->rb_right = parent;
 190                                 if (tmp)
 191                                         rb_set_parent_color(tmp, parent,
 192                                                             RB_BLACK);
 193                                 rb_set_parent_color(parent, node, RB_RED);
 194                                 parent = node;
 195                                 tmp = node->rb_left;
 196                         }
 197
 198                         /* Case 3 - left rotate at gparent */
 199                         gparent->rb_right = tmp;  /* == parent->rb_left */
 200                         parent->rb_left = gparent;
 201                         if (tmp)
 202                                 rb_set_parent_color(tmp, gparent, RB_BLACK);
 203                         __rb_rotate_set_parents(gparent, parent, root, RB_RED);
 204                         break;
 205                 }
 206         }
 207 }
 208 EXPORT_SYMBOL(rb_insert_color);
 209
 210 static void __rb_erase_color(struct rb_node *node, struct rb_node *parent,
 211                              struct rb_root *root)
 212 {
 213         struct rb_node *sibling, *tmp1, *tmp2;
 214
 215         while (true) {
 216                 /*
 217                  * Loop invariant: all leaf paths going through node have a
 218                  * black node count that is 1 lower than other leaf paths.
 219                  *
 220                  * If node is red, we can flip it to black to adjust.
 221                  * If node is the root, all leaf paths go through it.
 222                  * Otherwise, we need to adjust the tree through color flips
 223                  * and tree rotations as per one of the 4 cases below.
 224                  */
 225                 if (node && rb_is_red(node)) {
 226                         rb_set_parent_color(node, parent, RB_BLACK);
 227                         break;
 228                 } else if (!parent) {
 229                         break;
 230                 }
 231                 sibling = parent->rb_right;
 232                 if (node != sibling) {  /* node == parent->rb_left */
 233                         if (rb_is_red(sibling)) {
 234                                 /*
 235                                  * Case 1 - left rotate at parent
 236                                  *
 237                                  *     P               S
 238                                  *    / \             / \
 239                                  *   N   s    -->    p   Sr
 240                                  *      / \         / \
 241                                  *     Sl  Sr      N   Sl
 242                                  */
 243                                 parent->rb_right = tmp1 = sibling->rb_left;
 244                                 sibling->rb_left = parent;
 245                                 rb_set_parent_color(tmp1, parent, RB_BLACK);
 246                                 __rb_rotate_set_parents(parent, sibling, root,
 247                                                         RB_RED);
 248                                 sibling = tmp1;
 249                         }
 250                         tmp1 = sibling->rb_right;
 251                         if (!tmp1 || rb_is_black(tmp1)) {
 252                                 tmp2 = sibling->rb_left;
 253                                 if (!tmp2 || rb_is_black(tmp2)) {
 254                                         /*
 255                                          * Case 2 - sibling color flip
 256                                          * (p could be either color here)
 257                                          *
 258                                          *    (p)           (p)
 259                                          *    / \           / \
 260                                          *   N   S    -->  N   s
 261                                          *      / \           / \
 262                                          *     Sl  Sr        Sl  Sr
 263                                          *
 264                                          * This leaves us violating 5), so
 265                                          * recurse at p. If p is red, the
 266                                          * recursion will just flip it to black
 267                                          * and exit. If coming from Case 1,
 268                                          * p is known to be red.
 269                                          */
 270                                         rb_set_parent_color(sibling, parent,
 271                                                             RB_RED);
 272                                         node = parent;
 273                                         parent = rb_parent(node);
 274                                         continue;
 275                                 }
 276                                 /*
 277                                  * Case 3 - right rotate at sibling
 278                                  * (p could be either color here)
 279                                  *
 280                                  *   (p)           (p)
 281                                  *   / \           / \
 282                                  *  N   S    -->  N   Sl
 283                                  *     / \             \
 284                                  *    sl  Sr            s
 285                                  *                       \
 286                                  *                        Sr
 287                                  */
 288                                 sibling->rb_left = tmp1 = tmp2->rb_right;
 289                                 tmp2->rb_right = sibling;
 290                                 parent->rb_right = tmp2;
 291                                 if (tmp1)
 292                                         rb_set_parent_color(tmp1, sibling,
 293                                                             RB_BLACK);
 294                                 tmp1 = sibling;
 295                                 sibling = tmp2;
 296                         }
 297                         /*
 298                          * Case 4 - left rotate at parent + color flips
 299                          * (p and sl could be either color here.
 300                          *  After rotation, p becomes black, s acquires
 301                          *  p's color, and sl keeps its color)
 302                          *
 303                          *      (p)             (s)
 304                          *      / \             / \
 305                          *     N   S     -->   P   Sr
 306                          *        / \         / \
 307                          *      (sl) sr      N  (sl)
 308                          */
 309                         parent->rb_right = tmp2 = sibling->rb_left;
 310                         sibling->rb_left = parent;
 311                         rb_set_parent_color(tmp1, sibling, RB_BLACK);
 312                         if (tmp2)
 313                                 rb_set_parent(tmp2, parent);
 314                         __rb_rotate_set_parents(parent, sibling, root,
 315                                                 RB_BLACK);
 316                         break;
 317                 } else {
 318                         sibling = parent->rb_left;
 319                         if (rb_is_red(sibling)) {
 320                                 /* Case 1 - right rotate at parent */
 321                                 parent->rb_left = tmp1 = sibling->rb_right;
 322                                 sibling->rb_right = parent;
 323                                 rb_set_parent_color(tmp1, parent, RB_BLACK);
 324                                 __rb_rotate_set_parents(parent, sibling, root,
 325                                                         RB_RED);
 326                                 sibling = tmp1;
 327                         }
 328                         tmp1 = sibling->rb_left;
 329                         if (!tmp1 || rb_is_black(tmp1)) {
 330                                 tmp2 = sibling->rb_right;
 331                                 if (!tmp2 || rb_is_black(tmp2)) {
 332                                         /* Case 2 - sibling color flip */
 333                                         rb_set_parent_color(sibling, parent,
 334                                                             RB_RED);
 335                                         node = parent;
 336                                         parent = rb_parent(node);
 337                                         continue;
 338                                 }
 339                                 /* Case 3 - right rotate at sibling */
 340                                 sibling->rb_right = tmp1 = tmp2->rb_left;
 341                                 tmp2->rb_left = sibling;
 342                                 parent->rb_left = tmp2;
 343                                 if (tmp1)
 344                                         rb_set_parent_color(tmp1, sibling,
 345                                                             RB_BLACK);
 346                                 tmp1 = sibling;
 347                                 sibling = tmp2;
 348                         }
 349                         /* Case 4 - left rotate at parent + color flips */
 350                         parent->rb_left = tmp2 = sibling->rb_right;
 351                         sibling->rb_right = parent;
 352                         rb_set_parent_color(tmp1, sibling, RB_BLACK);
 353                         if (tmp2)
 354                                 rb_set_parent(tmp2, parent);
 355                         __rb_rotate_set_parents(parent, sibling, root,
 356                                                 RB_BLACK);
 357                         break;
 358                 }
 359         }
 360 }
 361
 362 void rb_erase(struct rb_node *node, struct rb_root *root)
 363 {
 364         struct rb_node *child, *parent;
 365         int color;
 366
 367         if (!node->rb_left)
 368                 child = node->rb_right;
 369         else if (!node->rb_right)
 370                 child = node->rb_left;
 371         else {
 372                 struct rb_node *old = node, *left;
 373
 374                 node = node->rb_right;
 375                 while ((left = node->rb_left) != NULL)
 376                         node = left;
 377
 378                 if (rb_parent(old)) {
 379                         if (rb_parent(old)->rb_left == old)
 380                                 rb_parent(old)->rb_left = node;
 381                         else
 382                                 rb_parent(old)->rb_right = node;
 383                 } else
 384                         root->rb_node = node;
 385
 386                 child = node->rb_right;
 387                 parent = rb_parent(node);
 388                 color = rb_color(node);
 389
 390                 if (parent == old) {
 391                         parent = node;
 392                 } else {
 393                         if (child)
 394                                 rb_set_parent(child, parent);
 395                         parent->rb_left = child;
 396
 397                         node->rb_right = old->rb_right;
 398                         rb_set_parent(old->rb_right, node);
 399                 }
 400
 401                 node->__rb_parent_color = old->__rb_parent_color;
 402                 node->rb_left = old->rb_left;
 403                 rb_set_parent(old->rb_left, node);
 404
 405                 goto color;
 406         }
 407
 408         parent = rb_parent(node);
 409         color = rb_color(node);
 410
 411         if (child)
 412                 rb_set_parent(child, parent);
 413         if (parent) {
 414                 if (parent->rb_left == node)
 415                         parent->rb_left = child;
 416                 else
 417                         parent->rb_right = child;
 418         } else
 419                 root->rb_node = child;
 420
 421 color:
 422         if (color == RB_BLACK)
 423                 __rb_erase_color(child, parent, root);
 424 }
 425 EXPORT_SYMBOL(rb_erase);
 426
 427 static void rb_augment_path(struct rb_node *node, rb_augment_f func, void *data)
 428 {
 429         struct rb_node *parent;
 430
 431 up:
 432         func(node, data);
 433         parent = rb_parent(node);
 434         if (!parent)
 435                 return;
 436
 437         if (node == parent->rb_left && parent->rb_right)
 438                 func(parent->rb_right, data);
 439         else if (parent->rb_left)
 440                 func(parent->rb_left, data);
 441
 442         node = parent;
 443         goto up;
 444 }
 445
 446 /*
 447  * after inserting @node into the tree, update the tree to account for
 448  * both the new entry and any damage done by rebalance
 449  */
 450 void rb_augment_insert(struct rb_node *node, rb_augment_f func, void *data)
 451 {
 452         if (node->rb_left)
 453                 node = node->rb_left;
 454         else if (node->rb_right)
 455                 node = node->rb_right;
 456
 457         rb_augment_path(node, func, data);
 458 }
 459 EXPORT_SYMBOL(rb_augment_insert);
 460
 461 /*
 462  * before removing the node, find the deepest node on the rebalance path
 463  * that will still be there after @node gets removed
 464  */
 465 struct rb_node *rb_augment_erase_begin(struct rb_node *node)
 466 {
 467         struct rb_node *deepest;
 468
 469         if (!node->rb_right && !node->rb_left)
 470                 deepest = rb_parent(node);
 471         else if (!node->rb_right)
 472                 deepest = node->rb_left;
 473         else if (!node->rb_left)
 474                 deepest = node->rb_right;
 475         else {
 476                 deepest = rb_next(node);
 477                 if (deepest->rb_right)
 478                         deepest = deepest->rb_right;
 479                 else if (rb_parent(deepest) != node)
 480                         deepest = rb_parent(deepest);
 481         }
 482
 483         return deepest;
 484 }
 485 EXPORT_SYMBOL(rb_augment_erase_begin);
 486
 487 /*
 488  * after removal, update the tree to account for the removed entry
 489  * and any rebalance damage.
 490  */
 491 void rb_augment_erase_end(struct rb_node *node, rb_augment_f func, void *data)
 492 {
 493         if (node)
 494                 rb_augment_path(node, func, data);
 495 }
 496 EXPORT_SYMBOL(rb_augment_erase_end);
 497
 498 /*
 499  * This function returns the first node (in sort order) of the tree.
 500  */
 501 struct rb_node *rb_first(const struct rb_root *root)
 502 {
 503         struct rb_node  *n;
 504
 505         n = root->rb_node;
 506         if (!n)
 507                 return NULL;
 508         while (n->rb_left)
 509                 n = n->rb_left;
 510         return n;
 511 }
 512 EXPORT_SYMBOL(rb_first);
 513
 514 struct rb_node *rb_last(const struct rb_root *root)
 515 {
 516         struct rb_node  *n;
 517
 518         n = root->rb_node;
 519         if (!n)
 520                 return NULL;
 521         while (n->rb_right)
 522                 n = n->rb_right;
 523         return n;
 524 }
 525 EXPORT_SYMBOL(rb_last);
 526
 527 struct rb_node *rb_next(const struct rb_node *node)
 528 {
 529         struct rb_node *parent;
 530
 531         if (RB_EMPTY_NODE(node))
 532                 return NULL;
 533
 534         /*
 535          * If we have a right-hand child, go down and then left as far
 536          * as we can.
 537          */
 538         if (node->rb_right) {
 539                 node = node->rb_right;
 540                 while (node->rb_left)
 541                         node=node->rb_left;
 542                 return (struct rb_node *)node;
 543         }
 544
 545         /*
 546          * No right-hand children. Everything down and left is smaller than us,
 547          * so any 'next' node must be in the general direction of our parent.
 548          * Go up the tree; any time the ancestor is a right-hand child of its
 549          * parent, keep going up. First time it's a left-hand child of its
 550          * parent, said parent is our 'next' node.
 551          */
 552         while ((parent = rb_parent(node)) && node == parent->rb_right)
 553                 node = parent;
 554
 555         return parent;
 556 }
 557 EXPORT_SYMBOL(rb_next);
 558
 559 struct rb_node *rb_prev(const struct rb_node *node)
 560 {
 561         struct rb_node *parent;
 562
 563         if (RB_EMPTY_NODE(node))
 564                 return NULL;
 565
 566         /*
 567          * If we have a left-hand child, go down and then right as far
 568          * as we can.
 569          */
 570         if (node->rb_left) {
 571                 node = node->rb_left;
 572                 while (node->rb_right)
 573                         node=node->rb_right;
 574                 return (struct rb_node *)node;
 575         }
 576
 577         /*
 578          * No left-hand children. Go up till we find an ancestor which
 579          * is a right-hand child of its parent.
 580          */
 581         while ((parent = rb_parent(node)) && node == parent->rb_left)
 582                 node = parent;
 583
 584         return parent;
 585 }
 586 EXPORT_SYMBOL(rb_prev);
 587
 588 void rb_replace_node(struct rb_node *victim, struct rb_node *new,
 589                      struct rb_root *root)
 590 {
 591         struct rb_node *parent = rb_parent(victim);
 592
 593         /* Set the surrounding nodes to point to the replacement */
 594         if (parent) {
 595                 if (victim == parent->rb_left)
 596                         parent->rb_left = new;
 597                 else
 598                         parent->rb_right = new;
 599         } else {
 600                 root->rb_node = new;
 601         }
 602         if (victim->rb_left)
 603                 rb_set_parent(victim->rb_left, new);
 604         if (victim->rb_right)
 605                 rb_set_parent(victim->rb_right, new);
 606
 607         /* Copy the pointers/colour from the victim to the replacement */
 608         *new = *victim;
 609 }
 610 EXPORT_SYMBOL(rb_replace_node);