]> asedeno.scripts.mit.edu Git - linux.git/blobdiff - tools/lib/rbtree.c
tools: Update rbtree implementation
[linux.git] / tools / lib / rbtree.c
index 17c2b596f0437f1f32d46441558f6f2329434351..904adb70a4f065acfda1c56b8b0de642796cb222 100644 (file)
@@ -22,6 +22,7 @@
 */
 
 #include <linux/rbtree_augmented.h>
+#include <linux/export.h>
 
 /*
  * red-black trees properties:  http://en.wikipedia.org/wiki/Rbtree
  *  parentheses and have some accompanying text comment.
  */
 
+/*
+ * Notes on lockless lookups:
+ *
+ * All stores to the tree structure (rb_left and rb_right) must be done using
+ * WRITE_ONCE(). And we must not inadvertently cause (temporary) loops in the
+ * tree structure as seen in program order.
+ *
+ * These two requirements will allow lockless iteration of the tree -- not
+ * correct iteration mind you, tree rotations are not atomic so a lookup might
+ * miss entire subtrees.
+ *
+ * But they do guarantee that any such traversal will only see valid elements
+ * and that it will indeed complete -- does not get stuck in a loop.
+ *
+ * It also guarantees that if the lookup returns an element it is the 'correct'
+ * one. But not returning an element does _NOT_ mean it's not present.
+ *
+ * NOTE:
+ *
+ * Stores to __rb_parent_color are not important for simple lookups so those
+ * are left undone as of now. Nor did I check for loops involving parent
+ * pointers.
+ */
+
 static inline void rb_set_black(struct rb_node *rb)
 {
        rb->__rb_parent_color |= RB_BLACK;
@@ -70,22 +95,35 @@ __rb_rotate_set_parents(struct rb_node *old, struct rb_node *new,
 
 static __always_inline void
 __rb_insert(struct rb_node *node, struct rb_root *root,
+           bool newleft, struct rb_node **leftmost,
            void (*augment_rotate)(struct rb_node *old, struct rb_node *new))
 {
        struct rb_node *parent = rb_red_parent(node), *gparent, *tmp;
 
+       if (newleft)
+               *leftmost = node;
+
        while (true) {
                /*
-                * Loop invariant: node is red
-                *
-                * If there is a black parent, we are done.
-                * Otherwise, take some corrective action as we don't
-                * want a red root or two consecutive red nodes.
+                * Loop invariant: node is red.
                 */
-               if (!parent) {
+               if (unlikely(!parent)) {
+                       /*
+                        * The inserted node is root. Either this is the
+                        * first node, or we recursed at Case 1 below and
+                        * are no longer violating 4).
+                        */
                        rb_set_parent_color(node, NULL, RB_BLACK);
                        break;
-               } else if (rb_is_black(parent))
+               }
+
+               /*
+                * If there is a black parent, we are done.
+                * Otherwise, take some corrective action as,
+                * per 4), we don't want a red root or two
+                * consecutive red nodes.
+                */
+               if(rb_is_black(parent))
                        break;
 
                gparent = rb_red_parent(parent);
@@ -94,7 +132,7 @@ __rb_insert(struct rb_node *node, struct rb_root *root,
                if (parent != tmp) {    /* parent == gparent->rb_left */
                        if (tmp && rb_is_red(tmp)) {
                                /*
-                                * Case 1 - color flips
+                                * Case 1 - node's uncle is red (color flips).
                                 *
                                 *       G            g
                                 *      / \          / \
@@ -117,7 +155,8 @@ __rb_insert(struct rb_node *node, struct rb_root *root,
                        tmp = parent->rb_right;
                        if (node == tmp) {
                                /*
-                                * Case 2 - left rotate at parent
+                                * Case 2 - node's uncle is black and node is
+                                * the parent's right child (left rotate at parent).
                                 *
                                 *      G             G
                                 *     / \           / \
@@ -128,8 +167,9 @@ __rb_insert(struct rb_node *node, struct rb_root *root,
                                 * This still leaves us in violation of 4), the
                                 * continuation into Case 3 will fix that.
                                 */
-                               parent->rb_right = tmp = node->rb_left;
-                               node->rb_left = parent;
+                               tmp = node->rb_left;
+                               WRITE_ONCE(parent->rb_right, tmp);
+                               WRITE_ONCE(node->rb_left, parent);
                                if (tmp)
                                        rb_set_parent_color(tmp, parent,
                                                            RB_BLACK);
@@ -140,7 +180,8 @@ __rb_insert(struct rb_node *node, struct rb_root *root,
                        }
 
                        /*
-                        * Case 3 - right rotate at gparent
+                        * Case 3 - node's uncle is black and node is
+                        * the parent's left child (right rotate at gparent).
                         *
                         *        G           P
                         *       / \         / \
@@ -148,8 +189,8 @@ __rb_insert(struct rb_node *node, struct rb_root *root,
                         *     /                 \
                         *    n                   U
                         */
-                       gparent->rb_left = tmp;  /* == parent->rb_right */
-                       parent->rb_right = gparent;
+                       WRITE_ONCE(gparent->rb_left, tmp); /* == parent->rb_right */
+                       WRITE_ONCE(parent->rb_right, gparent);
                        if (tmp)
                                rb_set_parent_color(tmp, gparent, RB_BLACK);
                        __rb_rotate_set_parents(gparent, parent, root, RB_RED);
@@ -170,8 +211,9 @@ __rb_insert(struct rb_node *node, struct rb_root *root,
                        tmp = parent->rb_left;
                        if (node == tmp) {
                                /* Case 2 - right rotate at parent */
-                               parent->rb_left = tmp = node->rb_right;
-                               node->rb_right = parent;
+                               tmp = node->rb_right;
+                               WRITE_ONCE(parent->rb_left, tmp);
+                               WRITE_ONCE(node->rb_right, parent);
                                if (tmp)
                                        rb_set_parent_color(tmp, parent,
                                                            RB_BLACK);
@@ -182,8 +224,8 @@ __rb_insert(struct rb_node *node, struct rb_root *root,
                        }
 
                        /* Case 3 - left rotate at gparent */
-                       gparent->rb_right = tmp;  /* == parent->rb_left */
-                       parent->rb_left = gparent;
+                       WRITE_ONCE(gparent->rb_right, tmp); /* == parent->rb_left */
+                       WRITE_ONCE(parent->rb_left, gparent);
                        if (tmp)
                                rb_set_parent_color(tmp, gparent, RB_BLACK);
                        __rb_rotate_set_parents(gparent, parent, root, RB_RED);
@@ -223,8 +265,9 @@ ____rb_erase_color(struct rb_node *parent, struct rb_root *root,
                                 *      / \         / \
                                 *     Sl  Sr      N   Sl
                                 */
-                               parent->rb_right = tmp1 = sibling->rb_left;
-                               sibling->rb_left = parent;
+                               tmp1 = sibling->rb_left;
+                               WRITE_ONCE(parent->rb_right, tmp1);
+                               WRITE_ONCE(sibling->rb_left, parent);
                                rb_set_parent_color(tmp1, parent, RB_BLACK);
                                __rb_rotate_set_parents(parent, sibling, root,
                                                        RB_RED);
@@ -268,15 +311,31 @@ ____rb_erase_color(struct rb_node *parent, struct rb_root *root,
                                 *
                                 *   (p)           (p)
                                 *   / \           / \
-                                *  N   S    -->  N   Sl
+                                *  N   S    -->  N   sl
                                 *     / \             \
-                                *    sl  Sr            s
+                                *    sl  Sr            S
                                 *                       \
                                 *                        Sr
+                                *
+                                * Note: p might be red, and then both
+                                * p and sl are red after rotation(which
+                                * breaks property 4). This is fixed in
+                                * Case 4 (in __rb_rotate_set_parents()
+                                *         which set sl the color of p
+                                *         and set p RB_BLACK)
+                                *
+                                *   (p)            (sl)
+                                *   / \            /  \
+                                *  N   sl   -->   P    S
+                                *       \        /      \
+                                *        S      N        Sr
+                                *         \
+                                *          Sr
                                 */
-                               sibling->rb_left = tmp1 = tmp2->rb_right;
-                               tmp2->rb_right = sibling;
-                               parent->rb_right = tmp2;
+                               tmp1 = tmp2->rb_right;
+                               WRITE_ONCE(sibling->rb_left, tmp1);
+                               WRITE_ONCE(tmp2->rb_right, sibling);
+                               WRITE_ONCE(parent->rb_right, tmp2);
                                if (tmp1)
                                        rb_set_parent_color(tmp1, sibling,
                                                            RB_BLACK);
@@ -296,8 +355,9 @@ ____rb_erase_color(struct rb_node *parent, struct rb_root *root,
                         *        / \         / \
                         *      (sl) sr      N  (sl)
                         */
-                       parent->rb_right = tmp2 = sibling->rb_left;
-                       sibling->rb_left = parent;
+                       tmp2 = sibling->rb_left;
+                       WRITE_ONCE(parent->rb_right, tmp2);
+                       WRITE_ONCE(sibling->rb_left, parent);
                        rb_set_parent_color(tmp1, sibling, RB_BLACK);
                        if (tmp2)
                                rb_set_parent(tmp2, parent);
@@ -309,8 +369,9 @@ ____rb_erase_color(struct rb_node *parent, struct rb_root *root,
                        sibling = parent->rb_left;
                        if (rb_is_red(sibling)) {
                                /* Case 1 - right rotate at parent */
-                               parent->rb_left = tmp1 = sibling->rb_right;
-                               sibling->rb_right = parent;
+                               tmp1 = sibling->rb_right;
+                               WRITE_ONCE(parent->rb_left, tmp1);
+                               WRITE_ONCE(sibling->rb_right, parent);
                                rb_set_parent_color(tmp1, parent, RB_BLACK);
                                __rb_rotate_set_parents(parent, sibling, root,
                                                        RB_RED);
@@ -334,10 +395,11 @@ ____rb_erase_color(struct rb_node *parent, struct rb_root *root,
                                        }
                                        break;
                                }
-                               /* Case 3 - right rotate at sibling */
-                               sibling->rb_right = tmp1 = tmp2->rb_left;
-                               tmp2->rb_left = sibling;
-                               parent->rb_left = tmp2;
+                               /* Case 3 - left rotate at sibling */
+                               tmp1 = tmp2->rb_left;
+                               WRITE_ONCE(sibling->rb_right, tmp1);
+                               WRITE_ONCE(tmp2->rb_left, sibling);
+                               WRITE_ONCE(parent->rb_left, tmp2);
                                if (tmp1)
                                        rb_set_parent_color(tmp1, sibling,
                                                            RB_BLACK);
@@ -345,9 +407,10 @@ ____rb_erase_color(struct rb_node *parent, struct rb_root *root,
                                tmp1 = sibling;
                                sibling = tmp2;
                        }
-                       /* Case 4 - left rotate at parent + color flips */
-                       parent->rb_left = tmp2 = sibling->rb_right;
-                       sibling->rb_right = parent;
+                       /* Case 4 - right rotate at parent + color flips */
+                       tmp2 = sibling->rb_right;
+                       WRITE_ONCE(parent->rb_left, tmp2);
+                       WRITE_ONCE(sibling->rb_right, parent);
                        rb_set_parent_color(tmp1, sibling, RB_BLACK);
                        if (tmp2)
                                rb_set_parent(tmp2, parent);
@@ -378,22 +441,41 @@ static inline void dummy_copy(struct rb_node *old, struct rb_node *new) {}
 static inline void dummy_rotate(struct rb_node *old, struct rb_node *new) {}
 
 static const struct rb_augment_callbacks dummy_callbacks = {
-       dummy_propagate, dummy_copy, dummy_rotate
+       .propagate = dummy_propagate,
+       .copy = dummy_copy,
+       .rotate = dummy_rotate
 };
 
 void rb_insert_color(struct rb_node *node, struct rb_root *root)
 {
-       __rb_insert(node, root, dummy_rotate);
+       __rb_insert(node, root, false, NULL, dummy_rotate);
 }
 
 void rb_erase(struct rb_node *node, struct rb_root *root)
 {
        struct rb_node *rebalance;
-       rebalance = __rb_erase_augmented(node, root, &dummy_callbacks);
+       rebalance = __rb_erase_augmented(node, root,
+                                        NULL, &dummy_callbacks);
        if (rebalance)
                ____rb_erase_color(rebalance, root, dummy_rotate);
 }
 
+void rb_insert_color_cached(struct rb_node *node,
+                           struct rb_root_cached *root, bool leftmost)
+{
+       __rb_insert(node, &root->rb_root, leftmost,
+                   &root->rb_leftmost, dummy_rotate);
+}
+
+void rb_erase_cached(struct rb_node *node, struct rb_root_cached *root)
+{
+       struct rb_node *rebalance;
+       rebalance = __rb_erase_augmented(node, &root->rb_root,
+                                        &root->rb_leftmost, &dummy_callbacks);
+       if (rebalance)
+               ____rb_erase_color(rebalance, &root->rb_root, dummy_rotate);
+}
+
 /*
  * Augmented rbtree manipulation functions.
  *
@@ -402,9 +484,10 @@ void rb_erase(struct rb_node *node, struct rb_root *root)
  */
 
 void __rb_insert_augmented(struct rb_node *node, struct rb_root *root,
+                          bool newleft, struct rb_node **leftmost,
        void (*augment_rotate)(struct rb_node *old, struct rb_node *new))
 {
-       __rb_insert(node, root, augment_rotate);
+       __rb_insert(node, root, newleft, leftmost, augment_rotate);
 }
 
 /*
@@ -498,15 +581,24 @@ void rb_replace_node(struct rb_node *victim, struct rb_node *new,
 {
        struct rb_node *parent = rb_parent(victim);
 
+       /* Copy the pointers/colour from the victim to the replacement */
+       *new = *victim;
+
        /* Set the surrounding nodes to point to the replacement */
-       __rb_change_child(victim, new, parent, root);
        if (victim->rb_left)
                rb_set_parent(victim->rb_left, new);
        if (victim->rb_right)
                rb_set_parent(victim->rb_right, new);
+       __rb_change_child(victim, new, parent, root);
+}
 
-       /* Copy the pointers/colour from the victim to the replacement */
-       *new = *victim;
+void rb_replace_node_cached(struct rb_node *victim, struct rb_node *new,
+                           struct rb_root_cached *root)
+{
+       rb_replace_node(victim, new, &root->rb_root);
+
+       if (root->rb_leftmost == victim)
+               root->rb_leftmost = new;
 }
 
 static struct rb_node *rb_left_deepest_node(const struct rb_node *node)