tools: Update rbtree implementation

[linux.git] / tools / include / linux / rbtree_augmented.h

diff --git a/tools/include/linux/rbtree_augmented.h b/tools/include/linux/rbtree_augmented.h
index 43be941db695958ff2dbd3c084860407976ea5b3..d008e14045802fb77cf1fc0138b340a4f2285647 100644 (file)
--- a/tools/include/linux/rbtree_augmented.h
+++ b/tools/include/linux/rbtree_augmented.h
@@ -44,7 +44,9 @@ struct rb_augment_callbacks {
        void (*rotate)(struct rb_node *old, struct rb_node *new);
 };
 
-extern void __rb_insert_augmented(struct rb_node *node, struct rb_root *root,
+extern 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));
 /*
  * Fixup the rbtree and update the augmented information when rebalancing.
@@ -60,7 +62,16 @@ static inline void
 rb_insert_augmented(struct rb_node *node, struct rb_root *root,
                    const struct rb_augment_callbacks *augment)
 {
-       __rb_insert_augmented(node, root, augment->rotate);
+       __rb_insert_augmented(node, root, false, NULL, augment->rotate);
+}
+
+static inline void
+rb_insert_augmented_cached(struct rb_node *node,
+                          struct rb_root_cached *root, bool newleft,
+                          const struct rb_augment_callbacks *augment)
+{
+       __rb_insert_augmented(node, &root->rb_root,
+                             newleft, &root->rb_leftmost, augment->rotate);
 }
 
 #define RB_DECLARE_CALLBACKS(rbstatic, rbname, rbstruct, rbfield,      \
@@ -93,7 +104,9 @@ rbname ## _rotate(struct rb_node *rb_old, struct rb_node *rb_new)    \
        old->rbaugmented = rbcompute(old);                              \
 }                                                                      \
 rbstatic const struct rb_augment_callbacks rbname = {                  \
-       rbname ## _propagate, rbname ## _copy, rbname ## _rotate        \
+       .propagate = rbname ## _propagate,                              \
+       .copy = rbname ## _copy,                                        \
+       .rotate = rbname ## _rotate                                     \
 };
 
 
@@ -126,11 +139,11 @@ __rb_change_child(struct rb_node *old, struct rb_node *new,
 {
        if (parent) {
                if (parent->rb_left == old)
-                       parent->rb_left = new;
+                       WRITE_ONCE(parent->rb_left, new);
                else
-                       parent->rb_right = new;
+                       WRITE_ONCE(parent->rb_right, new);
        } else
-               root->rb_node = new;
+               WRITE_ONCE(root->rb_node, new);
 }
 
 extern void __rb_erase_color(struct rb_node *parent, struct rb_root *root,
@@ -138,12 +151,17 @@ extern void __rb_erase_color(struct rb_node *parent, struct rb_root *root,
 
 static __always_inline struct rb_node *
 __rb_erase_augmented(struct rb_node *node, struct rb_root *root,
+                    struct rb_node **leftmost,
                     const struct rb_augment_callbacks *augment)
 {
-       struct rb_node *child = node->rb_right, *tmp = node->rb_left;
+       struct rb_node *child = node->rb_right;
+       struct rb_node *tmp = node->rb_left;
        struct rb_node *parent, *rebalance;
        unsigned long pc;
 
+       if (leftmost && node == *leftmost)
+               *leftmost = rb_next(node);
+
        if (!tmp) {
                /*
                 * Case 1: node to erase has no more than 1 child (easy!)
@@ -170,6 +188,7 @@ __rb_erase_augmented(struct rb_node *node, struct rb_root *root,
                tmp = parent;
        } else {
                struct rb_node *successor = child, *child2;
+
                tmp = child->rb_left;
                if (!tmp) {
                        /*
@@ -183,6 +202,7 @@ __rb_erase_augmented(struct rb_node *node, struct rb_root *root,
                         */
                        parent = successor;
                        child2 = successor->rb_right;
+
                        augment->copy(node, successor);
                } else {
                        /*
@@ -204,19 +224,23 @@ __rb_erase_augmented(struct rb_node *node, struct rb_root *root,
                                successor = tmp;
                                tmp = tmp->rb_left;
                        } while (tmp);
-                       parent->rb_left = child2 = successor->rb_right;
-                       successor->rb_right = child;
+                       child2 = successor->rb_right;
+                       WRITE_ONCE(parent->rb_left, child2);
+                       WRITE_ONCE(successor->rb_right, child);
                        rb_set_parent(child, successor);
+
                        augment->copy(node, successor);
                        augment->propagate(parent, successor);
                }
 
-               successor->rb_left = tmp = node->rb_left;
+               tmp = node->rb_left;
+               WRITE_ONCE(successor->rb_left, tmp);
                rb_set_parent(tmp, successor);
 
                pc = node->__rb_parent_color;
                tmp = __rb_parent(pc);
                __rb_change_child(node, successor, tmp, root);
+
                if (child2) {
                        successor->__rb_parent_color = pc;
                        rb_set_parent_color(child2, parent, RB_BLACK);
@@ -237,9 +261,21 @@ static __always_inline void
 rb_erase_augmented(struct rb_node *node, struct rb_root *root,
                   const struct rb_augment_callbacks *augment)
 {
-       struct rb_node *rebalance = __rb_erase_augmented(node, root, augment);
+       struct rb_node *rebalance = __rb_erase_augmented(node, root,
+                                                        NULL, augment);
        if (rebalance)
                __rb_erase_color(rebalance, root, augment->rotate);
 }
 
-#endif /* _TOOLS_LINUX_RBTREE_AUGMENTED_H */
+static __always_inline void
+rb_erase_augmented_cached(struct rb_node *node, struct rb_root_cached *root,
+                         const struct rb_augment_callbacks *augment)
+{
+       struct rb_node *rebalance = __rb_erase_augmented(node, &root->rb_root,
+                                                        &root->rb_leftmost,
+                                                        augment);
+       if (rebalance)
+               __rb_erase_color(rebalance, &root->rb_root, augment->rotate);
+}
+
+#endif /* _TOOLS_LINUX_RBTREE_AUGMENTED_H */