tools/lib/rbtree.c

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