LXR linux/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  (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  linux/lib/rbtree.c
  22*/
  23
  24#include <linux/rbtree_augmented.h>
  25#include <linux/export.h>
  26
  27/*
  28 * red-black trees properties:  http://en.wikipedia.org/wiki/Rbtree
  29 *
  30 *  1) A node is either red or black
  31 *  2) The root is black
  32 *  3) All leaves (NULL) are black
  33 *  4) Both children of every red node are black
  34 *  5) Every simple path from root to leaves contains the same number
  35 *     of black nodes.
  36 *
  37 *  4 and 5 give the O(log n) guarantee, since 4 implies you cannot have two
  38 *  consecutive red nodes in a path and every red node is therefore followed by
  39 *  a black. So if B is the number of black nodes on every simple path (as per
  40 *  5), then the longest possible path due to 4 is 2B.
  41 *
  42 *  We shall indicate color with case, where black nodes are uppercase and red
  43 *  nodes will be lowercase. Unknown color nodes shall be drawn as red within
  44 *  parentheses and have some accompanying text comment.
  45 */
  46
  47/*
  48 * Notes on lockless lookups:
  49 *
  50 * All stores to the tree structure (rb_left and rb_right) must be done using
  51 * WRITE_ONCE(). And we must not inadvertently cause (temporary) loops in the
  52 * tree structure as seen in program order.
  53 *
  54 * These two requirements will allow lockless iteration of the tree -- not
  55 * correct iteration mind you, tree rotations are not atomic so a lookup might
  56 * miss entire subtrees.
  57 *
  58 * But they do guarantee that any such traversal will only see valid elements
  59 * and that it will indeed complete -- does not get stuck in a loop.
  60 *
  61 * It also guarantees that if the lookup returns an element it is the 'correct'
  62 * one. But not returning an element does _NOT_ mean it's not present.
  63 *
  64 * NOTE:
  65 *
  66 * Stores to __rb_parent_color are not important for simple lookups so those
  67 * are left undone as of now. Nor did I check for loops involving parent
  68 * pointers.
  69 */
  70
  71static inline void rb_set_black(struct rb_node *rb)
  72{
  73        rb->__rb_parent_color |= RB_BLACK;
  74}
  75
  76static inline struct rb_node *rb_red_parent(struct rb_node *red)
  77{
  78        return (struct rb_node *)red->__rb_parent_color;
  79}
  80
  81/*
  82 * Helper function for rotations:
  83 * - old's parent and color get assigned to new
  84 * - old gets assigned new as a parent and 'color' as a color.
  85 */
  86static inline void
  87__rb_rotate_set_parents(struct rb_node *old, struct rb_node *new,
  88                        struct rb_root *root, int color)
  89{
  90        struct rb_node *parent = rb_parent(old);
  91        new->__rb_parent_color = old->__rb_parent_color;
  92        rb_set_parent_color(old, new, color);
  93        __rb_change_child(old, new, parent, root);
  94}
  95
  96static __always_inline void
  97__rb_insert(struct rb_node *node, struct rb_root *root,
  98            void (*augment_rotate)(struct rb_node *old, struct rb_node *new))
  99{
 100        struct rb_node *parent = rb_red_parent(node), *gparent, *tmp;
 101
 102        while (true) {
 103                /*
 104                 * Loop invariant: node is red
 105                 *
 106                 * If there is a black parent, we are done.
 107                 * Otherwise, take some corrective action as we don't
 108                 * want a red root or two consecutive red nodes.
 109                 */
 110                if (!parent) {
 111                        rb_set_parent_color(node, NULL, RB_BLACK);
 112                        break;
 113                } else if (rb_is_black(parent))
 114                        break;
 115
 116                gparent = rb_red_parent(parent);
 117
 118                tmp = gparent->rb_right;
 119                if (parent != tmp) {    /* parent == gparent->rb_left */
 120                        if (tmp && rb_is_red(tmp)) {
 121                                /*
 122                                 * Case 1 - color flips
 123                                 *
 124                                 *       G            g
 125                                 *      / \          / \
 126                                 *     p   u  -->   P   U
 127                                 *    /            /
 128                                 *   n            n
 129                                 *
 130                                 * However, since g's parent might be red, and
 131                                 * 4) does not allow this, we need to recurse
 132                                 * at g.
 133                                 */
 134                                rb_set_parent_color(tmp, gparent, RB_BLACK);
 135                                rb_set_parent_color(parent, gparent, RB_BLACK);
 136                                node = gparent;
 137                                parent = rb_parent(node);
 138                                rb_set_parent_color(node, parent, RB_RED);
 139                                continue;
 140                        }
 141
 142                        tmp = parent->rb_right;
 143                        if (node == tmp) {
 144                                /*
 145                                 * Case 2 - left rotate at parent
 146                                 *
 147                                 *      G             G
 148                                 *     / \           / \
 149                                 *    p   U  -->    n   U
 150                                 *     \           /
 151                                 *      n         p
 152                                 *
 153                                 * This still leaves us in violation of 4), the
 154                                 * continuation into Case 3 will fix that.
 155                                 */
 156                                tmp = node->rb_left;
 157                                WRITE_ONCE(parent->rb_right, tmp);
 158                                WRITE_ONCE(node->rb_left, parent);
 159                                if (tmp)
 160                                        rb_set_parent_color(tmp, parent,
 161                                                            RB_BLACK);
 162                                rb_set_parent_color(parent, node, RB_RED);
 163                                augment_rotate(parent, node);
 164                                parent = node;
 165                                tmp = node->rb_right;
 166                        }
 167
 168                        /*
 169                         * Case 3 - right rotate at gparent
 170                         *
 171                         *        G           P
 172                         *       / \         / \
 173                         *      p   U  -->  n   g
 174                         *     /                 \
 175                         *    n                   U
 176                         */
 177                        WRITE_ONCE(gparent->rb_left, tmp); /* == parent->rb_right */
 178                        WRITE_ONCE(parent->rb_right, gparent);
 179                        if (tmp)
 180                                rb_set_parent_color(tmp, gparent, RB_BLACK);
 181                        __rb_rotate_set_parents(gparent, parent, root, RB_RED);
 182                        augment_rotate(gparent, parent);
 183                        break;
 184                } else {
 185                        tmp = gparent->rb_left;
 186                        if (tmp && rb_is_red(tmp)) {
 187                                /* Case 1 - color flips */
 188                                rb_set_parent_color(tmp, gparent, RB_BLACK);
 189                                rb_set_parent_color(parent, gparent, RB_BLACK);
 190                                node = gparent;
 191                                parent = rb_parent(node);
 192                                rb_set_parent_color(node, parent, RB_RED);
 193                                continue;
 194                        }
 195
 196                        tmp = parent->rb_left;
 197                        if (node == tmp) {
 198                                /* Case 2 - right rotate at parent */
 199                                tmp = node->rb_right;
 200                                WRITE_ONCE(parent->rb_left, tmp);
 201                                WRITE_ONCE(node->rb_right, parent);
 202                                if (tmp)
 203                                        rb_set_parent_color(tmp, parent,
 204                                                            RB_BLACK);
 205                                rb_set_parent_color(parent, node, RB_RED);
 206                                augment_rotate(parent, node);
 207                                parent = node;
 208                                tmp = node->rb_left;
 209                        }
 210
 211                        /* Case 3 - left rotate at gparent */
 212                        WRITE_ONCE(gparent->rb_right, tmp); /* == parent->rb_left */
 213                        WRITE_ONCE(parent->rb_left, gparent);
 214                        if (tmp)
 215                                rb_set_parent_color(tmp, gparent, RB_BLACK);
 216                        __rb_rotate_set_parents(gparent, parent, root, RB_RED);
 217                        augment_rotate(gparent, parent);
 218                        break;
 219                }
 220        }
 221}
 222
 223/*
 224 * Inline version for rb_erase() use - we want to be able to inline
 225 * and eliminate the dummy_rotate callback there
 226 */
 227static __always_inline void
 228____rb_erase_color(struct rb_node *parent, struct rb_root *root,
 229        void (*augment_rotate)(struct rb_node *old, struct rb_node *new))
 230{
 231        struct rb_node *node = NULL, *sibling, *tmp1, *tmp2;
 232
 233        while (true) {
 234                /*
 235                 * Loop invariants:
 236                 * - node is black (or NULL on first iteration)
 237                 * - node is not the root (parent is not NULL)
 238                 * - All leaf paths going through parent and node have a
 239                 *   black node count that is 1 lower than other leaf paths.
 240                 */
 241                sibling = parent->rb_right;
 242                if (node != sibling) {  /* node == parent->rb_left */
 243                        if (rb_is_red(sibling)) {
 244                                /*
 245                                 * Case 1 - left rotate at parent
 246                                 *
 247                                 *     P               S
 248                                 *    / \             / \
 249                                 *   N   s    -->    p   Sr
 250                                 *      / \         / \
 251                                 *     Sl  Sr      N   Sl
 252                                 */
 253                                tmp1 = sibling->rb_left;
 254                                WRITE_ONCE(parent->rb_right, tmp1);
 255                                WRITE_ONCE(sibling->rb_left, parent);
 256                                rb_set_parent_color(tmp1, parent, RB_BLACK);
 257                                __rb_rotate_set_parents(parent, sibling, root,
 258                                                        RB_RED);
 259                                augment_rotate(parent, sibling);
 260                                sibling = tmp1;
 261                        }
 262                        tmp1 = sibling->rb_right;
 263                        if (!tmp1 || rb_is_black(tmp1)) {
 264                                tmp2 = sibling->rb_left;
 265                                if (!tmp2 || rb_is_black(tmp2)) {
 266                                        /*
 267                                         * Case 2 - sibling color flip
 268                                         * (p could be either color here)
 269                                         *
 270                                         *    (p)           (p)
 271                                         *    / \           / \
 272                                         *   N   S    -->  N   s
 273                                         *      / \           / \
 274                                         *     Sl  Sr        Sl  Sr
 275                                         *
 276                                         * This leaves us violating 5) which
 277                                         * can be fixed by flipping p to black
 278                                         * if it was red, or by recursing at p.
 279                                         * p is red when coming from Case 1.
 280                                         */
 281                                        rb_set_parent_color(sibling, parent,
 282                                                            RB_RED);
 283                                        if (rb_is_red(parent))
 284                                                rb_set_black(parent);
 285                                        else {
 286                                                node = parent;
 287                                                parent = rb_parent(node);
 288                                                if (parent)
 289                                                        continue;
 290                                        }
 291                                        break;
 292                                }
 293                                /*
 294                                 * Case 3 - right rotate at sibling
 295                                 * (p could be either color here)
 296                                 *
 297                                 *   (p)           (p)
 298                                 *   / \           / \
 299                                 *  N   S    -->  N   Sl
 300                                 *     / \             \
 301                                 *    sl  Sr            s
 302                                 *                       \
 303                                 *                        Sr
 304                                 */
 305                                tmp1 = tmp2->rb_right;
 306                                WRITE_ONCE(sibling->rb_left, tmp1);
 307                                WRITE_ONCE(tmp2->rb_right, sibling);
 308                                WRITE_ONCE(parent->rb_right, tmp2);
 309                                if (tmp1)
 310                                        rb_set_parent_color(tmp1, sibling,
 311                                                            RB_BLACK);
 312                                augment_rotate(sibling, tmp2);
 313                                tmp1 = sibling;
 314                                sibling = tmp2;
 315                        }
 316                        /*
 317                         * Case 4 - left rotate at parent + color flips
 318                         * (p and sl could be either color here.
 319                         *  After rotation, p becomes black, s acquires
 320                         *  p's color, and sl keeps its color)
 321                         *
 322                         *      (p)             (s)
 323                         *      / \             / \
 324                         *     N   S     -->   P   Sr
 325                         *        / \         / \
 326                         *      (sl) sr      N  (sl)
 327                         */
 328                        tmp2 = sibling->rb_left;
 329                        WRITE_ONCE(parent->rb_right, tmp2);
 330                        WRITE_ONCE(sibling->rb_left, parent);
 331                        rb_set_parent_color(tmp1, sibling, RB_BLACK);
 332                        if (tmp2)
 333                                rb_set_parent(tmp2, parent);
 334                        __rb_rotate_set_parents(parent, sibling, root,
 335                                                RB_BLACK);
 336                        augment_rotate(parent, sibling);
 337                        break;
 338                } else {
 339                        sibling = parent->rb_left;
 340                        if (rb_is_red(sibling)) {
 341                                /* Case 1 - right rotate at parent */
 342                                tmp1 = sibling->rb_right;
 343                                WRITE_ONCE(parent->rb_left, tmp1);
 344                                WRITE_ONCE(sibling->rb_right, parent);
 345                                rb_set_parent_color(tmp1, parent, RB_BLACK);
 346                                __rb_rotate_set_parents(parent, sibling, root,
 347                                                        RB_RED);
 348                                augment_rotate(parent, sibling);
 349                                sibling = tmp1;
 350                        }
 351                        tmp1 = sibling->rb_left;
 352                        if (!tmp1 || rb_is_black(tmp1)) {
 353                                tmp2 = sibling->rb_right;
 354                                if (!tmp2 || rb_is_black(tmp2)) {
 355                                        /* Case 2 - sibling color flip */
 356                                        rb_set_parent_color(sibling, parent,
 357                                                            RB_RED);
 358                                        if (rb_is_red(parent))
 359                                                rb_set_black(parent);
 360                                        else {
 361                                                node = parent;
 362                                                parent = rb_parent(node);
 363                                                if (parent)
 364                                                        continue;
 365                                        }
 366                                        break;
 367                                }
 368                                /* Case 3 - right rotate at sibling */
 369                                tmp1 = tmp2->rb_left;
 370                                WRITE_ONCE(sibling->rb_right, tmp1);
 371                                WRITE_ONCE(tmp2->rb_left, sibling);
 372                                WRITE_ONCE(parent->rb_left, tmp2);
 373                                if (tmp1)
 374                                        rb_set_parent_color(tmp1, sibling,
 375                                                            RB_BLACK);
 376                                augment_rotate(sibling, tmp2);
 377                                tmp1 = sibling;
 378                                sibling = tmp2;
 379                        }
 380                        /* Case 4 - left rotate at parent + color flips */
 381                        tmp2 = sibling->rb_right;
 382                        WRITE_ONCE(parent->rb_left, tmp2);
 383                        WRITE_ONCE(sibling->rb_right, parent);
 384                        rb_set_parent_color(tmp1, sibling, RB_BLACK);
 385                        if (tmp2)
 386                                rb_set_parent(tmp2, parent);
 387                        __rb_rotate_set_parents(parent, sibling, root,
 388                                                RB_BLACK);
 389                        augment_rotate(parent, sibling);
 390                        break;
 391                }
 392        }
 393}
 394
 395/* Non-inline version for rb_erase_augmented() use */
 396void __rb_erase_color(struct rb_node *parent, struct rb_root *root,
 397        void (*augment_rotate)(struct rb_node *old, struct rb_node *new))
 398{
 399        ____rb_erase_color(parent, root, augment_rotate);
 400}
 401EXPORT_SYMBOL(__rb_erase_color);
 402
 403/*
 404 * Non-augmented rbtree manipulation functions.
 405 *
 406 * We use dummy augmented callbacks here, and have the compiler optimize them
 407 * out of the rb_insert_color() and rb_erase() function definitions.
 408 */
 409
 410static inline void dummy_propagate(struct rb_node *node, struct rb_node *stop) {}
 411static inline void dummy_copy(struct rb_node *old, struct rb_node *new) {}
 412static inline void dummy_rotate(struct rb_node *old, struct rb_node *new) {}
 413
 414static const struct rb_augment_callbacks dummy_callbacks = {
 415        dummy_propagate, dummy_copy, dummy_rotate
 416};
 417
 418void rb_insert_color(struct rb_node *node, struct rb_root *root)
 419{
 420        __rb_insert(node, root, dummy_rotate);
 421}
 422EXPORT_SYMBOL(rb_insert_color);
 423
 424void rb_erase(struct rb_node *node, struct rb_root *root)
 425{
 426        struct rb_node *rebalance;
 427        rebalance = __rb_erase_augmented(node, root, &dummy_callbacks);
 428        if (rebalance)
 429                ____rb_erase_color(rebalance, root, dummy_rotate);
 430}
 431EXPORT_SYMBOL(rb_erase);
 432
 433/*
 434 * Augmented rbtree manipulation functions.
 435 *
 436 * This instantiates the same __always_inline functions as in the non-augmented
 437 * case, but this time with user-defined callbacks.
 438 */
 439
 440void __rb_insert_augmented(struct rb_node *node, struct rb_root *root,
 441        void (*augment_rotate)(struct rb_node *old, struct rb_node *new))
 442{
 443        __rb_insert(node, root, augment_rotate);
 444}
 445EXPORT_SYMBOL(__rb_insert_augmented);
 446
 447/*
 448 * This function returns the first node (in sort order) of the tree.
 449 */
 450struct rb_node *rb_first(const struct rb_root *root)
 451{
 452        struct rb_node  *n;
 453
 454        n = root->rb_node;
 455        if (!n)
 456                return NULL;
 457        while (n->rb_left)
 458                n = n->rb_left;
 459        return n;
 460}
 461EXPORT_SYMBOL(rb_first);
 462
 463struct rb_node *rb_last(const struct rb_root *root)
 464{
 465        struct rb_node  *n;
 466
 467        n = root->rb_node;
 468        if (!n)
 469                return NULL;
 470        while (n->rb_right)
 471                n = n->rb_right;
 472        return n;
 473}
 474EXPORT_SYMBOL(rb_last);
 475
 476struct rb_node *rb_next(const struct rb_node *node)
 477{
 478        struct rb_node *parent;
 479
 480        if (RB_EMPTY_NODE(node))
 481                return NULL;
 482
 483        /*
 484         * If we have a right-hand child, go down and then left as far
 485         * as we can.
 486         */
 487        if (node->rb_right) {
 488                node = node->rb_right; 
 489                while (node->rb_left)
 490                        node=node->rb_left;
 491                return (struct rb_node *)node;
 492        }
 493
 494        /*
 495         * No right-hand children. Everything down and left is smaller than us,
 496         * so any 'next' node must be in the general direction of our parent.
 497         * Go up the tree; any time the ancestor is a right-hand child of its
 498         * parent, keep going up. First time it's a left-hand child of its
 499         * parent, said parent is our 'next' node.
 500         */
 501        while ((parent = rb_parent(node)) && node == parent->rb_right)
 502                node = parent;
 503
 504        return parent;
 505}
 506EXPORT_SYMBOL(rb_next);
 507
 508struct rb_node *rb_prev(const struct rb_node *node)
 509{
 510        struct rb_node *parent;
 511
 512        if (RB_EMPTY_NODE(node))
 513                return NULL;
 514
 515        /*
 516         * If we have a left-hand child, go down and then right as far
 517         * as we can.
 518         */
 519        if (node->rb_left) {
 520                node = node->rb_left; 
 521                while (node->rb_right)
 522                        node=node->rb_right;
 523                return (struct rb_node *)node;
 524        }
 525
 526        /*
 527         * No left-hand children. Go up till we find an ancestor which
 528         * is a right-hand child of its parent.
 529         */
 530        while ((parent = rb_parent(node)) && node == parent->rb_left)
 531                node = parent;
 532
 533        return parent;
 534}
 535EXPORT_SYMBOL(rb_prev);
 536
 537void rb_replace_node(struct rb_node *victim, struct rb_node *new,
 538                     struct rb_root *root)
 539{
 540        struct rb_node *parent = rb_parent(victim);
 541
 542        /* Set the surrounding nodes to point to the replacement */
 543        __rb_change_child(victim, new, parent, root);
 544        if (victim->rb_left)
 545                rb_set_parent(victim->rb_left, new);
 546        if (victim->rb_right)
 547                rb_set_parent(victim->rb_right, new);
 548
 549        /* Copy the pointers/colour from the victim to the replacement */
 550        *new = *victim;
 551}
 552EXPORT_SYMBOL(rb_replace_node);
 553
 554static struct rb_node *rb_left_deepest_node(const struct rb_node *node)
 555{
 556        for (;;) {
 557                if (node->rb_left)
 558                        node = node->rb_left;
 559                else if (node->rb_right)
 560                        node = node->rb_right;
 561                else
 562                        return (struct rb_node *)node;
 563        }
 564}
 565
 566struct rb_node *rb_next_postorder(const struct rb_node *node)
 567{
 568        const struct rb_node *parent;
 569        if (!node)
 570                return NULL;
 571        parent = rb_parent(node);
 572
 573        /* If we're sitting on node, we've already seen our children */
 574        if (parent && node == parent->rb_left && parent->rb_right) {
 575                /* If we are the parent's left node, go to the parent's right
 576                 * node then all the way down to the left */
 577                return rb_left_deepest_node(parent->rb_right);
 578        } else
 579                /* Otherwise we are the parent's right node, and the parent
 580                 * should be next */
 581                return (struct rb_node *)parent;
 582}
 583EXPORT_SYMBOL(rb_next_postorder);
 584
 585struct rb_node *rb_first_postorder(const struct rb_root *root)
 586{
 587        if (!root->rb_node)
 588                return NULL;
 589
 590        return rb_left_deepest_node(root->rb_node);
 591}
 592EXPORT_SYMBOL(rb_first_postorder);
 593