Back to home page

LXR

 
 

    


0001 /*
0002   Red Black Trees
0003   (C) 1999  Andrea Arcangeli <andrea@suse.de>
0004   (C) 2002  David Woodhouse <dwmw2@infradead.org>
0005   (C) 2012  Michel Lespinasse <walken@google.com>
0006 
0007   This program is free software; you can redistribute it and/or modify
0008   it under the terms of the GNU General Public License as published by
0009   the Free Software Foundation; either version 2 of the License, or
0010   (at your option) any later version.
0011 
0012   This program is distributed in the hope that it will be useful,
0013   but WITHOUT ANY WARRANTY; without even the implied warranty of
0014   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
0015   GNU General Public License for more details.
0016 
0017   You should have received a copy of the GNU General Public License
0018   along with this program; if not, write to the Free Software
0019   Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
0020 
0021   linux/lib/rbtree.c
0022 */
0023 
0024 #include <linux/rbtree_augmented.h>
0025 #include <linux/export.h>
0026 
0027 /*
0028  * red-black trees properties:  http://en.wikipedia.org/wiki/Rbtree
0029  *
0030  *  1) A node is either red or black
0031  *  2) The root is black
0032  *  3) All leaves (NULL) are black
0033  *  4) Both children of every red node are black
0034  *  5) Every simple path from root to leaves contains the same number
0035  *     of black nodes.
0036  *
0037  *  4 and 5 give the O(log n) guarantee, since 4 implies you cannot have two
0038  *  consecutive red nodes in a path and every red node is therefore followed by
0039  *  a black. So if B is the number of black nodes on every simple path (as per
0040  *  5), then the longest possible path due to 4 is 2B.
0041  *
0042  *  We shall indicate color with case, where black nodes are uppercase and red
0043  *  nodes will be lowercase. Unknown color nodes shall be drawn as red within
0044  *  parentheses and have some accompanying text comment.
0045  */
0046 
0047 /*
0048  * Notes on lockless lookups:
0049  *
0050  * All stores to the tree structure (rb_left and rb_right) must be done using
0051  * WRITE_ONCE(). And we must not inadvertently cause (temporary) loops in the
0052  * tree structure as seen in program order.
0053  *
0054  * These two requirements will allow lockless iteration of the tree -- not
0055  * correct iteration mind you, tree rotations are not atomic so a lookup might
0056  * miss entire subtrees.
0057  *
0058  * But they do guarantee that any such traversal will only see valid elements
0059  * and that it will indeed complete -- does not get stuck in a loop.
0060  *
0061  * It also guarantees that if the lookup returns an element it is the 'correct'
0062  * one. But not returning an element does _NOT_ mean it's not present.
0063  *
0064  * NOTE:
0065  *
0066  * Stores to __rb_parent_color are not important for simple lookups so those
0067  * are left undone as of now. Nor did I check for loops involving parent
0068  * pointers.
0069  */
0070 
0071 static inline void rb_set_black(struct rb_node *rb)
0072 {
0073     rb->__rb_parent_color |= RB_BLACK;
0074 }
0075 
0076 static inline struct rb_node *rb_red_parent(struct rb_node *red)
0077 {
0078     return (struct rb_node *)red->__rb_parent_color;
0079 }
0080 
0081 /*
0082  * Helper function for rotations:
0083  * - old's parent and color get assigned to new
0084  * - old gets assigned new as a parent and 'color' as a color.
0085  */
0086 static inline void
0087 __rb_rotate_set_parents(struct rb_node *old, struct rb_node *new,
0088             struct rb_root *root, int color)
0089 {
0090     struct rb_node *parent = rb_parent(old);
0091     new->__rb_parent_color = old->__rb_parent_color;
0092     rb_set_parent_color(old, new, color);
0093     __rb_change_child(old, new, parent, root);
0094 }
0095 
0096 static __always_inline void
0097 __rb_insert(struct rb_node *node, struct rb_root *root,
0098         void (*augment_rotate)(struct rb_node *old, struct rb_node *new))
0099 {
0100     struct rb_node *parent = rb_red_parent(node), *gparent, *tmp;
0101 
0102     while (true) {
0103         /*
0104          * Loop invariant: node is red
0105          *
0106          * If there is a black parent, we are done.
0107          * Otherwise, take some corrective action as we don't
0108          * want a red root or two consecutive red nodes.
0109          */
0110         if (!parent) {
0111             rb_set_parent_color(node, NULL, RB_BLACK);
0112             break;
0113         } else if (rb_is_black(parent))
0114             break;
0115 
0116         gparent = rb_red_parent(parent);
0117 
0118         tmp = gparent->rb_right;
0119         if (parent != tmp) {    /* parent == gparent->rb_left */
0120             if (tmp && rb_is_red(tmp)) {
0121                 /*
0122                  * Case 1 - color flips
0123                  *
0124                  *       G            g
0125                  *      / \          / \
0126                  *     p   u  -->   P   U
0127                  *    /            /
0128                  *   n            n
0129                  *
0130                  * However, since g's parent might be red, and
0131                  * 4) does not allow this, we need to recurse
0132                  * at g.
0133                  */
0134                 rb_set_parent_color(tmp, gparent, RB_BLACK);
0135                 rb_set_parent_color(parent, gparent, RB_BLACK);
0136                 node = gparent;
0137                 parent = rb_parent(node);
0138                 rb_set_parent_color(node, parent, RB_RED);
0139                 continue;
0140             }
0141 
0142             tmp = parent->rb_right;
0143             if (node == tmp) {
0144                 /*
0145                  * Case 2 - left rotate at parent
0146                  *
0147                  *      G             G
0148                  *     / \           / \
0149                  *    p   U  -->    n   U
0150                  *     \           /
0151                  *      n         p
0152                  *
0153                  * This still leaves us in violation of 4), the
0154                  * continuation into Case 3 will fix that.
0155                  */
0156                 tmp = node->rb_left;
0157                 WRITE_ONCE(parent->rb_right, tmp);
0158                 WRITE_ONCE(node->rb_left, parent);
0159                 if (tmp)
0160                     rb_set_parent_color(tmp, parent,
0161                                 RB_BLACK);
0162                 rb_set_parent_color(parent, node, RB_RED);
0163                 augment_rotate(parent, node);
0164                 parent = node;
0165                 tmp = node->rb_right;
0166             }
0167 
0168             /*
0169              * Case 3 - right rotate at gparent
0170              *
0171              *        G           P
0172              *       / \         / \
0173              *      p   U  -->  n   g
0174              *     /                 \
0175              *    n                   U
0176              */
0177             WRITE_ONCE(gparent->rb_left, tmp); /* == parent->rb_right */
0178             WRITE_ONCE(parent->rb_right, gparent);
0179             if (tmp)
0180                 rb_set_parent_color(tmp, gparent, RB_BLACK);
0181             __rb_rotate_set_parents(gparent, parent, root, RB_RED);
0182             augment_rotate(gparent, parent);
0183             break;
0184         } else {
0185             tmp = gparent->rb_left;
0186             if (tmp && rb_is_red(tmp)) {
0187                 /* Case 1 - color flips */
0188                 rb_set_parent_color(tmp, gparent, RB_BLACK);
0189                 rb_set_parent_color(parent, gparent, RB_BLACK);
0190                 node = gparent;
0191                 parent = rb_parent(node);
0192                 rb_set_parent_color(node, parent, RB_RED);
0193                 continue;
0194             }
0195 
0196             tmp = parent->rb_left;
0197             if (node == tmp) {
0198                 /* Case 2 - right rotate at parent */
0199                 tmp = node->rb_right;
0200                 WRITE_ONCE(parent->rb_left, tmp);
0201                 WRITE_ONCE(node->rb_right, parent);
0202                 if (tmp)
0203                     rb_set_parent_color(tmp, parent,
0204                                 RB_BLACK);
0205                 rb_set_parent_color(parent, node, RB_RED);
0206                 augment_rotate(parent, node);
0207                 parent = node;
0208                 tmp = node->rb_left;
0209             }
0210 
0211             /* Case 3 - left rotate at gparent */
0212             WRITE_ONCE(gparent->rb_right, tmp); /* == parent->rb_left */
0213             WRITE_ONCE(parent->rb_left, gparent);
0214             if (tmp)
0215                 rb_set_parent_color(tmp, gparent, RB_BLACK);
0216             __rb_rotate_set_parents(gparent, parent, root, RB_RED);
0217             augment_rotate(gparent, parent);
0218             break;
0219         }
0220     }
0221 }
0222 
0223 /*
0224  * Inline version for rb_erase() use - we want to be able to inline
0225  * and eliminate the dummy_rotate callback there
0226  */
0227 static __always_inline void
0228 ____rb_erase_color(struct rb_node *parent, struct rb_root *root,
0229     void (*augment_rotate)(struct rb_node *old, struct rb_node *new))
0230 {
0231     struct rb_node *node = NULL, *sibling, *tmp1, *tmp2;
0232 
0233     while (true) {
0234         /*
0235          * Loop invariants:
0236          * - node is black (or NULL on first iteration)
0237          * - node is not the root (parent is not NULL)
0238          * - All leaf paths going through parent and node have a
0239          *   black node count that is 1 lower than other leaf paths.
0240          */
0241         sibling = parent->rb_right;
0242         if (node != sibling) {  /* node == parent->rb_left */
0243             if (rb_is_red(sibling)) {
0244                 /*
0245                  * Case 1 - left rotate at parent
0246                  *
0247                  *     P               S
0248                  *    / \             / \
0249                  *   N   s    -->    p   Sr
0250                  *      / \         / \
0251                  *     Sl  Sr      N   Sl
0252                  */
0253                 tmp1 = sibling->rb_left;
0254                 WRITE_ONCE(parent->rb_right, tmp1);
0255                 WRITE_ONCE(sibling->rb_left, parent);
0256                 rb_set_parent_color(tmp1, parent, RB_BLACK);
0257                 __rb_rotate_set_parents(parent, sibling, root,
0258                             RB_RED);
0259                 augment_rotate(parent, sibling);
0260                 sibling = tmp1;
0261             }
0262             tmp1 = sibling->rb_right;
0263             if (!tmp1 || rb_is_black(tmp1)) {
0264                 tmp2 = sibling->rb_left;
0265                 if (!tmp2 || rb_is_black(tmp2)) {
0266                     /*
0267                      * Case 2 - sibling color flip
0268                      * (p could be either color here)
0269                      *
0270                      *    (p)           (p)
0271                      *    / \           / \
0272                      *   N   S    -->  N   s
0273                      *      / \           / \
0274                      *     Sl  Sr        Sl  Sr
0275                      *
0276                      * This leaves us violating 5) which
0277                      * can be fixed by flipping p to black
0278                      * if it was red, or by recursing at p.
0279                      * p is red when coming from Case 1.
0280                      */
0281                     rb_set_parent_color(sibling, parent,
0282                                 RB_RED);
0283                     if (rb_is_red(parent))
0284                         rb_set_black(parent);
0285                     else {
0286                         node = parent;
0287                         parent = rb_parent(node);
0288                         if (parent)
0289                             continue;
0290                     }
0291                     break;
0292                 }
0293                 /*
0294                  * Case 3 - right rotate at sibling
0295                  * (p could be either color here)
0296                  *
0297                  *   (p)           (p)
0298                  *   / \           / \
0299                  *  N   S    -->  N   sl
0300                  *     / \             \
0301                  *    sl  Sr            S
0302                  *                       \
0303                  *                        Sr
0304                  *
0305                  * Note: p might be red, and then both
0306                  * p and sl are red after rotation(which
0307                  * breaks property 4). This is fixed in
0308                  * Case 4 (in __rb_rotate_set_parents()
0309                  *         which set sl the color of p
0310                  *         and set p RB_BLACK)
0311                  *
0312                  *   (p)            (sl)
0313                  *   / \            /  \
0314                  *  N   sl   -->   P    S
0315                  *       \        /      \
0316                  *        S      N        Sr
0317                  *         \
0318                  *          Sr
0319                  */
0320                 tmp1 = tmp2->rb_right;
0321                 WRITE_ONCE(sibling->rb_left, tmp1);
0322                 WRITE_ONCE(tmp2->rb_right, sibling);
0323                 WRITE_ONCE(parent->rb_right, tmp2);
0324                 if (tmp1)
0325                     rb_set_parent_color(tmp1, sibling,
0326                                 RB_BLACK);
0327                 augment_rotate(sibling, tmp2);
0328                 tmp1 = sibling;
0329                 sibling = tmp2;
0330             }
0331             /*
0332              * Case 4 - left rotate at parent + color flips
0333              * (p and sl could be either color here.
0334              *  After rotation, p becomes black, s acquires
0335              *  p's color, and sl keeps its color)
0336              *
0337              *      (p)             (s)
0338              *      / \             / \
0339              *     N   S     -->   P   Sr
0340              *        / \         / \
0341              *      (sl) sr      N  (sl)
0342              */
0343             tmp2 = sibling->rb_left;
0344             WRITE_ONCE(parent->rb_right, tmp2);
0345             WRITE_ONCE(sibling->rb_left, parent);
0346             rb_set_parent_color(tmp1, sibling, RB_BLACK);
0347             if (tmp2)
0348                 rb_set_parent(tmp2, parent);
0349             __rb_rotate_set_parents(parent, sibling, root,
0350                         RB_BLACK);
0351             augment_rotate(parent, sibling);
0352             break;
0353         } else {
0354             sibling = parent->rb_left;
0355             if (rb_is_red(sibling)) {
0356                 /* Case 1 - right rotate at parent */
0357                 tmp1 = sibling->rb_right;
0358                 WRITE_ONCE(parent->rb_left, tmp1);
0359                 WRITE_ONCE(sibling->rb_right, parent);
0360                 rb_set_parent_color(tmp1, parent, RB_BLACK);
0361                 __rb_rotate_set_parents(parent, sibling, root,
0362                             RB_RED);
0363                 augment_rotate(parent, sibling);
0364                 sibling = tmp1;
0365             }
0366             tmp1 = sibling->rb_left;
0367             if (!tmp1 || rb_is_black(tmp1)) {
0368                 tmp2 = sibling->rb_right;
0369                 if (!tmp2 || rb_is_black(tmp2)) {
0370                     /* Case 2 - sibling color flip */
0371                     rb_set_parent_color(sibling, parent,
0372                                 RB_RED);
0373                     if (rb_is_red(parent))
0374                         rb_set_black(parent);
0375                     else {
0376                         node = parent;
0377                         parent = rb_parent(node);
0378                         if (parent)
0379                             continue;
0380                     }
0381                     break;
0382                 }
0383                 /* Case 3 - left rotate at sibling */
0384                 tmp1 = tmp2->rb_left;
0385                 WRITE_ONCE(sibling->rb_right, tmp1);
0386                 WRITE_ONCE(tmp2->rb_left, sibling);
0387                 WRITE_ONCE(parent->rb_left, tmp2);
0388                 if (tmp1)
0389                     rb_set_parent_color(tmp1, sibling,
0390                                 RB_BLACK);
0391                 augment_rotate(sibling, tmp2);
0392                 tmp1 = sibling;
0393                 sibling = tmp2;
0394             }
0395             /* Case 4 - right rotate at parent + color flips */
0396             tmp2 = sibling->rb_right;
0397             WRITE_ONCE(parent->rb_left, tmp2);
0398             WRITE_ONCE(sibling->rb_right, parent);
0399             rb_set_parent_color(tmp1, sibling, RB_BLACK);
0400             if (tmp2)
0401                 rb_set_parent(tmp2, parent);
0402             __rb_rotate_set_parents(parent, sibling, root,
0403                         RB_BLACK);
0404             augment_rotate(parent, sibling);
0405             break;
0406         }
0407     }
0408 }
0409 
0410 /* Non-inline version for rb_erase_augmented() use */
0411 void __rb_erase_color(struct rb_node *parent, struct rb_root *root,
0412     void (*augment_rotate)(struct rb_node *old, struct rb_node *new))
0413 {
0414     ____rb_erase_color(parent, root, augment_rotate);
0415 }
0416 EXPORT_SYMBOL(__rb_erase_color);
0417 
0418 /*
0419  * Non-augmented rbtree manipulation functions.
0420  *
0421  * We use dummy augmented callbacks here, and have the compiler optimize them
0422  * out of the rb_insert_color() and rb_erase() function definitions.
0423  */
0424 
0425 static inline void dummy_propagate(struct rb_node *node, struct rb_node *stop) {}
0426 static inline void dummy_copy(struct rb_node *old, struct rb_node *new) {}
0427 static inline void dummy_rotate(struct rb_node *old, struct rb_node *new) {}
0428 
0429 static const struct rb_augment_callbacks dummy_callbacks = {
0430     dummy_propagate, dummy_copy, dummy_rotate
0431 };
0432 
0433 void rb_insert_color(struct rb_node *node, struct rb_root *root)
0434 {
0435     __rb_insert(node, root, dummy_rotate);
0436 }
0437 EXPORT_SYMBOL(rb_insert_color);
0438 
0439 void rb_erase(struct rb_node *node, struct rb_root *root)
0440 {
0441     struct rb_node *rebalance;
0442     rebalance = __rb_erase_augmented(node, root, &dummy_callbacks);
0443     if (rebalance)
0444         ____rb_erase_color(rebalance, root, dummy_rotate);
0445 }
0446 EXPORT_SYMBOL(rb_erase);
0447 
0448 /*
0449  * Augmented rbtree manipulation functions.
0450  *
0451  * This instantiates the same __always_inline functions as in the non-augmented
0452  * case, but this time with user-defined callbacks.
0453  */
0454 
0455 void __rb_insert_augmented(struct rb_node *node, struct rb_root *root,
0456     void (*augment_rotate)(struct rb_node *old, struct rb_node *new))
0457 {
0458     __rb_insert(node, root, augment_rotate);
0459 }
0460 EXPORT_SYMBOL(__rb_insert_augmented);
0461 
0462 /*
0463  * This function returns the first node (in sort order) of the tree.
0464  */
0465 struct rb_node *rb_first(const struct rb_root *root)
0466 {
0467     struct rb_node  *n;
0468 
0469     n = root->rb_node;
0470     if (!n)
0471         return NULL;
0472     while (n->rb_left)
0473         n = n->rb_left;
0474     return n;
0475 }
0476 EXPORT_SYMBOL(rb_first);
0477 
0478 struct rb_node *rb_last(const struct rb_root *root)
0479 {
0480     struct rb_node  *n;
0481 
0482     n = root->rb_node;
0483     if (!n)
0484         return NULL;
0485     while (n->rb_right)
0486         n = n->rb_right;
0487     return n;
0488 }
0489 EXPORT_SYMBOL(rb_last);
0490 
0491 struct rb_node *rb_next(const struct rb_node *node)
0492 {
0493     struct rb_node *parent;
0494 
0495     if (RB_EMPTY_NODE(node))
0496         return NULL;
0497 
0498     /*
0499      * If we have a right-hand child, go down and then left as far
0500      * as we can.
0501      */
0502     if (node->rb_right) {
0503         node = node->rb_right; 
0504         while (node->rb_left)
0505             node=node->rb_left;
0506         return (struct rb_node *)node;
0507     }
0508 
0509     /*
0510      * No right-hand children. Everything down and left is smaller than us,
0511      * so any 'next' node must be in the general direction of our parent.
0512      * Go up the tree; any time the ancestor is a right-hand child of its
0513      * parent, keep going up. First time it's a left-hand child of its
0514      * parent, said parent is our 'next' node.
0515      */
0516     while ((parent = rb_parent(node)) && node == parent->rb_right)
0517         node = parent;
0518 
0519     return parent;
0520 }
0521 EXPORT_SYMBOL(rb_next);
0522 
0523 struct rb_node *rb_prev(const struct rb_node *node)
0524 {
0525     struct rb_node *parent;
0526 
0527     if (RB_EMPTY_NODE(node))
0528         return NULL;
0529 
0530     /*
0531      * If we have a left-hand child, go down and then right as far
0532      * as we can.
0533      */
0534     if (node->rb_left) {
0535         node = node->rb_left; 
0536         while (node->rb_right)
0537             node=node->rb_right;
0538         return (struct rb_node *)node;
0539     }
0540 
0541     /*
0542      * No left-hand children. Go up till we find an ancestor which
0543      * is a right-hand child of its parent.
0544      */
0545     while ((parent = rb_parent(node)) && node == parent->rb_left)
0546         node = parent;
0547 
0548     return parent;
0549 }
0550 EXPORT_SYMBOL(rb_prev);
0551 
0552 void rb_replace_node(struct rb_node *victim, struct rb_node *new,
0553              struct rb_root *root)
0554 {
0555     struct rb_node *parent = rb_parent(victim);
0556 
0557     /* Copy the pointers/colour from the victim to the replacement */
0558     *new = *victim;
0559 
0560     /* Set the surrounding nodes to point to the replacement */
0561     if (victim->rb_left)
0562         rb_set_parent(victim->rb_left, new);
0563     if (victim->rb_right)
0564         rb_set_parent(victim->rb_right, new);
0565     __rb_change_child(victim, new, parent, root);
0566 }
0567 EXPORT_SYMBOL(rb_replace_node);
0568 
0569 void rb_replace_node_rcu(struct rb_node *victim, struct rb_node *new,
0570              struct rb_root *root)
0571 {
0572     struct rb_node *parent = rb_parent(victim);
0573 
0574     /* Copy the pointers/colour from the victim to the replacement */
0575     *new = *victim;
0576 
0577     /* Set the surrounding nodes to point to the replacement */
0578     if (victim->rb_left)
0579         rb_set_parent(victim->rb_left, new);
0580     if (victim->rb_right)
0581         rb_set_parent(victim->rb_right, new);
0582 
0583     /* Set the parent's pointer to the new node last after an RCU barrier
0584      * so that the pointers onwards are seen to be set correctly when doing
0585      * an RCU walk over the tree.
0586      */
0587     __rb_change_child_rcu(victim, new, parent, root);
0588 }
0589 EXPORT_SYMBOL(rb_replace_node_rcu);
0590 
0591 static struct rb_node *rb_left_deepest_node(const struct rb_node *node)
0592 {
0593     for (;;) {
0594         if (node->rb_left)
0595             node = node->rb_left;
0596         else if (node->rb_right)
0597             node = node->rb_right;
0598         else
0599             return (struct rb_node *)node;
0600     }
0601 }
0602 
0603 struct rb_node *rb_next_postorder(const struct rb_node *node)
0604 {
0605     const struct rb_node *parent;
0606     if (!node)
0607         return NULL;
0608     parent = rb_parent(node);
0609 
0610     /* If we're sitting on node, we've already seen our children */
0611     if (parent && node == parent->rb_left && parent->rb_right) {
0612         /* If we are the parent's left node, go to the parent's right
0613          * node then all the way down to the left */
0614         return rb_left_deepest_node(parent->rb_right);
0615     } else
0616         /* Otherwise we are the parent's right node, and the parent
0617          * should be next */
0618         return (struct rb_node *)parent;
0619 }
0620 EXPORT_SYMBOL(rb_next_postorder);
0621 
0622 struct rb_node *rb_first_postorder(const struct rb_root *root)
0623 {
0624     if (!root->rb_node)
0625         return NULL;
0626 
0627     return rb_left_deepest_node(root->rb_node);
0628 }
0629 EXPORT_SYMBOL(rb_first_postorder);