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_right || rb_is_black(other->rb_right)) {
 287                                 if (!other->rb_left ||
 288                                     rb_is_black(other->rb_left)) {
 289                                         rb_set_red(other);
 290                                         node = parent;
 291                                         parent = rb_parent(node);
 292                                         continue;
 293                                 }
 294                                 rb_set_black(other->rb_left);
 295                                 rb_set_red(other);
 296                                 __rb_rotate_right(other, root);
 297                                 other = parent->rb_right;
 298                         }
 299                         rb_set_color(other, rb_color(parent));
 300                         rb_set_black(parent);
 301                         rb_set_black(other->rb_right);
 302                         __rb_rotate_left(parent, root);
 303                         break;
 304                 } else {
 305                         other = parent->rb_left;
 306                         if (rb_is_red(other))
 307                         {
 308                                 rb_set_black(other);
 309                                 rb_set_red(parent);
 310                                 __rb_rotate_right(parent, root);
 311                                 other = parent->rb_left;
 312                         }
 313                         if (!other->rb_left || rb_is_black(other->rb_left)) {
 314                                 if (!other->rb_right ||
 315                                     rb_is_black(other->rb_right)) {
 316                                         rb_set_red(other);
 317                                         node = parent;
 318                                         parent = rb_parent(node);
 319                                         continue;
 320                                 }
 321                                 rb_set_black(other->rb_right);
 322                                 rb_set_red(other);
 323                                 __rb_rotate_left(other, root);
 324                                 other = parent->rb_left;
 325                         }
 326                         rb_set_color(other, rb_color(parent));
 327                         rb_set_black(parent);
 328                         rb_set_black(other->rb_left);
 329                         __rb_rotate_right(parent, root);
 330                         break;
 331                 }
 332         }
 333 }
 334
 335 void rb_erase(struct rb_node *node, struct rb_root *root)
 336 {
 337         struct rb_node *child, *parent;
 338         int color;
 339
 340         if (!node->rb_left)
 341                 child = node->rb_right;
 342         else if (!node->rb_right)
 343                 child = node->rb_left;
 344         else
 345         {
 346                 struct rb_node *old = node, *left;
 347
 348                 node = node->rb_right;
 349                 while ((left = node->rb_left) != NULL)
 350                         node = left;
 351
 352                 if (rb_parent(old)) {
 353                         if (rb_parent(old)->rb_left == old)
 354                                 rb_parent(old)->rb_left = node;
 355                         else
 356                                 rb_parent(old)->rb_right = node;
 357                 } else
 358                         root->rb_node = node;
 359
 360                 child = node->rb_right;
 361                 parent = rb_parent(node);
 362                 color = rb_color(node);
 363
 364                 if (parent == old) {
 365                         parent = node;
 366                 } else {
 367                         if (child)
 368                                 rb_set_parent(child, parent);
 369                         parent->rb_left = child;
 370
 371                         node->rb_right = old->rb_right;
 372                         rb_set_parent(old->rb_right, node);
 373                 }
 374
 375                 node->__rb_parent_color = old->__rb_parent_color;
 376                 node->rb_left = old->rb_left;
 377                 rb_set_parent(old->rb_left, node);
 378
 379                 goto color;
 380         }
 381
 382         parent = rb_parent(node);
 383         color = rb_color(node);
 384
 385         if (child)
 386                 rb_set_parent(child, parent);
 387         if (parent)
 388         {
 389                 if (parent->rb_left == node)
 390                         parent->rb_left = child;
 391                 else
 392                         parent->rb_right = child;
 393         }
 394         else
 395                 root->rb_node = child;
 396
 397  color:
 398         if (color == RB_BLACK)
 399                 __rb_erase_color(child, parent, root);
 400 }
 401 EXPORT_SYMBOL(rb_erase);
 402
 403 static void rb_augment_path(struct rb_node *node, rb_augment_f func, void *data)
 404 {
 405         struct rb_node *parent;
 406
 407 up:
 408         func(node, data);
 409         parent = rb_parent(node);
 410         if (!parent)
 411                 return;
 412
 413         if (node == parent->rb_left && parent->rb_right)
 414                 func(parent->rb_right, data);
 415         else if (parent->rb_left)
 416                 func(parent->rb_left, data);
 417
 418         node = parent;
 419         goto up;
 420 }
 421
 422 /*
 423  * after inserting @node into the tree, update the tree to account for
 424  * both the new entry and any damage done by rebalance
 425  */
 426 void rb_augment_insert(struct rb_node *node, rb_augment_f func, void *data)
 427 {
 428         if (node->rb_left)
 429                 node = node->rb_left;
 430         else if (node->rb_right)
 431                 node = node->rb_right;
 432
 433         rb_augment_path(node, func, data);
 434 }
 435 EXPORT_SYMBOL(rb_augment_insert);
 436
 437 /*
 438  * before removing the node, find the deepest node on the rebalance path
 439  * that will still be there after @node gets removed
 440  */
 441 struct rb_node *rb_augment_erase_begin(struct rb_node *node)
 442 {
 443         struct rb_node *deepest;
 444
 445         if (!node->rb_right && !node->rb_left)
 446                 deepest = rb_parent(node);
 447         else if (!node->rb_right)
 448                 deepest = node->rb_left;
 449         else if (!node->rb_left)
 450                 deepest = node->rb_right;
 451         else {
 452                 deepest = rb_next(node);
 453                 if (deepest->rb_right)
 454                         deepest = deepest->rb_right;
 455                 else if (rb_parent(deepest) != node)
 456                         deepest = rb_parent(deepest);
 457         }
 458
 459         return deepest;
 460 }
 461 EXPORT_SYMBOL(rb_augment_erase_begin);
 462
 463 /*
 464  * after removal, update the tree to account for the removed entry
 465  * and any rebalance damage.
 466  */
 467 void rb_augment_erase_end(struct rb_node *node, rb_augment_f func, void *data)
 468 {
 469         if (node)
 470                 rb_augment_path(node, func, data);
 471 }
 472 EXPORT_SYMBOL(rb_augment_erase_end);
 473
 474 /*
 475  * This function returns the first node (in sort order) of the tree.
 476  */
 477 struct rb_node *rb_first(const struct rb_root *root)
 478 {
 479         struct rb_node  *n;
 480
 481         n = root->rb_node;
 482         if (!n)
 483                 return NULL;
 484         while (n->rb_left)
 485                 n = n->rb_left;
 486         return n;
 487 }
 488 EXPORT_SYMBOL(rb_first);
 489
 490 struct rb_node *rb_last(const struct rb_root *root)
 491 {
 492         struct rb_node  *n;
 493
 494         n = root->rb_node;
 495         if (!n)
 496                 return NULL;
 497         while (n->rb_right)
 498                 n = n->rb_right;
 499         return n;
 500 }
 501 EXPORT_SYMBOL(rb_last);
 502
 503 struct rb_node *rb_next(const struct rb_node *node)
 504 {
 505         struct rb_node *parent;
 506
 507         if (RB_EMPTY_NODE(node))
 508                 return NULL;
 509
 510         /* If we have a right-hand child, go down and then left as far
 511            as we can. */
 512         if (node->rb_right) {
 513                 node = node->rb_right;
 514                 while (node->rb_left)
 515                         node=node->rb_left;
 516                 return (struct rb_node *)node;
 517         }
 518
 519         /* No right-hand children.  Everything down and left is
 520            smaller than us, so any 'next' node must be in the general
 521            direction of our parent. Go up the tree; any time the
 522            ancestor is a right-hand child of its parent, keep going
 523            up. First time it's a left-hand child of its parent, said
 524            parent is our 'next' node. */
 525         while ((parent = rb_parent(node)) && node == parent->rb_right)
 526                 node = parent;
 527
 528         return parent;
 529 }
 530 EXPORT_SYMBOL(rb_next);
 531
 532 struct rb_node *rb_prev(const struct rb_node *node)
 533 {
 534         struct rb_node *parent;
 535
 536         if (RB_EMPTY_NODE(node))
 537                 return NULL;
 538
 539         /* If we have a left-hand child, go down and then right as far
 540            as we can. */
 541         if (node->rb_left) {
 542                 node = node->rb_left;
 543                 while (node->rb_right)
 544                         node=node->rb_right;
 545                 return (struct rb_node *)node;
 546         }
 547
 548         /* No left-hand children. Go up till we find an ancestor which
 549            is a right-hand child of its parent */
 550         while ((parent = rb_parent(node)) && node == parent->rb_left)
 551                 node = parent;
 552
 553         return parent;
 554 }
 555 EXPORT_SYMBOL(rb_prev);
 556
 557 void rb_replace_node(struct rb_node *victim, struct rb_node *new,
 558                      struct rb_root *root)
 559 {
 560         struct rb_node *parent = rb_parent(victim);
 561
 562         /* Set the surrounding nodes to point to the replacement */
 563         if (parent) {
 564                 if (victim == parent->rb_left)
 565                         parent->rb_left = new;
 566                 else
 567                         parent->rb_right = new;
 568         } else {
 569                 root->rb_node = new;
 570         }
 571         if (victim->rb_left)
 572                 rb_set_parent(victim->rb_left, new);
 573         if (victim->rb_right)
 574                 rb_set_parent(victim->rb_right, new);
 575
 576         /* Copy the pointers/colour from the victim to the replacement */
 577         *new = *victim;
 578 }
 579 EXPORT_SYMBOL(rb_replace_node);