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