2 files changed, 18 insertions, 11 deletions

diff --git a/lib/rbtree.c b/lib/rbtree.c
index 41cf19b2fe5..baf7c835c57 100644
--- a/lib/rbtree.c
+++ b/lib/rbtree.c
@@ -259,10 +259,22 @@ static void __rb_erase_color(struct rb_node *node, struct rb_node *parent,
 {
         struct rb_node *other;
 
-        while ((!node || rb_is_black(node)) && node != root->rb_node)
-        {
-                if (parent->rb_left == node)
-                {
+        while (true) {
+                /*
+                 * Loop invariant: all leaf paths going through node have a
+                 * black node count that is 1 lower than other leaf paths.
+                 *
+                 * If node is red, we can flip it to black to adjust.
+                 * If node is the root, all leaf paths go through it.
+                 * Otherwise, we need to adjust the tree through color flips
+                 * and tree rotations as per one of the 4 cases below.
+                 */
+                if (node && rb_is_red(node)) {
+                        rb_set_black(node);
+                        break;
+                } else if (!parent) {
+                        break;
+                } else if (parent->rb_left == node) {
                         other = parent->rb_right;
                         if (rb_is_red(other))
                         {
@@ -291,12 +303,9 @@ static void __rb_erase_color(struct rb_node *node, struct rb_node *parent,
                                 rb_set_black(parent);
                                 rb_set_black(other->rb_right);
                                 __rb_rotate_left(parent, root);
-                                node = root->rb_node;
                                 break;
                         }
-                }
-                else
-                {
+                } else {
                         other = parent->rb_left;
                         if (rb_is_red(other))
                         {
@@ -325,13 +334,10 @@ static void __rb_erase_color(struct rb_node *node, struct rb_node *parent,
                                 rb_set_black(parent);
                                 rb_set_black(other->rb_left);
                                 __rb_rotate_right(parent, root);
-                                node = root->rb_node;
                                 break;
                         }
                 }
         }
-        if (node)
-                rb_set_black(node);
 }
 
 void rb_erase(struct rb_node *node, struct rb_root *root)
diff --git a/tools/perf/util/include/linux/rbtree.h b/tools/perf/util/include/linux/rbtree.h
index 2a030c5af3a..9bcdc844b33 100644
--- a/tools/perf/util/include/linux/rbtree.h
+++ b/tools/perf/util/include/linux/rbtree.h
@@ -1,2 +1,3 @@
 #include <stdbool.h>
+#include <stdbool.h>
 #include "../../../../include/linux/rbtree.h"