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