wimlib.net Git - wimlib/blob - src/rbtree.c

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