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 static inline void
  70 __rb_change_child(struct rb_node *old, struct rb_node *new,
  71                   struct rb_node *parent, struct rb_root *root)
  72 {
  73         if (parent) {
  74                 if (parent->rb_left == old)
  75                         parent->rb_left = new;
  76                 else
  77                         parent->rb_right = new;
  78         } else
  79                 root->rb_node = new;
  80 }
  81
  82 /*
  83  * Helper function for rotations:
  84  * - old's parent and color get assigned to new
  85  * - old gets assigned new as a parent and 'color' as a color.
  86  */
  87 static inline void
  88 __rb_rotate_set_parents(struct rb_node *old, struct rb_node *new,
  89                         struct rb_root *root, int color)
  90 {
  91         struct rb_node *parent = rb_parent(old);
  92         new->__rb_parent_color = old->__rb_parent_color;
  93         rb_set_parent_color(old, new, color);
  94         __rb_change_child(old, new, parent, root);
  95 }
  96
  97 void rb_insert_color(struct rb_node *node, struct rb_root *root)
  98 {
  99         struct rb_node *parent = rb_red_parent(node), *gparent, *tmp;
 100
 101         while (true) {
 102                 /*
 103                  * Loop invariant: node is red
 104                  *
 105                  * If there is a black parent, we are done.
 106                  * Otherwise, take some corrective action as we don't
 107                  * want a red root or two consecutive red nodes.
 108                  */
 109                 if (!parent) {
 110                         rb_set_parent_color(node, NULL, RB_BLACK);
 111                         break;
 112                 } else if (rb_is_black(parent))
 113                         break;
 114
 115                 gparent = rb_red_parent(parent);
 116
 117                 tmp = gparent->rb_right;
 118                 if (parent != tmp) {    /* parent == gparent->rb_left */
 119                         if (tmp && rb_is_red(tmp)) {
 120                                 /*
 121                                  * Case 1 - color flips
 122                                  *
 123                                  *       G            g
 124                                  *      / \          / \
 125                                  *     p   u  -->   P   U
 126                                  *    /            /
 127                                  *   n            N
 128                                  *
 129                                  * However, since g's parent might be red, and
 130                                  * 4) does not allow this, we need to recurse
 131                                  * at g.
 132                                  */
 133                                 rb_set_parent_color(tmp, gparent, RB_BLACK);
 134                                 rb_set_parent_color(parent, gparent, RB_BLACK);
 135                                 node = gparent;
 136                                 parent = rb_parent(node);
 137                                 rb_set_parent_color(node, parent, RB_RED);
 138                                 continue;
 139                         }
 140
 141                         tmp = parent->rb_right;
 142                         if (node == tmp) {
 143                                 /*
 144                                  * Case 2 - left rotate at parent
 145                                  *
 146                                  *      G             G
 147                                  *     / \           / \
 148                                  *    p   U  -->    n   U
 149                                  *     \           /
 150                                  *      n         p
 151                                  *
 152                                  * This still leaves us in violation of 4), the
 153                                  * continuation into Case 3 will fix that.
 154                                  */
 155                                 parent->rb_right = tmp = node->rb_left;
 156                                 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                                 parent = node;
 162                                 tmp = node->rb_right;
 163                         }
 164
 165                         /*
 166                          * Case 3 - right rotate at gparent
 167                          *
 168                          *        G           P
 169                          *       / \         / \
 170                          *      p   U  -->  n   g
 171                          *     /                 \
 172                          *    n                   U
 173                          */
 174                         gparent->rb_left = tmp;  /* == parent->rb_right */
 175                         parent->rb_right = gparent;
 176                         if (tmp)
 177                                 rb_set_parent_color(tmp, gparent, RB_BLACK);
 178                         __rb_rotate_set_parents(gparent, parent, root, RB_RED);
 179                         break;
 180                 } else {
 181                         tmp = gparent->rb_left;
 182                         if (tmp && rb_is_red(tmp)) {
 183                                 /* Case 1 - color flips */
 184                                 rb_set_parent_color(tmp, gparent, RB_BLACK);
 185                                 rb_set_parent_color(parent, gparent, RB_BLACK);
 186                                 node = gparent;
 187                                 parent = rb_parent(node);
 188                                 rb_set_parent_color(node, parent, RB_RED);
 189                                 continue;
 190                         }
 191
 192                         tmp = parent->rb_left;
 193                         if (node == tmp) {
 194                                 /* Case 2 - right rotate at parent */
 195                                 parent->rb_left = tmp = node->rb_right;
 196                                 node->rb_right = parent;
 197                                 if (tmp)
 198                                         rb_set_parent_color(tmp, parent,
 199                                                             RB_BLACK);
 200                                 rb_set_parent_color(parent, node, RB_RED);
 201                                 parent = node;
 202                                 tmp = node->rb_left;
 203                         }
 204
 205                         /* Case 3 - left rotate at gparent */
 206                         gparent->rb_right = tmp;  /* == parent->rb_left */
 207                         parent->rb_left = gparent;
 208                         if (tmp)
 209                                 rb_set_parent_color(tmp, gparent, RB_BLACK);
 210                         __rb_rotate_set_parents(gparent, parent, root, RB_RED);
 211                         break;
 212                 }
 213         }
 214 }
 215 EXPORT_SYMBOL(rb_insert_color);
 216
 217 static void __rb_erase_color(struct rb_node *node, struct rb_node *parent,
 218                              struct rb_root *root)
 219 {
 220         struct rb_node *sibling, *tmp1, *tmp2;
 221
 222         while (true) {
 223                 /*
 224                  * Loop invariant: all leaf paths going through node have a
 225                  * black node count that is 1 lower than other leaf paths.
 226                  *
 227                  * If node is red, we can flip it to black to adjust.
 228                  * If node is the root, all leaf paths go through it.
 229                  * Otherwise, we need to adjust the tree through color flips
 230                  * and tree rotations as per one of the 4 cases below.
 231                  */
 232                 if (node && rb_is_red(node)) {
 233                         rb_set_parent_color(node, parent, RB_BLACK);
 234                         break;
 235                 } else if (!parent) {
 236                         break;
 237                 }
 238                 sibling = parent->rb_right;
 239                 if (node != sibling) {  /* node == parent->rb_left */
 240                         if (rb_is_red(sibling)) {
 241                                 /*
 242                                  * Case 1 - left rotate at parent
 243                                  *
 244                                  *     P               S
 245                                  *    / \             / \
 246                                  *   N   s    -->    p   Sr
 247                                  *      / \         / \
 248                                  *     Sl  Sr      N   Sl
 249                                  */
 250                                 parent->rb_right = tmp1 = sibling->rb_left;
 251                                 sibling->rb_left = parent;
 252                                 rb_set_parent_color(tmp1, parent, RB_BLACK);
 253                                 __rb_rotate_set_parents(parent, sibling, root,
 254                                                         RB_RED);
 255                                 sibling = tmp1;
 256                         }
 257                         tmp1 = sibling->rb_right;
 258                         if (!tmp1 || rb_is_black(tmp1)) {
 259                                 tmp2 = sibling->rb_left;
 260                                 if (!tmp2 || rb_is_black(tmp2)) {
 261                                         /*
 262                                          * Case 2 - sibling color flip
 263                                          * (p could be either color here)
 264                                          *
 265                                          *    (p)           (p)
 266                                          *    / \           / \
 267                                          *   N   S    -->  N   s
 268                                          *      / \           / \
 269                                          *     Sl  Sr        Sl  Sr
 270                                          *
 271                                          * This leaves us violating 5), so
 272                                          * recurse at p. If p is red, the
 273                                          * recursion will just flip it to black
 274                                          * and exit. If coming from Case 1,
 275                                          * p is known to be red.
 276                                          */
 277                                         rb_set_parent_color(sibling, parent,
 278                                                             RB_RED);
 279                                         node = parent;
 280                                         parent = rb_parent(node);
 281                                         continue;
 282                                 }
 283                                 /*
 284                                  * Case 3 - right rotate at sibling
 285                                  * (p could be either color here)
 286                                  *
 287                                  *   (p)           (p)
 288                                  *   / \           / \
 289                                  *  N   S    -->  N   Sl
 290                                  *     / \             \
 291                                  *    sl  Sr            s
 292                                  *                       \
 293                                  *                        Sr
 294                                  */
 295                                 sibling->rb_left = tmp1 = tmp2->rb_right;
 296                                 tmp2->rb_right = sibling;
 297                                 parent->rb_right = tmp2;
 298                                 if (tmp1)
 299                                         rb_set_parent_color(tmp1, sibling,
 300                                                             RB_BLACK);
 301                                 tmp1 = sibling;
 302                                 sibling = tmp2;
 303                         }
 304                         /*
 305                          * Case 4 - left rotate at parent + color flips
 306                          * (p and sl could be either color here.
 307                          *  After rotation, p becomes black, s acquires
 308                          *  p's color, and sl keeps its color)
 309                          *
 310                          *      (p)             (s)
 311                          *      / \             / \
 312                          *     N   S     -->   P   Sr
 313                          *        / \         / \
 314                          *      (sl) sr      N  (sl)
 315                          */
 316                         parent->rb_right = tmp2 = sibling->rb_left;
 317                         sibling->rb_left = parent;
 318                         rb_set_parent_color(tmp1, sibling, RB_BLACK);
 319                         if (tmp2)
 320                                 rb_set_parent(tmp2, parent);
 321                         __rb_rotate_set_parents(parent, sibling, root,
 322                                                 RB_BLACK);
 323                         break;
 324                 } else {
 325                         sibling = parent->rb_left;
 326                         if (rb_is_red(sibling)) {
 327                                 /* Case 1 - right rotate at parent */
 328                                 parent->rb_left = tmp1 = sibling->rb_right;
 329                                 sibling->rb_right = parent;
 330                                 rb_set_parent_color(tmp1, parent, RB_BLACK);
 331                                 __rb_rotate_set_parents(parent, sibling, root,
 332                                                         RB_RED);
 333                                 sibling = tmp1;
 334                         }
 335                         tmp1 = sibling->rb_left;
 336                         if (!tmp1 || rb_is_black(tmp1)) {
 337                                 tmp2 = sibling->rb_right;
 338                                 if (!tmp2 || rb_is_black(tmp2)) {
 339                                         /* Case 2 - sibling color flip */
 340                                         rb_set_parent_color(sibling, parent,
 341                                                             RB_RED);
 342                                         node = parent;
 343                                         parent = rb_parent(node);
 344                                         continue;
 345                                 }
 346                                 /* Case 3 - right rotate at sibling */
 347                                 sibling->rb_right = tmp1 = tmp2->rb_left;
 348                                 tmp2->rb_left = sibling;
 349                                 parent->rb_left = tmp2;
 350                                 if (tmp1)
 351                                         rb_set_parent_color(tmp1, sibling,
 352                                                             RB_BLACK);
 353                                 tmp1 = sibling;
 354                                 sibling = tmp2;
 355                         }
 356                         /* Case 4 - left rotate at parent + color flips */
 357                         parent->rb_left = tmp2 = sibling->rb_right;
 358                         sibling->rb_right = parent;
 359                         rb_set_parent_color(tmp1, sibling, RB_BLACK);
 360                         if (tmp2)
 361                                 rb_set_parent(tmp2, parent);
 362                         __rb_rotate_set_parents(parent, sibling, root,
 363                                                 RB_BLACK);
 364                         break;
 365                 }
 366         }
 367 }
 368
 369 void rb_erase(struct rb_node *node, struct rb_root *root)
 370 {
 371         struct rb_node *child = node->rb_right, *tmp = node->rb_left;
 372         struct rb_node *parent;
 373         int color;
 374
 375         if (!tmp) {
 376         case1:
 377                 /* Case 1: node to erase has no more than 1 child (easy!) */
 378
 379                 parent = rb_parent(node);
 380                 color = rb_color(node);
 381
 382                 if (child)
 383                         rb_set_parent(child, parent);
 384                 __rb_change_child(node, child, parent, root);
 385         } else if (!child) {
 386                 /* Still case 1, but this time the child is node->rb_left */
 387                 child = tmp;
 388                 goto case1;
 389         } else {
 390                 struct rb_node *old = node, *left;
 391
 392                 node = child;
 393                 while ((left = node->rb_left) != NULL)
 394                         node = left;
 395
 396                 __rb_change_child(old, node, rb_parent(old), root);
 397
 398                 child = node->rb_right;
 399                 parent = rb_parent(node);
 400                 color = rb_color(node);
 401
 402                 if (parent == old) {
 403                         parent = node;
 404                 } else {
 405                         if (child)
 406                                 rb_set_parent(child, parent);
 407                         parent->rb_left = child;
 408
 409                         node->rb_right = old->rb_right;
 410                         rb_set_parent(old->rb_right, node);
 411                 }
 412
 413                 node->__rb_parent_color = old->__rb_parent_color;
 414                 node->rb_left = old->rb_left;
 415                 rb_set_parent(old->rb_left, node);
 416         }
 417
 418         if (color == RB_BLACK)
 419                 __rb_erase_color(child, parent, root);
 420 }
 421 EXPORT_SYMBOL(rb_erase);
 422
 423 static void rb_augment_path(struct rb_node *node, rb_augment_f func, void *data)
 424 {
 425         struct rb_node *parent;
 426
 427 up:
 428         func(node, data);
 429         parent = rb_parent(node);
 430         if (!parent)
 431                 return;
 432
 433         if (node == parent->rb_left && parent->rb_right)
 434                 func(parent->rb_right, data);
 435         else if (parent->rb_left)
 436                 func(parent->rb_left, data);
 437
 438         node = parent;
 439         goto up;
 440 }
 441
 442 /*
 443  * after inserting @node into the tree, update the tree to account for
 444  * both the new entry and any damage done by rebalance
 445  */
 446 void rb_augment_insert(struct rb_node *node, rb_augment_f func, void *data)
 447 {
 448         if (node->rb_left)
 449                 node = node->rb_left;
 450         else if (node->rb_right)
 451                 node = node->rb_right;
 452
 453         rb_augment_path(node, func, data);
 454 }
 455 EXPORT_SYMBOL(rb_augment_insert);
 456
 457 /*
 458  * before removing the node, find the deepest node on the rebalance path
 459  * that will still be there after @node gets removed
 460  */
 461 struct rb_node *rb_augment_erase_begin(struct rb_node *node)
 462 {
 463         struct rb_node *deepest;
 464
 465         if (!node->rb_right && !node->rb_left)
 466                 deepest = rb_parent(node);
 467         else if (!node->rb_right)
 468                 deepest = node->rb_left;
 469         else if (!node->rb_left)
 470                 deepest = node->rb_right;
 471         else {
 472                 deepest = rb_next(node);
 473                 if (deepest->rb_right)
 474                         deepest = deepest->rb_right;
 475                 else if (rb_parent(deepest) != node)
 476                         deepest = rb_parent(deepest);
 477         }
 478
 479         return deepest;
 480 }
 481 EXPORT_SYMBOL(rb_augment_erase_begin);
 482
 483 /*
 484  * after removal, update the tree to account for the removed entry
 485  * and any rebalance damage.
 486  */
 487 void rb_augment_erase_end(struct rb_node *node, rb_augment_f func, void *data)
 488 {
 489         if (node)
 490                 rb_augment_path(node, func, data);
 491 }
 492 EXPORT_SYMBOL(rb_augment_erase_end);
 493
 494 /*
 495  * This function returns the first node (in sort order) of the tree.
 496  */
 497 struct rb_node *rb_first(const struct rb_root *root)
 498 {
 499         struct rb_node  *n;
 500
 501         n = root->rb_node;
 502         if (!n)
 503                 return NULL;
 504         while (n->rb_left)
 505                 n = n->rb_left;
 506         return n;
 507 }
 508 EXPORT_SYMBOL(rb_first);
 509
 510 struct rb_node *rb_last(const struct rb_root *root)
 511 {
 512         struct rb_node  *n;
 513
 514         n = root->rb_node;
 515         if (!n)
 516                 return NULL;
 517         while (n->rb_right)
 518                 n = n->rb_right;
 519         return n;
 520 }
 521 EXPORT_SYMBOL(rb_last);
 522
 523 struct rb_node *rb_next(const struct rb_node *node)
 524 {
 525         struct rb_node *parent;
 526
 527         if (RB_EMPTY_NODE(node))
 528                 return NULL;
 529
 530         /*
 531          * If we have a right-hand child, go down and then left as far
 532          * as we can.
 533          */
 534         if (node->rb_right) {
 535                 node = node->rb_right;
 536                 while (node->rb_left)
 537                         node=node->rb_left;
 538                 return (struct rb_node *)node;
 539         }
 540
 541         /*
 542          * No right-hand children. Everything down and left is smaller than us,
 543          * so any 'next' node must be in the general direction of our parent.
 544          * Go up the tree; any time the ancestor is a right-hand child of its
 545          * parent, keep going up. First time it's a left-hand child of its
 546          * parent, said parent is our 'next' node.
 547          */
 548         while ((parent = rb_parent(node)) && node == parent->rb_right)
 549                 node = parent;
 550
 551         return parent;
 552 }
 553 EXPORT_SYMBOL(rb_next);
 554
 555 struct rb_node *rb_prev(const struct rb_node *node)
 556 {
 557         struct rb_node *parent;
 558
 559         if (RB_EMPTY_NODE(node))
 560                 return NULL;
 561
 562         /*
 563          * If we have a left-hand child, go down and then right as far
 564          * as we can.
 565          */
 566         if (node->rb_left) {
 567                 node = node->rb_left;
 568                 while (node->rb_right)
 569                         node=node->rb_right;
 570                 return (struct rb_node *)node;
 571         }
 572
 573         /*
 574          * No left-hand children. Go up till we find an ancestor which
 575          * is a right-hand child of its parent.
 576          */
 577         while ((parent = rb_parent(node)) && node == parent->rb_left)
 578                 node = parent;
 579
 580         return parent;
 581 }
 582 EXPORT_SYMBOL(rb_prev);
 583
 584 void rb_replace_node(struct rb_node *victim, struct rb_node *new,
 585                      struct rb_root *root)
 586 {
 587         struct rb_node *parent = rb_parent(victim);
 588
 589         /* Set the surrounding nodes to point to the replacement */
 590         __rb_change_child(victim, new, parent, root);
 591         if (victim->rb_left)
 592                 rb_set_parent(victim->rb_left, new);
 593         if (victim->rb_right)
 594                 rb_set_parent(victim->rb_right, new);
 595
 596         /* Copy the pointers/colour from the victim to the replacement */
 597         *new = *victim;
 598 }
 599 EXPORT_SYMBOL(rb_replace_node);