[RFC PATCH 0/6] augmented rbtree changes

[RFC PATCH 0/6] augmented rbtree changes

[Michel Lespinasse]

I've been looking at rbtrees with an eye towards improving the augmented
rbtree support, and even though I don't consider this work done, I am
getting to the stage where I would like to get feedback.

Patches 1-2 are generic rbtree improvements I came up with after sending
the previous patch series. Patch 1 makes rb_erase easier to read (IMO),
while patch 2 is another minor optimization in the rbtree rebalancing code.

Patch 3 adds an augmented rbtree test, where nodes get a new 'value' field
in addition to the existing sort key, and we want to maintain an 'augmented'
field for each node, which should be equal to the max of all 'value' fields
for nodes within that node's subtree.

Patch 4 speeds up augmented rbtree insertion. We make the handcoded search
function responsible for updating the augmented values on the path to the
insertion point, then we use a new version of rb_insert_color() which has
an added callback for updating augmented values on a tree rotation.

Patch 5 speeds up the augmented rbtree erase. Here again we use a tree
rotation callback during rebalancing; however we also have to propagate
the augmented node information above nodes being erased and/or stitched,
and I haven't found a nice enough way to do that. So for now I am proposing
the simple-stupid way of propagating all the way to the root. More on
this later.

Patch 6 removes the prior augmented API interface, and migrates its single
current user to the proposed new interface.


IMO patches 1-2 are ready for inclusion in -mm tree along with the
previous rbtree series.


I would like feedback on the rest - but first, I think I should mention
what use cases I envision for this augmented rbtree support. The way
I see it, augmented rbtree could be used in:

- Some kind of generic interval tree support, where nodes have explicit
(start, end) values. The rbtree is sorted by start order, and nodes
are augmented with a max(node->end for node in subtree) property.
arch/x86/mm/pat_rbtree.c and mm/kmemleak.c could make use of that
(instead of an ad-hoc interval tree implementations based on augmented
rbtree and prio tree respectively);

- The rbtree for all VMAs in a MM could be augmented, as suggested by Rik,
to maintain a max(empty gap before node's VMA for node in subtree) property
and support fast virtual address space allocation;

- The prio tree of all VMAs mapping a given file (struct address_space)
could be switched to an augmented rbtree based interval tree (thus removing
the prio tree library in favor of augmented rbtrees)

- I would like to introduce yet another interval tree to hold all VMAs
sharing the same anon_vma (so the linked list of all AVCs with a given
anon_vma would be replaced with an interval tree).

With these plans, each VMA could be on up to 3 separate augmented rbtrees,
so that's why I want them to be fast :)


As they stand, patches 3-6 don't seem to make a difference for basic rbtree
support, and they improve my augmented rbtree insertion/erase benchmark
by a factor of ~2.1 to ~2.3 depending on test machines.


In addition to the usual feedback about code sanity or lack thereof, I would
like to ask if I stroke the right balance on pure speed vs code size.
I think having generic functions for augmented rbtree rebalancing, with
a callback for tree rotations, is probably a decent choice given that we
might have to update several augmented rbtrees in sequence when adding or
removing VMAs.


Another point I am not fully happy with is the way I am propagating augmented
subtree information in rb_erase_augmented(). I initially tried to stop
propagating updates up as soon as a node was reached that already had the
proper augmented value. However, there are a few complications with that.

In case 2 of rb_erase_augmented(), if 'old' had the highest augmented value
in the subtree and 'node' had the next highest value, node's augmented
value is still correct after stitching it as the new root of that subtree,
but its parent's augmented value must still be adjusted as 'old', which
had the highest value in that subtree, has been removed from the subtree.

In case 3 of rb_erase_augmented(), 'node' gets stitched out of its place
and relocated to the subtree root, in place of 'old'. As a result, we
might have to propagate augmented values a few levels above node's old
location, before reaching a node that already has the right augmented
value, but is still below the point where 'old' was replaced with 'new'
which might have a lower augmented value (if 'old' had the highest
value for that subtree).

So while it would be possible to handle all these cases without propagating
all the way to the root (and it should be more efficient too - most of the
nodes in a balanced tree are on the last few levels, so having to go all
the way back to the root really is wasteful), I have not found a nice
elegant way to do that yet, let alone in a generic way. If someone wants
to try their hand at that problem, I would be very interested to see what
they can come up with.


Michel Lespinasse (6):
  rbtree: rb_erase updates and comments
  rbtree: optimize fetching of sibling node
  augmented rbtree test
  rbtree: faster augmented insert
  rbtree: faster augmented erase
  rbtree: remove prior augmented rbtree implementation

 arch/x86/mm/pat_rbtree.c        |   52 ++++++----
 include/linux/rbtree.h          |    8 --
 include/linux/rbtree_internal.h |  131 +++++++++++++++++++++++++
 lib/rbtree.c                    |  206 ++++++++-------------------------------
 lib/rbtree_test.c               |  120 ++++++++++++++++++++++-
 5 files changed, 322 insertions(+), 195 deletions(-)
 create mode 100644 include/linux/rbtree_internal.h

-- 
1.7.7.3

[PATCH 1/6] rbtree: rb_erase updates and comments

[Michel Lespinasse]

Minor updates to the rb_erase() function:
- Reorder code to put simplest / common case (no more than 1 child) first.
- Fetch the parent first, since it ends up being required in all 3 cases.
- Add a few comments to illustrate case 2 (node to remove has 2 childs,
  but one of them is the successor) and case 3 (node to remove has 2 childs,
  successor is a left-descendant of the right child).

Signed-off-by: Michel Lespinasse <walken@google.com>
---
 lib/rbtree.c |  115 ++++++++++++++++++++++++++++++++++++----------------------
 1 files changed, 72 insertions(+), 43 deletions(-)

diff --git a/lib/rbtree.c b/lib/rbtree.c
index 0892670..3b6ec98 100644
--- a/lib/rbtree.c
+++ b/lib/rbtree.c
@@ -2,7 +2,8 @@
   Red Black Trees
   (C) 1999  Andrea Arcangeli <andrea@suse.de>
   (C) 2002  David Woodhouse <dwmw2@infradead.org>
-  
+  (C) 2012  Michel Lespinasse <walken@google.com>
+
   This program is free software; you can redistribute it and/or modify
   it under the terms of the GNU General Public License as published by
   the Free Software Foundation; either version 2 of the License, or
@@ -356,65 +357,93 @@ static void __rb_erase_color(struct rb_node *node, struct rb_node *parent,
 
 void rb_erase(struct rb_node *node, struct rb_root *root)
 {
-	struct rb_node *child, *parent;
-	int color;
+	struct rb_node *parent = rb_parent(node);
+	struct rb_node *child = node->rb_right;
+	struct rb_node *tmp = node->rb_left;
+	bool black;
+
+	if (!tmp) {
+		/* Case 1: node to erase has no more than 1 child (easy!) */
+		if (child)
+one_child:
+			rb_set_parent(child, parent);
+		if (parent) {
+			if (parent->rb_left == node)
+				parent->rb_left = child;
+			else
+				parent->rb_right = child;
+		} else
+			root->rb_node = child;
 
-	if (!node->rb_left)
-		child = node->rb_right;
-	else if (!node->rb_right)
-		child = node->rb_left;
-	else {
-		struct rb_node *old = node, *left;
+		black = rb_is_black(node);
+	} else if (!child) {
+		/* Still case 1, but this time the child is node->rb_left */
+		child = tmp;
+		goto one_child;
+	} else {
+		struct rb_node *old = node;
 
+		/*
+		 * Old is the node we want to erase. It's got left and right
+		 * children, which makes things difficult. Let's find the
+		 * next node in the tree to have it fill old's position.
+		 */
 		node = node->rb_right;
-		while ((left = node->rb_left) != NULL)
-			node = left;
+		while ((tmp = node->rb_left) != NULL)
+			node = tmp;
 
-		if (rb_parent(old)) {
-			if (rb_parent(old)->rb_left == old)
-				rb_parent(old)->rb_left = node;
+		/* Graft node (old's successor) under old's parent. */
+		if (parent) {
+			if (parent->rb_left == old)
+				parent->rb_left = node;
 			else
-				rb_parent(old)->rb_right = node;
+				parent->rb_right = node;
 		} else
 			root->rb_node = node;
 
-		child = node->rb_right;
 		parent = rb_parent(node);
-		color = rb_color(node);
+		black = rb_is_black(node);
+		node->__rb_parent_color = old->__rb_parent_color;
+
+		/*
+		 * Node doesn't have a left child, since it is old's successor,
+		 * so we can take old's left child and graft it under node.
+		 */
+		node->rb_left = tmp = old->rb_left;
+		rb_set_parent(tmp, node);
 
+		child = node->rb_right;
 		if (parent == old) {
+			/*
+			 * Case 2: old is node's parent (we are done!)
+			 *
+			 *    (o)          (n)
+			 *    / \          / \
+			 *  (x) (n)  ->  (x) (c)
+			 *        \
+			 *        (c)
+			 */
 			parent = node;
 		} else {
+			/*
+			 * Case 3: old is node's ancestor but not its parent
+			 *
+			 *    (o)          (n)
+			 *    / \          / \
+			 *  (x) (y)  ->  (x) (y)
+			 *      /            /
+			 *    (n)          (c)
+			 *      \
+			 *      (c)
+			 */
+			node->rb_right = tmp = old->rb_right;
+			parent->rb_left = child;
+			rb_set_parent(tmp, node);
 			if (child)
 				rb_set_parent(child, parent);
-			parent->rb_left = child;
-
-			node->rb_right = old->rb_right;
-			rb_set_parent(old->rb_right, node);
 		}
-
-		node->__rb_parent_color = old->__rb_parent_color;
-		node->rb_left = old->rb_left;
-		rb_set_parent(old->rb_left, node);
-
-		goto color;
 	}
-
-	parent = rb_parent(node);
-	color = rb_color(node);
-
-	if (child)
-		rb_set_parent(child, parent);
-	if (parent) {
-		if (parent->rb_left == node)
-			parent->rb_left = child;
-		else
-			parent->rb_right = child;
-	} else
-		root->rb_node = child;
-
-color:
-	if (color == RB_BLACK)
+	if (black)
 		__rb_erase_color(child, parent, root);
 }
 EXPORT_SYMBOL(rb_erase);
-- 
1.7.7.3

[PATCH 2/6] rbtree: optimize fetching of sibling node

[Michel Lespinasse]

When looking to fetch a node's sibling, we went through a sequence of:
- check if node is the parent's left child
- if it is, then fetch the parent's right child

This can be replaced with:
- fetch the parent's right child as an assumed sibling
- check that node is NOT the fetched child

This avoids fetching the parent's left child when node is actually
that child. Saves a bit on code size, though it doesn't seem to make
a large difference in speed.

Signed-off-by: Michel Lespinasse <walken@google.com>
---
 lib/rbtree.c |   21 +++++++++++++--------
 1 files changed, 13 insertions(+), 8 deletions(-)

diff --git a/lib/rbtree.c b/lib/rbtree.c
index 3b6ec98..8b111cc 100644
--- a/lib/rbtree.c
+++ b/lib/rbtree.c
@@ -108,8 +108,8 @@ void rb_insert_color(struct rb_node *node, struct rb_root *root)
 
 		gparent = rb_red_parent(parent);
 
-		if (parent == gparent->rb_left) {
-			tmp = gparent->rb_right;
+		tmp = gparent->rb_right;
+		if (parent != tmp) {	/* parent == gparent->rb_left */
 			if (tmp && rb_is_red(tmp)) {
 				/*
 				 * Case 1 - color flips
@@ -132,7 +132,8 @@ void rb_insert_color(struct rb_node *node, struct rb_root *root)
 				continue;
 			}
 
-			if (parent->rb_right == node) {
+			tmp = parent->rb_right;
+			if (node == tmp) {
 				/*
 				 * Case 2 - left rotate at parent
 				 *
@@ -152,6 +153,7 @@ void rb_insert_color(struct rb_node *node, struct rb_root *root)
 							    RB_BLACK);
 				rb_set_parent_color(parent, node, RB_RED);
 				parent = node;
+				tmp = node->rb_right;
 			}
 
 			/*
@@ -163,7 +165,7 @@ void rb_insert_color(struct rb_node *node, struct rb_root *root)
 			 *     /                 \
 			 *    n                   U
 			 */
-			gparent->rb_left = tmp = parent->rb_right;
+			gparent->rb_left = tmp;  /* == parent->rb_right */
 			parent->rb_right = gparent;
 			if (tmp)
 				rb_set_parent_color(tmp, gparent, RB_BLACK);
@@ -181,7 +183,8 @@ void rb_insert_color(struct rb_node *node, struct rb_root *root)
 				continue;
 			}
 
-			if (parent->rb_left == node) {
+			tmp = parent->rb_left;
+			if (node == tmp) {
 				/* Case 2 - right rotate at parent */
 				parent->rb_left = tmp = node->rb_right;
 				node->rb_right = parent;
@@ -190,10 +193,11 @@ void rb_insert_color(struct rb_node *node, struct rb_root *root)
 							    RB_BLACK);
 				rb_set_parent_color(parent, node, RB_RED);
 				parent = node;
+				tmp = node->rb_left;
 			}
 
 			/* Case 3 - left rotate at gparent */
-			gparent->rb_right = tmp = parent->rb_left;
+			gparent->rb_right = tmp;  /* == parent->rb_left */
 			parent->rb_left = gparent;
 			if (tmp)
 				rb_set_parent_color(tmp, gparent, RB_BLACK);
@@ -224,8 +228,9 @@ static void __rb_erase_color(struct rb_node *node, struct rb_node *parent,
 			break;
 		} else if (!parent) {
 			break;
-		} else if (parent->rb_left == node) {
-			sibling = parent->rb_right;
+		}
+		sibling = parent->rb_right;
+		if (node != sibling) {	/* node == parent->rb_left */
 			if (rb_is_red(sibling)) {
 				/*
 				 * Case 1 - left rotate at parent
-- 
1.7.7.3

[PATCH 3/6] augmented rbtree test

[Michel Lespinasse]

Signed-off-by: Michel Lespinasse <walken@google.com>
---
 lib/rbtree_test.c |  103 +++++++++++++++++++++++++++++++++++++++++++++++++++-
 1 files changed, 101 insertions(+), 2 deletions(-)

diff --git a/lib/rbtree_test.c b/lib/rbtree_test.c
index 4c6d250..2dfafe4 100644
--- a/lib/rbtree_test.c
+++ b/lib/rbtree_test.c
@@ -10,6 +10,10 @@
 struct test_node {
 	struct rb_node rb;
 	u32 key;
+
+	/* following fields used for testing augmented rbtree functionality */
+	u32 val;
+	u32 augmented;
 };
 
 static struct rb_root root = RB_ROOT;
@@ -20,10 +24,11 @@ static struct rnd_state rnd;
 static void insert(struct test_node *node, struct rb_root *root)
 {
 	struct rb_node **new = &root->rb_node, *parent = NULL;
+	u32 key = node->key;
 
 	while (*new) {
 		parent = *new;
-		if (node->key < rb_entry(parent, struct test_node, rb)->key)
+		if (key < rb_entry(parent, struct test_node, rb)->key)
 			new = &parent->rb_left;
 		else
 			new = &parent->rb_right;
@@ -38,11 +43,62 @@ static inline void erase(struct test_node *node, struct rb_root *root)
 	rb_erase(&node->rb, root);
 }
 
+static inline u32 augment_recompute(struct test_node *node)
+{
+	u32 max = node->val, child_augmented;
+	if (node->rb.rb_left) {
+		child_augmented = rb_entry(node->rb.rb_left, struct test_node,
+					   rb)->augmented;
+		if (max < child_augmented)
+			max = child_augmented;
+	}
+	if (node->rb.rb_right) {
+		child_augmented = rb_entry(node->rb.rb_right, struct test_node,
+					   rb)->augmented;
+		if (max < child_augmented)
+			max = child_augmented;
+	}
+	return max;
+}
+
+static void augment_callback(struct rb_node *rb, void *unused)
+{
+	struct test_node *node = rb_entry(rb, struct test_node, rb);
+	node->augmented = augment_recompute(node);
+}
+
+static void insert_augmented(struct test_node *node, struct rb_root *root)
+{
+	struct rb_node **new = &root->rb_node, *parent = NULL;
+	u32 key = node->key;
+
+	while (*new) {
+		parent = *new;
+		if (key < rb_entry(parent, struct test_node, rb)->key)
+			new = &parent->rb_left;
+		else
+			new = &parent->rb_right;
+	}
+
+	rb_link_node(&node->rb, parent, new);
+	rb_insert_color(&node->rb, root);
+	rb_augment_insert(&node->rb, augment_callback, NULL);
+}
+
+static void erase_augmented(struct test_node *node, struct rb_root *root)
+{
+	struct rb_node *deepest = rb_augment_erase_begin(&node->rb);
+	rb_erase(&node->rb, root);
+	rb_augment_erase_end(deepest, augment_callback, NULL);
+}
+
 static void init(void)
 {
 	int i;
-	for (i = 0; i < NODES; i++)
+	for (i = 0; i < NODES; i++) {
 		nodes[i].key = prandom32(&rnd);
+		nodes[i].val = prandom32(&rnd);
+	}
 }
 
 static bool is_red(struct rb_node *rb)
@@ -81,6 +137,17 @@ static void check(int nr_nodes)
 	WARN_ON_ONCE(count != nr_nodes);
 }
 
+static void check_augmented(int nr_nodes)
+{
+	struct rb_node *rb;
+
+	check(nr_nodes);
+	for (rb = rb_first(&root); rb; rb = rb_next(rb)) {
+		struct test_node *node = rb_entry(rb, struct test_node, rb);
+		WARN_ON_ONCE(node->augmented != augment_recompute(node));
+	}
+}
+
 static int rbtree_test_init(void)
 {
 	int i, j;
@@ -119,6 +186,38 @@ static int rbtree_test_init(void)
 		check(0);
 	}
 
+	printk(KERN_ALERT "augmented rbtree testing");
+
+	init();
+
+	time1 = get_cycles();
+
+	for (i = 0; i < PERF_LOOPS; i++) {
+		for (j = 0; j < NODES; j++)
+			insert_augmented(nodes + j, &root);
+		for (j = 0; j < NODES; j++)
+			erase_augmented(nodes + j, &root);
+	}
+
+	time2 = get_cycles();
+	time = time2 - time1;
+
+	time = div_u64(time, PERF_LOOPS);
+	printk(" -> %llu cycles\n", time);
+
+	for (i = 0; i < CHECK_LOOPS; i++) {
+		init();
+		for (j = 0; j < NODES; j++) {
+			check_augmented(j);
+			insert_augmented(nodes + j, &root);
+		}
+		for (j = 0; j < NODES; j++) {
+			check_augmented(NODES - j);
+			erase_augmented(nodes + j, &root);
+		}
+		check_augmented(0);
+	}
+
 	return -EAGAIN; /* Fail will directly unload the module */
 }
 
-- 
1.7.7.3

[PATCH 4/6] rbtree: faster augmented insert

[Michel Lespinasse]

Introduce rb_insert_augmented(), which is a version of rb_insert_color()
with an added callback on tree rotations. This can be used for insertion
into an augmented tree: the handcoded search phase must be updated to
maintain the augmented information on insertion, and then the rbtree
coloring/rebalancing algorithms keep it up to date.

rb_insert_color() is now a special case of rb_insert_augmented() with
a do-nothing callback. I used inlining to optimize out the callback,
with the intent that this would generate the same code as previously
for rb_insert_augmented(). This didn't fully work, as my compiler output
is now *smaller* than before for that function. Speed wise, they seem
comparable though.

Signed-off-by: Michel Lespinasse <walken@google.com>
---
 include/linux/rbtree.h |    5 +++++
 lib/rbtree.c           |   14 +++++++++++++-
 lib/rbtree_test.c      |   31 +++++++++++++++++++++++--------
 3 files changed, 41 insertions(+), 9 deletions(-)

diff --git a/include/linux/rbtree.h b/include/linux/rbtree.h
index bf836a2..1364b81 100644
--- a/include/linux/rbtree.h
+++ b/include/linux/rbtree.h
@@ -61,6 +61,11 @@ struct rb_root {
 extern void rb_insert_color(struct rb_node *, struct rb_root *);
 extern void rb_erase(struct rb_node *, struct rb_root *);
 
+typedef void rb_augment_rotate(struct rb_node *old, struct rb_node *new);
+
+extern void rb_insert_augmented(struct rb_node *node, struct rb_root *root,
+				rb_augment_rotate *augment);
+
 typedef void (*rb_augment_f)(struct rb_node *node, void *data);
 
 extern void rb_augment_insert(struct rb_node *node,
diff --git a/lib/rbtree.c b/lib/rbtree.c
index 8b111cc..a6ae4c5 100644
--- a/lib/rbtree.c
+++ b/lib/rbtree.c
@@ -88,7 +88,8 @@ __rb_rotate_set_parents(struct rb_node *old, struct rb_node *new,
 		root->rb_node = new;
 }
 
-void rb_insert_color(struct rb_node *node, struct rb_root *root)
+inline void rb_insert_augmented(struct rb_node *node, struct rb_root *root,
+				rb_augment_rotate *augment)
 {
 	struct rb_node *parent = rb_red_parent(node), *gparent, *tmp;
 
@@ -152,6 +153,7 @@ void rb_insert_color(struct rb_node *node, struct rb_root *root)
 					rb_set_parent_color(tmp, parent,
 							    RB_BLACK);
 				rb_set_parent_color(parent, node, RB_RED);
+				augment(parent, node);
 				parent = node;
 				tmp = node->rb_right;
 			}
@@ -170,6 +172,7 @@ void rb_insert_color(struct rb_node *node, struct rb_root *root)
 			if (tmp)
 				rb_set_parent_color(tmp, gparent, RB_BLACK);
 			__rb_rotate_set_parents(gparent, parent, root, RB_RED);
+			augment(gparent, parent);
 			break;
 		} else {
 			tmp = gparent->rb_left;
@@ -192,6 +195,7 @@ void rb_insert_color(struct rb_node *node, struct rb_root *root)
 					rb_set_parent_color(tmp, parent,
 							    RB_BLACK);
 				rb_set_parent_color(parent, node, RB_RED);
+				augment(parent, node);
 				parent = node;
 				tmp = node->rb_left;
 			}
@@ -202,10 +206,18 @@ void rb_insert_color(struct rb_node *node, struct rb_root *root)
 			if (tmp)
 				rb_set_parent_color(tmp, gparent, RB_BLACK);
 			__rb_rotate_set_parents(gparent, parent, root, RB_RED);
+			augment(gparent, parent);
 			break;
 		}
 	}
 }
+EXPORT_SYMBOL(rb_insert_augmented);
+
+static inline void dummy(struct rb_node *old, struct rb_node *new) {}
+
+void rb_insert_color(struct rb_node *node, struct rb_root *root) {
+	rb_insert_augmented(node, root, dummy);
+}
 EXPORT_SYMBOL(rb_insert_color);
 
 static void __rb_erase_color(struct rb_node *node, struct rb_node *parent,
diff --git a/lib/rbtree_test.c b/lib/rbtree_test.c
index 2dfafe4..5ace332 100644
--- a/lib/rbtree_test.c
+++ b/lib/rbtree_test.c
@@ -67,22 +67,37 @@ static void augment_callback(struct rb_node *rb, void *unused)
 	node->augmented = augment_recompute(node);
 }
 
+static void augment_rotate(struct rb_node *rb_old, struct rb_node *rb_new)
+{
+	struct test_node *old = rb_entry(rb_old, struct test_node, rb);
+	struct test_node *new = rb_entry(rb_new, struct test_node, rb);
+
+	/* Rotation doesn't change subtree's augmented value */
+	new->augmented = old->augmented;
+	old->augmented = augment_recompute(old);
+}
+
 static void insert_augmented(struct test_node *node, struct rb_root *root)
 {
-	struct rb_node **new = &root->rb_node, *parent = NULL;
+	struct rb_node **new = &root->rb_node, *rb_parent = NULL;
 	u32 key = node->key;
+	u32 val = node->val;
+	struct test_node *parent;
 
 	while (*new) {
-		parent = *new;
-		if (key < rb_entry(parent, struct test_node, rb)->key)
-			new = &parent->rb_left;
+		rb_parent = *new;
+		parent = rb_entry(rb_parent, struct test_node, rb);
+		if (parent->augmented < val)
+			parent->augmented = val;
+		if (key < parent->key)
+			new = &parent->rb.rb_left;
 		else
-			new = &parent->rb_right;
+			new = &parent->rb.rb_right;
 	}
 
-	rb_link_node(&node->rb, parent, new);
-	rb_insert_color(&node->rb, root);
-	rb_augment_insert(&node->rb, augment_callback, NULL);
+	node->augmented = val;
+	rb_link_node(&node->rb, rb_parent, new);
+	rb_insert_augmented(&node->rb, root, augment_rotate);
 }
 
 static void erase_augmented(struct test_node *node, struct rb_root *root)
-- 
1.7.7.3

[PATCH 5/6] rbtree: faster augmented erase

[Michel Lespinasse]

Add an augmented tree rotation callback to __rb_erase_color(), so that
augmented tree information can be maintained while rebalancing.

Also introduce rb_erase_augmented(), which is a version of rb_erase()
with augmented tree callbacks. We need two callbacks here: one to propagate
the augmented values up after removing a node, and one to pass up to
__rb_erase_color() to handle rebalancing.

Things are set up so that rb_erase() uses dummy do-nothing callbacks,
which get inlined and eliminated by the compiler, and also inlines the
__rb_erase_color() call so as to generate similar code than before
(once again, the compiler somehow generates smaller code than before
with all that inlining, but the speed seems to be on par). For the
augmented version rb_erase_augmented(), however, we use partial
inlining: we want rb_erase_augmented() and its augmented propagation
callback to get inlined together, but we still call into a generic
__rb_erase_color() (passing a non-inlined callback function) for the
rebalancing work. This is intended to strike a reasonable compromise
between speed and compiled code size.

Signed-off-by: Michel Lespinasse <walken@google.com>
---
 include/linux/rbtree.h          |    5 --
 include/linux/rbtree_internal.h |  131 ++++++++++++++++++++++++++++++++++++
 lib/rbtree.c                    |  139 +++++---------------------------------
 lib/rbtree_test.c               |   20 +++---
 4 files changed, 161 insertions(+), 134 deletions(-)
 create mode 100644 include/linux/rbtree_internal.h

diff --git a/include/linux/rbtree.h b/include/linux/rbtree.h
index 1364b81..bf836a2 100644
--- a/include/linux/rbtree.h
+++ b/include/linux/rbtree.h
@@ -61,11 +61,6 @@ struct rb_root {
 extern void rb_insert_color(struct rb_node *, struct rb_root *);
 extern void rb_erase(struct rb_node *, struct rb_root *);
 
-typedef void rb_augment_rotate(struct rb_node *old, struct rb_node *new);
-
-extern void rb_insert_augmented(struct rb_node *node, struct rb_root *root,
-				rb_augment_rotate *augment);
-
 typedef void (*rb_augment_f)(struct rb_node *node, void *data);
 
 extern void rb_augment_insert(struct rb_node *node,
diff --git a/include/linux/rbtree_internal.h b/include/linux/rbtree_internal.h
new file mode 100644
index 0000000..46aada0
--- /dev/null
+++ b/include/linux/rbtree_internal.h
@@ -0,0 +1,131 @@
+#ifndef _LINUX_RBTREE_INTERNAL_H
+#define _LINUX_RBTREE_INTERNAL_H
+
+#define	RB_RED		0
+#define	RB_BLACK	1
+
+#define rb_color(r)   ((r)->__rb_parent_color & 1)
+#define rb_is_red(r)   (!rb_color(r))
+#define rb_is_black(r) rb_color(r)
+
+static inline void rb_set_parent(struct rb_node *rb, struct rb_node *p)
+{
+	rb->__rb_parent_color = rb_color(rb) | (unsigned long)p;
+}
+
+static inline void rb_set_parent_color(struct rb_node *rb,
+				       struct rb_node *p, int color)
+{
+	rb->__rb_parent_color = (unsigned long)p | color;
+}
+
+static inline struct rb_node *rb_red_parent(struct rb_node *red)
+{
+	return (struct rb_node *)red->__rb_parent_color;
+}
+
+typedef void rb_augment_rotate(struct rb_node *old, struct rb_node *new);
+typedef void rb_augment_propagate(struct rb_node *node);
+
+extern void rb_insert_augmented(struct rb_node *node, struct rb_root *root,
+				rb_augment_rotate *augment);
+extern void __rb_erase_color(struct rb_node *node, struct rb_node *parent,
+			     struct rb_root *root, rb_augment_rotate *augment);
+
+static inline void
+rb_erase_augmented(struct rb_node *node, struct rb_root *root,
+		   rb_augment_propagate *augment_propagate,
+		   rb_augment_rotate *augment_rotate)
+{
+	struct rb_node *parent = rb_parent(node);
+	struct rb_node *child = node->rb_right;
+	struct rb_node *tmp = node->rb_left;
+	bool black;
+
+	if (!tmp) {
+		/* Case 1: node to erase has no more than 1 child (easy!) */
+		if (child)
+one_child:
+			rb_set_parent(child, parent);
+		if (parent) {
+			if (parent->rb_left == node)
+				parent->rb_left = child;
+			else
+				parent->rb_right = child;
+		} else
+			root->rb_node = child;
+
+		black = rb_is_black(node);
+	} else if (!child) {
+		/* Still case 1, but this time the child is node->rb_left */
+		child = tmp;
+		goto one_child;
+	} else {
+		struct rb_node *old = node;
+
+		/*
+		 * Old is the node we want to erase. It's got left and right
+		 * children, which makes things difficult. Let's find the
+		 * next node in the tree to have it fill old's position.
+		 */
+		node = node->rb_right;
+		while ((tmp = node->rb_left) != NULL)
+			node = tmp;
+
+		/* Graft node (old's successor) under old's parent. */
+		if (parent) {
+			if (parent->rb_left == old)
+				parent->rb_left = node;
+			else
+				parent->rb_right = node;
+		} else
+			root->rb_node = node;
+
+		parent = rb_parent(node);
+		black = rb_is_black(node);
+		node->__rb_parent_color = old->__rb_parent_color;
+
+		/*
+		 * Node doesn't have a left child, since it is old's successor,
+		 * so we can take old's left child and graft it under node.
+		 */
+		node->rb_left = tmp = old->rb_left;
+		rb_set_parent(tmp, node);
+
+		child = node->rb_right;
+		if (parent == old) {
+			/*
+			 * Case 2: old is node's parent (we are done!)
+			 *
+			 *    (o)          (n)
+			 *    / \          / \
+			 *  (x) (n)  ->  (x) (c)
+			 *        \
+			 *        (c)
+			 */
+			parent = node;
+		} else {
+			/*
+			 * Case 3: old is node's ancestor but not its parent
+			 *
+			 *    (o)          (n)
+			 *    / \          / \
+			 *  (x) (y)  ->  (x) (y)
+			 *      /            /
+			 *    (n)          (c)
+			 *      \
+			 *      (c)
+			 */
+			node->rb_right = tmp = old->rb_right;
+			parent->rb_left = child;
+			rb_set_parent(tmp, node);
+			if (child)
+				rb_set_parent(child, parent);
+		}
+	}
+	augment_propagate(parent);
+	if (black)
+		__rb_erase_color(child, parent, root, augment_rotate);
+}
+
+#endif	/* _LINUX_RBTREE_INTERNAL_H */
diff --git a/lib/rbtree.c b/lib/rbtree.c
index a6ae4c5..7ea3e4c 100644
--- a/lib/rbtree.c
+++ b/lib/rbtree.c
@@ -22,6 +22,7 @@
 */
 
 #include <linux/rbtree.h>
+#include <linux/rbtree_internal.h>
 #include <linux/export.h>
 
 /*
@@ -44,29 +45,6 @@
  *  parentheses and have some accompanying text comment.
  */
 
-#define	RB_RED		0
-#define	RB_BLACK	1
-
-#define rb_color(r)   ((r)->__rb_parent_color & 1)
-#define rb_is_red(r)   (!rb_color(r))
-#define rb_is_black(r) rb_color(r)
-
-static inline void rb_set_parent(struct rb_node *rb, struct rb_node *p)
-{
-	rb->__rb_parent_color = rb_color(rb) | (unsigned long)p;
-}
-
-static inline void rb_set_parent_color(struct rb_node *rb,
-				       struct rb_node *p, int color)
-{
-	rb->__rb_parent_color = (unsigned long)p | color;
-}
-
-static inline struct rb_node *rb_red_parent(struct rb_node *red)
-{
-	return (struct rb_node *)red->__rb_parent_color;
-}
-
 /*
  * Helper function for rotations:
  * - old's parent and color get assigned to new
@@ -213,15 +191,8 @@ inline void rb_insert_augmented(struct rb_node *node, struct rb_root *root,
 }
 EXPORT_SYMBOL(rb_insert_augmented);
 
-static inline void dummy(struct rb_node *old, struct rb_node *new) {}
-
-void rb_insert_color(struct rb_node *node, struct rb_root *root) {
-	rb_insert_augmented(node, root, dummy);
-}
-EXPORT_SYMBOL(rb_insert_color);
-
-static void __rb_erase_color(struct rb_node *node, struct rb_node *parent,
-			     struct rb_root *root)
+inline void __rb_erase_color(struct rb_node *node, struct rb_node *parent,
+			     struct rb_root *root, rb_augment_rotate *augment)
 {
 	struct rb_node *sibling, *tmp1, *tmp2;
 
@@ -258,6 +229,7 @@ static void __rb_erase_color(struct rb_node *node, struct rb_node *parent,
 				rb_set_parent_color(tmp1, parent, RB_BLACK);
 				__rb_rotate_set_parents(parent, sibling, root,
 							RB_RED);
+				augment(parent, sibling);
 				sibling = tmp1;
 			}
 			tmp1 = sibling->rb_right;
@@ -304,6 +276,7 @@ static void __rb_erase_color(struct rb_node *node, struct rb_node *parent,
 				if (tmp1)
 					rb_set_parent_color(tmp1, sibling,
 							    RB_BLACK);
+				augment(sibling, tmp2);
 				tmp1 = sibling;
 				sibling = tmp2;
 			}
@@ -326,6 +299,7 @@ static void __rb_erase_color(struct rb_node *node, struct rb_node *parent,
 				rb_set_parent(tmp2, parent);
 			__rb_rotate_set_parents(parent, sibling, root,
 						RB_BLACK);
+			augment(parent, sibling);
 			break;
 		} else {
 			sibling = parent->rb_left;
@@ -336,6 +310,7 @@ static void __rb_erase_color(struct rb_node *node, struct rb_node *parent,
 				rb_set_parent_color(tmp1, parent, RB_BLACK);
 				__rb_rotate_set_parents(parent, sibling, root,
 							RB_RED);
+				augment(parent, sibling);
 				sibling = tmp1;
 			}
 			tmp1 = sibling->rb_left;
@@ -356,6 +331,7 @@ static void __rb_erase_color(struct rb_node *node, struct rb_node *parent,
 				if (tmp1)
 					rb_set_parent_color(tmp1, sibling,
 							    RB_BLACK);
+				augment(sibling, tmp2);
 				tmp1 = sibling;
 				sibling = tmp2;
 			}
@@ -367,101 +343,24 @@ static void __rb_erase_color(struct rb_node *node, struct rb_node *parent,
 				rb_set_parent(tmp2, parent);
 			__rb_rotate_set_parents(parent, sibling, root,
 						RB_BLACK);
+			augment(parent, sibling);
 			break;
 		}
 	}
 }
+EXPORT_SYMBOL(__rb_erase_color);
 
-void rb_erase(struct rb_node *node, struct rb_root *root)
-{
-	struct rb_node *parent = rb_parent(node);
-	struct rb_node *child = node->rb_right;
-	struct rb_node *tmp = node->rb_left;
-	bool black;
-
-	if (!tmp) {
-		/* Case 1: node to erase has no more than 1 child (easy!) */
-		if (child)
-one_child:
-			rb_set_parent(child, parent);
-		if (parent) {
-			if (parent->rb_left == node)
-				parent->rb_left = child;
-			else
-				parent->rb_right = child;
-		} else
-			root->rb_node = child;
-
-		black = rb_is_black(node);
-	} else if (!child) {
-		/* Still case 1, but this time the child is node->rb_left */
-		child = tmp;
-		goto one_child;
-	} else {
-		struct rb_node *old = node;
+static inline void dummy_rotate(struct rb_node *old, struct rb_node *new) {}
+static inline void dummy_propagate(struct rb_node *node) {}
 
-		/*
-		 * Old is the node we want to erase. It's got left and right
-		 * children, which makes things difficult. Let's find the
-		 * next node in the tree to have it fill old's position.
-		 */
-		node = node->rb_right;
-		while ((tmp = node->rb_left) != NULL)
-			node = tmp;
-
-		/* Graft node (old's successor) under old's parent. */
-		if (parent) {
-			if (parent->rb_left == old)
-				parent->rb_left = node;
-			else
-				parent->rb_right = node;
-		} else
-			root->rb_node = node;
-
-		parent = rb_parent(node);
-		black = rb_is_black(node);
-		node->__rb_parent_color = old->__rb_parent_color;
-
-		/*
-		 * Node doesn't have a left child, since it is old's successor,
-		 * so we can take old's left child and graft it under node.
-		 */
-		node->rb_left = tmp = old->rb_left;
-		rb_set_parent(tmp, node);
+void rb_insert_color(struct rb_node *node, struct rb_root *root) {
+	rb_insert_augmented(node, root, dummy_rotate);
+}
+EXPORT_SYMBOL(rb_insert_color);
 
-		child = node->rb_right;
-		if (parent == old) {
-			/*
-			 * Case 2: old is node's parent (we are done!)
-			 *
-			 *    (o)          (n)
-			 *    / \          / \
-			 *  (x) (n)  ->  (x) (c)
-			 *        \
-			 *        (c)
-			 */
-			parent = node;
-		} else {
-			/*
-			 * Case 3: old is node's ancestor but not its parent
-			 *
-			 *    (o)          (n)
-			 *    / \          / \
-			 *  (x) (y)  ->  (x) (y)
-			 *      /            /
-			 *    (n)          (c)
-			 *      \
-			 *      (c)
-			 */
-			node->rb_right = tmp = old->rb_right;
-			parent->rb_left = child;
-			rb_set_parent(tmp, node);
-			if (child)
-				rb_set_parent(child, parent);
-		}
-	}
-	if (black)
-		__rb_erase_color(child, parent, root);
+void rb_erase(struct rb_node *node, struct rb_root *root)
+{
+	rb_erase_augmented(node, root, dummy_propagate, dummy_rotate);
 }
 EXPORT_SYMBOL(rb_erase);
 
diff --git a/lib/rbtree_test.c b/lib/rbtree_test.c
index 5ace332..fcea907 100644
--- a/lib/rbtree_test.c
+++ b/lib/rbtree_test.c
@@ -1,5 +1,6 @@
 #include <linux/module.h>
 #include <linux/rbtree.h>
+#include <linux/rbtree_internal.h>
 #include <linux/random.h>
 #include <asm/timex.h>
 
@@ -61,12 +62,6 @@ static inline u32 augment_recompute(struct test_node *node)
 	return max;
 }
 
-static void augment_callback(struct rb_node *rb, void *unused)
-{
-	struct test_node *node = rb_entry(rb, struct test_node, rb);
-	node->augmented = augment_recompute(node);
-}
-
 static void augment_rotate(struct rb_node *rb_old, struct rb_node *rb_new)
 {
 	struct test_node *old = rb_entry(rb_old, struct test_node, rb);
@@ -100,11 +95,18 @@ static void insert_augmented(struct test_node *node, struct rb_root *root)
 	rb_insert_augmented(&node->rb, root, augment_rotate);
 }
 
+static inline void augment_propagate(struct rb_node *rb)
+{
+	while (rb) {
+		struct test_node *node = rb_entry(rb, struct test_node, rb);
+		node->augmented = augment_recompute(node);
+		rb = rb_parent(&node->rb);
+	}
+}
+
 static void erase_augmented(struct test_node *node, struct rb_root *root)
 {
-	struct rb_node *deepest = rb_augment_erase_begin(&node->rb);
-	rb_erase(&node->rb, root);
-	rb_augment_erase_end(deepest, augment_callback, NULL);
+	rb_erase_augmented(&node->rb, root, augment_propagate, augment_rotate);
 }
 
 static void init(void)
-- 
1.7.7.3

[PATCH 6/6] rbtree: remove prior augmented rbtree implementation

[Michel Lespinasse]

convert arch/x86/mm/pat_rbtree.c to the proposed augmented rbtree api
and remove the old augmented rbtree implementation.

timings:
lpa36 ~168 / 282 cycles
lph71 ~168 / 271 cycles
lpk18 ~ 72 / 122 cycles

   text    data     bss     dec     hex filename
   2969       0       0    2969     b99 lib/rbtree.o

0000000000000450 l     F .text  0000000000000006 dummy_rotate
00000000000001d0 g     F .text  0000000000000272 __rb_erase_color
0000000000000480 g     F .text  000000000000001e rb_last
0000000000000510 g     F .text  0000000000000047 rb_next
0000000000000000 g     F .text  00000000000001cd rb_insert_augmented
00000000000005b0 g     F .text  0000000000000165 rb_insert_color
0000000000000560 g     F .text  0000000000000047 rb_prev
0000000000000720 g     F .text  0000000000000332 rb_erase
0000000000000460 g     F .text  000000000000001e rb_first
00000000000004a0 g     F .text  000000000000006e rb_replace_node

Signed-off-by: Michel Lespinasse <walken@google.com>
---
 arch/x86/mm/pat_rbtree.c |   52 +++++++++++++++++++++------------
 include/linux/rbtree.h   |    8 -----
 lib/rbtree.c             |   71 ----------------------------------------------
 3 files changed, 33 insertions(+), 98 deletions(-)

diff --git a/arch/x86/mm/pat_rbtree.c b/arch/x86/mm/pat_rbtree.c
index 8acaddd..fad1d58 100644
--- a/arch/x86/mm/pat_rbtree.c
+++ b/arch/x86/mm/pat_rbtree.c
@@ -13,6 +13,7 @@
 #include <linux/kernel.h>
 #include <linux/module.h>
 #include <linux/rbtree.h>
+#include <linux/rbtree_internal.h>
 #include <linux/sched.h>
 #include <linux/gfp.h>
 
@@ -54,27 +55,40 @@ static u64 get_subtree_max_end(struct rb_node *node)
 	return ret;
 }
 
-/* Update 'subtree_max_end' for a node, based on node and its children */
-static void memtype_rb_augment_cb(struct rb_node *node, void *__unused)
+static u64 compute_subtree_max_end(struct memtype *data)
 {
-	struct memtype *data;
-	u64 max_end, child_max_end;
-
-	if (!node)
-		return;
+	u64 max_end = data->end, child_max_end;
 
-	data = container_of(node, struct memtype, rb);
-	max_end = data->end;
-
-	child_max_end = get_subtree_max_end(node->rb_right);
+	child_max_end = get_subtree_max_end(data->rb.rb_right);
 	if (child_max_end > max_end)
 		max_end = child_max_end;
 
-	child_max_end = get_subtree_max_end(node->rb_left);
+	child_max_end = get_subtree_max_end(data->rb.rb_left);
 	if (child_max_end > max_end)
 		max_end = child_max_end;
 
-	data->subtree_max_end = max_end;
+	return max_end;
+}
+
+/* Update 'subtree_max_end' after tree rotation. old and new are the
+ * former and current subtree roots */
+static void memtype_rb_rotate_cb(struct rb_node *old, struct rb_node *new)
+{
+	struct memtype *old_data = container_of(old, struct memtype, rb);
+	struct memtype *new_data = container_of(new, struct memtype, rb);
+
+	new_data->subtree_max_end = old_data->subtree_max_end;
+	old_data->subtree_max_end = compute_subtree_max_end(old_data);
+}
+
+/* Update 'subtree_max_end' for node and its parents */
+static void memtype_rb_propagate_cb(struct rb_node *node)
+{
+	while (node) {
+		struct memtype *data = container_of(node, struct memtype, rb);
+		data->subtree_max_end = compute_subtree_max_end(data);
+		node = rb_parent(&data->rb);
+	}
 }
 
 /* Find the first (lowest start addr) overlapping range from rb tree */
@@ -179,15 +193,17 @@ static void memtype_rb_insert(struct rb_root *root, struct memtype *newdata)
 		struct memtype *data = container_of(*node, struct memtype, rb);
 
 		parent = *node;
+		if (data->subtree_max_end < newdata->end)
+			data->subtree_max_end = newdata->end;
 		if (newdata->start <= data->start)
 			node = &((*node)->rb_left);
 		else if (newdata->start > data->start)
 			node = &((*node)->rb_right);
 	}
 
+	newdata->subtree_max_end = newdata->end;
 	rb_link_node(&newdata->rb, parent, node);
-	rb_insert_color(&newdata->rb, root);
-	rb_augment_insert(&newdata->rb, memtype_rb_augment_cb, NULL);
+	rb_insert_augmented(&newdata->rb, root, memtype_rb_rotate_cb);
 }
 
 int rbt_memtype_check_insert(struct memtype *new, unsigned long *ret_type)
@@ -209,16 +225,14 @@ int rbt_memtype_check_insert(struct memtype *new, unsigned long *ret_type)
 
 struct memtype *rbt_memtype_erase(u64 start, u64 end)
 {
-	struct rb_node *deepest;
 	struct memtype *data;
 
 	data = memtype_rb_exact_match(&memtype_rbroot, start, end);
 	if (!data)
 		goto out;
 
-	deepest = rb_augment_erase_begin(&data->rb);
-	rb_erase(&data->rb, &memtype_rbroot);
-	rb_augment_erase_end(deepest, memtype_rb_augment_cb, NULL);
+	rb_erase_augmented(&data->rb, &memtype_rbroot,
+			   memtype_rb_propagate_cb, memtype_rb_rotate_cb);
 out:
 	return data;
 }
diff --git a/include/linux/rbtree.h b/include/linux/rbtree.h
index bf836a2..487f00b 100644
--- a/include/linux/rbtree.h
+++ b/include/linux/rbtree.h
@@ -61,14 +61,6 @@ struct rb_root {
 extern void rb_insert_color(struct rb_node *, struct rb_root *);
 extern void rb_erase(struct rb_node *, struct rb_root *);
 
-typedef void (*rb_augment_f)(struct rb_node *node, void *data);
-
-extern void rb_augment_insert(struct rb_node *node,
-			      rb_augment_f func, void *data);
-extern struct rb_node *rb_augment_erase_begin(struct rb_node *node);
-extern void rb_augment_erase_end(struct rb_node *node,
-				 rb_augment_f func, void *data);
-
 /* Find logical next and previous nodes in a tree */
 extern struct rb_node *rb_next(const struct rb_node *);
 extern struct rb_node *rb_prev(const struct rb_node *);
diff --git a/lib/rbtree.c b/lib/rbtree.c
index 7ea3e4c..8f40f9d 100644
--- a/lib/rbtree.c
+++ b/lib/rbtree.c
@@ -364,77 +364,6 @@ void rb_erase(struct rb_node *node, struct rb_root *root)
 }
 EXPORT_SYMBOL(rb_erase);
 
-static void rb_augment_path(struct rb_node *node, rb_augment_f func, void *data)
-{
-	struct rb_node *parent;
-
-up:
-	func(node, data);
-	parent = rb_parent(node);
-	if (!parent)
-		return;
-
-	if (node == parent->rb_left && parent->rb_right)
-		func(parent->rb_right, data);
-	else if (parent->rb_left)
-		func(parent->rb_left, data);
-
-	node = parent;
-	goto up;
-}
-
-/*
- * after inserting @node into the tree, update the tree to account for
- * both the new entry and any damage done by rebalance
- */
-void rb_augment_insert(struct rb_node *node, rb_augment_f func, void *data)
-{
-	if (node->rb_left)
-		node = node->rb_left;
-	else if (node->rb_right)
-		node = node->rb_right;
-
-	rb_augment_path(node, func, data);
-}
-EXPORT_SYMBOL(rb_augment_insert);
-
-/*
- * before removing the node, find the deepest node on the rebalance path
- * that will still be there after @node gets removed
- */
-struct rb_node *rb_augment_erase_begin(struct rb_node *node)
-{
-	struct rb_node *deepest;
-
-	if (!node->rb_right && !node->rb_left)
-		deepest = rb_parent(node);
-	else if (!node->rb_right)
-		deepest = node->rb_left;
-	else if (!node->rb_left)
-		deepest = node->rb_right;
-	else {
-		deepest = rb_next(node);
-		if (deepest->rb_right)
-			deepest = deepest->rb_right;
-		else if (rb_parent(deepest) != node)
-			deepest = rb_parent(deepest);
-	}
-
-	return deepest;
-}
-EXPORT_SYMBOL(rb_augment_erase_begin);
-
-/*
- * after removal, update the tree to account for the removed entry
- * and any rebalance damage.
- */
-void rb_augment_erase_end(struct rb_node *node, rb_augment_f func, void *data)
-{
-	if (node)
-		rb_augment_path(node, func, data);
-}
-EXPORT_SYMBOL(rb_augment_erase_end);
-
 /*
  * This function returns the first node (in sort order) of the tree.
  */
-- 
1.7.7.3

Re: [RFC PATCH 0/6] augmented rbtree changes

[Michel Lespinasse]

On Fri, Jul 20, 2012 at 05:31:01AM -0700, Michel Lespinasse wrote:
> Patch 5 speeds up the augmented rbtree erase. Here again we use a tree
> rotation callback during rebalancing; however we also have to propagate
> the augmented node information above nodes being erased and/or stitched,
> and I haven't found a nice enough way to do that. So for now I am proposing
> the simple-stupid way of propagating all the way to the root. More on
> this later.

So, I looked at it again and finally figured out a decent way to avoid
unnecessary propagation here. Going to resend patches 5/6 as replies to
their original postings.

> - The prio tree of all VMAs mapping a given file (struct address_space)
> could be switched to an augmented rbtree based interval tree (thus removing
> the prio tree library in favor of augmented rbtrees)

I actually have a prototype for that already. The augmented rbtree based
implementation is slightly faster than prio tree on insert/erase, and
considerably faster on lookups. However, this is with a synthetic test
exercising prio and rbtrees directly, not with a realistic workload going
through the MM layers. Do we know of situations where prio tree performance
is currently a concern ?

> As they stand, patches 3-6 don't seem to make a difference for basic rbtree
> support, and they improve my augmented rbtree insertion/erase benchmark
> by a factor of ~2.1 to ~2.3 depending on test machines.

After rewriting patches 5-6 as discussed above, augmented rbtrees are now
~2.5 - ~2.7 times faster than before this patch series.

-- 
Michel "Walken" Lespinasse
A program is never fully debugged until the last user dies.

Re: [PATCH 5/6] rbtree: faster augmented erase

[Michel Lespinasse]

Add an augmented tree rotation callback to __rb_erase_color(), so that
augmented tree information can be maintained while rebalancing.

Also introduce rb_erase_augmented(), which is a version of rb_erase()
with augmented tree callbacks. We need three callbacks here: one to
copy the subtree's augmented value after stitching in a new node as
the subtree root (rb_erase_augmented cases 2 and 3), one to propagate
the augmented values up after removing a node, and one to pass up to
__rb_erase_color() to handle rebalancing.

Things are set up so that rb_erase() uses dummy do-nothing callbacks,
which get inlined and eliminated by the compiler, and also inlines the
__rb_erase_color() call so as to generate similar code than before
(once again, the compiler somehow generates smaller code than before
with all that inlining, but the speed seems to be on par). For the
augmented version rb_erase_augmented(), however, we use partial
inlining: we want rb_erase_augmented() and its augmented copy and
propagation callbacks to get inlined together, but we still call into
a generic __rb_erase_color() (passing a non-inlined callback function)
for the rebalancing work. This is intended to strike a reasonable
compromise between speed and compiled code size.

Signed-off-by: Michel Lespinasse <walken@google.com>
---
 include/linux/rbtree.h          |    5 --
 include/linux/rbtree_internal.h |  137 +++++++++++++++++++++++++++++++++++++
 lib/rbtree.c                    |  141 ++++++---------------------------------
 lib/rbtree_test.c               |   31 ++++++---
 4 files changed, 180 insertions(+), 134 deletions(-)
 create mode 100644 include/linux/rbtree_internal.h

diff --git a/include/linux/rbtree.h b/include/linux/rbtree.h
index 1364b81..bf836a2 100644
--- a/include/linux/rbtree.h
+++ b/include/linux/rbtree.h
@@ -61,11 +61,6 @@ struct rb_root {
 extern void rb_insert_color(struct rb_node *, struct rb_root *);
 extern void rb_erase(struct rb_node *, struct rb_root *);
 
-typedef void rb_augment_rotate(struct rb_node *old, struct rb_node *new);
-
-extern void rb_insert_augmented(struct rb_node *node, struct rb_root *root,
-				rb_augment_rotate *augment);
-
 typedef void (*rb_augment_f)(struct rb_node *node, void *data);
 
 extern void rb_augment_insert(struct rb_node *node,
diff --git a/include/linux/rbtree_internal.h b/include/linux/rbtree_internal.h
new file mode 100644
index 0000000..82d2864
--- /dev/null
+++ b/include/linux/rbtree_internal.h
@@ -0,0 +1,137 @@
+#ifndef _LINUX_RBTREE_INTERNAL_H
+#define _LINUX_RBTREE_INTERNAL_H
+
+#define	RB_RED		0
+#define	RB_BLACK	1
+
+#define rb_color(r)   ((r)->__rb_parent_color & 1)
+#define rb_is_red(r)   (!rb_color(r))
+#define rb_is_black(r) rb_color(r)
+
+static inline void rb_set_parent(struct rb_node *rb, struct rb_node *p)
+{
+	rb->__rb_parent_color = rb_color(rb) | (unsigned long)p;
+}
+
+static inline void rb_set_parent_color(struct rb_node *rb,
+				       struct rb_node *p, int color)
+{
+	rb->__rb_parent_color = (unsigned long)p | color;
+}
+
+static inline struct rb_node *rb_red_parent(struct rb_node *red)
+{
+	return (struct rb_node *)red->__rb_parent_color;
+}
+
+typedef void rb_augment_rotate(struct rb_node *old, struct rb_node *new);
+typedef void rb_augment_copy(struct rb_node *old, struct rb_node *new);
+typedef void rb_augment_propagate(struct rb_node *node, struct rb_node *stop);
+
+extern void rb_insert_augmented(struct rb_node *node, struct rb_root *root,
+				rb_augment_rotate *augment);
+extern void __rb_erase_color(struct rb_node *node, struct rb_node *parent,
+			     struct rb_root *root, rb_augment_rotate *augment);
+
+static inline void
+rb_erase_augmented(struct rb_node *node, struct rb_root *root,
+		   rb_augment_copy *augment_copy,
+		   rb_augment_propagate *augment_propagate,
+		   rb_augment_rotate *augment_rotate)
+{
+	struct rb_node *parent = rb_parent(node);
+	struct rb_node *child = node->rb_right;
+	struct rb_node *tmp = node->rb_left;
+	bool black;
+
+	if (!tmp) {
+		/* Case 1: node to erase has no more than 1 child (easy!) */
+		if (child)
+one_child:
+			rb_set_parent(child, parent);
+		if (parent) {
+			if (parent->rb_left == node)
+				parent->rb_left = child;
+			else
+				parent->rb_right = child;
+		} else
+			root->rb_node = child;
+
+		tmp = parent;
+		black = rb_is_black(node);
+	} else if (!child) {
+		/* Still case 1, but this time the child is node->rb_left */
+		child = tmp;
+		goto one_child;
+	} else {
+		struct rb_node *old = node;
+
+		/*
+		 * Old is the node we want to erase. It's got left and right
+		 * children, which makes things difficult. Let's find the
+		 * next node in the tree to have it fill old's position.
+		 */
+		node = node->rb_right;
+		while ((tmp = node->rb_left) != NULL)
+			node = tmp;
+
+		/* Graft node (old's successor) under old's parent. */
+		if (parent) {
+			if (parent->rb_left == old)
+				parent->rb_left = node;
+			else
+				parent->rb_right = node;
+		} else
+			root->rb_node = node;
+
+		parent = rb_parent(node);
+		black = rb_is_black(node);
+		node->__rb_parent_color = old->__rb_parent_color;
+		augment_copy(old, node);
+
+		/*
+		 * Node doesn't have a left child, since it is old's successor,
+		 * so we can take old's left child and graft it under node.
+		 */
+		node->rb_left = tmp = old->rb_left;
+		rb_set_parent(tmp, node);
+
+		child = node->rb_right;
+		if (parent == old) {
+			/*
+			 * Case 2: old is node's parent (we are done!)
+			 *
+			 *    (o)          (n)
+			 *    / \          / \
+			 *  (x) (n)  ->  (x) (c)
+			 *        \
+			 *        (c)
+			 */
+			parent = tmp = node;
+		} else {
+			/*
+			 * Case 3: old is node's ancestor but not its parent
+			 *
+			 *    (o)          (n)
+			 *    / \          / \
+			 *  (x) (y)  ->  (x) (y)
+			 *      /            /
+			 *    (n)          (c)
+			 *      \
+			 *      (c)
+			 */
+			node->rb_right = tmp = old->rb_right;
+			parent->rb_left = child;
+			rb_set_parent(tmp, node);
+			if (child)
+				rb_set_parent(child, parent);
+			augment_propagate(parent, node);
+			tmp = node;
+		}
+	}
+	augment_propagate(tmp, NULL);
+	if (black)
+		__rb_erase_color(child, parent, root, augment_rotate);
+}
+
+#endif	/* _LINUX_RBTREE_INTERNAL_H */
diff --git a/lib/rbtree.c b/lib/rbtree.c
index a6ae4c5..90349ec 100644
--- a/lib/rbtree.c
+++ b/lib/rbtree.c
@@ -22,6 +22,7 @@
 */
 
 #include <linux/rbtree.h>
+#include <linux/rbtree_internal.h>
 #include <linux/export.h>
 
 /*
@@ -44,29 +45,6 @@
  *  parentheses and have some accompanying text comment.
  */
 
-#define	RB_RED		0
-#define	RB_BLACK	1
-
-#define rb_color(r)   ((r)->__rb_parent_color & 1)
-#define rb_is_red(r)   (!rb_color(r))
-#define rb_is_black(r) rb_color(r)
-
-static inline void rb_set_parent(struct rb_node *rb, struct rb_node *p)
-{
-	rb->__rb_parent_color = rb_color(rb) | (unsigned long)p;
-}
-
-static inline void rb_set_parent_color(struct rb_node *rb,
-				       struct rb_node *p, int color)
-{
-	rb->__rb_parent_color = (unsigned long)p | color;
-}
-
-static inline struct rb_node *rb_red_parent(struct rb_node *red)
-{
-	return (struct rb_node *)red->__rb_parent_color;
-}
-
 /*
  * Helper function for rotations:
  * - old's parent and color get assigned to new
@@ -213,15 +191,8 @@ inline void rb_insert_augmented(struct rb_node *node, struct rb_root *root,
 }
 EXPORT_SYMBOL(rb_insert_augmented);
 
-static inline void dummy(struct rb_node *old, struct rb_node *new) {}
-
-void rb_insert_color(struct rb_node *node, struct rb_root *root) {
-	rb_insert_augmented(node, root, dummy);
-}
-EXPORT_SYMBOL(rb_insert_color);
-
-static void __rb_erase_color(struct rb_node *node, struct rb_node *parent,
-			     struct rb_root *root)
+inline void __rb_erase_color(struct rb_node *node, struct rb_node *parent,
+			     struct rb_root *root, rb_augment_rotate *augment)
 {
 	struct rb_node *sibling, *tmp1, *tmp2;
 
@@ -258,6 +229,7 @@ static void __rb_erase_color(struct rb_node *node, struct rb_node *parent,
 				rb_set_parent_color(tmp1, parent, RB_BLACK);
 				__rb_rotate_set_parents(parent, sibling, root,
 							RB_RED);
+				augment(parent, sibling);
 				sibling = tmp1;
 			}
 			tmp1 = sibling->rb_right;
@@ -304,6 +276,7 @@ static void __rb_erase_color(struct rb_node *node, struct rb_node *parent,
 				if (tmp1)
 					rb_set_parent_color(tmp1, sibling,
 							    RB_BLACK);
+				augment(sibling, tmp2);
 				tmp1 = sibling;
 				sibling = tmp2;
 			}
@@ -326,6 +299,7 @@ static void __rb_erase_color(struct rb_node *node, struct rb_node *parent,
 				rb_set_parent(tmp2, parent);
 			__rb_rotate_set_parents(parent, sibling, root,
 						RB_BLACK);
+			augment(parent, sibling);
 			break;
 		} else {
 			sibling = parent->rb_left;
@@ -336,6 +310,7 @@ static void __rb_erase_color(struct rb_node *node, struct rb_node *parent,
 				rb_set_parent_color(tmp1, parent, RB_BLACK);
 				__rb_rotate_set_parents(parent, sibling, root,
 							RB_RED);
+				augment(parent, sibling);
 				sibling = tmp1;
 			}
 			tmp1 = sibling->rb_left;
@@ -356,6 +331,7 @@ static void __rb_erase_color(struct rb_node *node, struct rb_node *parent,
 				if (tmp1)
 					rb_set_parent_color(tmp1, sibling,
 							    RB_BLACK);
+				augment(sibling, tmp2);
 				tmp1 = sibling;
 				sibling = tmp2;
 			}
@@ -367,101 +343,26 @@ static void __rb_erase_color(struct rb_node *node, struct rb_node *parent,
 				rb_set_parent(tmp2, parent);
 			__rb_rotate_set_parents(parent, sibling, root,
 						RB_BLACK);
+			augment(parent, sibling);
 			break;
 		}
 	}
 }
+EXPORT_SYMBOL(__rb_erase_color);
 
-void rb_erase(struct rb_node *node, struct rb_root *root)
-{
-	struct rb_node *parent = rb_parent(node);
-	struct rb_node *child = node->rb_right;
-	struct rb_node *tmp = node->rb_left;
-	bool black;
-
-	if (!tmp) {
-		/* Case 1: node to erase has no more than 1 child (easy!) */
-		if (child)
-one_child:
-			rb_set_parent(child, parent);
-		if (parent) {
-			if (parent->rb_left == node)
-				parent->rb_left = child;
-			else
-				parent->rb_right = child;
-		} else
-			root->rb_node = child;
-
-		black = rb_is_black(node);
-	} else if (!child) {
-		/* Still case 1, but this time the child is node->rb_left */
-		child = tmp;
-		goto one_child;
-	} else {
-		struct rb_node *old = node;
+static inline void dummy_rotate(struct rb_node *old, struct rb_node *new) {}
+static inline void dummy_copy(struct rb_node *old, struct rb_node *new) {}
+static inline void dummy_propagate(struct rb_node *node, struct rb_node *stop) {}
 
-		/*
-		 * Old is the node we want to erase. It's got left and right
-		 * children, which makes things difficult. Let's find the
-		 * next node in the tree to have it fill old's position.
-		 */
-		node = node->rb_right;
-		while ((tmp = node->rb_left) != NULL)
-			node = tmp;
-
-		/* Graft node (old's successor) under old's parent. */
-		if (parent) {
-			if (parent->rb_left == old)
-				parent->rb_left = node;
-			else
-				parent->rb_right = node;
-		} else
-			root->rb_node = node;
-
-		parent = rb_parent(node);
-		black = rb_is_black(node);
-		node->__rb_parent_color = old->__rb_parent_color;
-
-		/*
-		 * Node doesn't have a left child, since it is old's successor,
-		 * so we can take old's left child and graft it under node.
-		 */
-		node->rb_left = tmp = old->rb_left;
-		rb_set_parent(tmp, node);
+void rb_insert_color(struct rb_node *node, struct rb_root *root) {
+	rb_insert_augmented(node, root, dummy_rotate);
+}
+EXPORT_SYMBOL(rb_insert_color);
 
-		child = node->rb_right;
-		if (parent == old) {
-			/*
-			 * Case 2: old is node's parent (we are done!)
-			 *
-			 *    (o)          (n)
-			 *    / \          / \
-			 *  (x) (n)  ->  (x) (c)
-			 *        \
-			 *        (c)
-			 */
-			parent = node;
-		} else {
-			/*
-			 * Case 3: old is node's ancestor but not its parent
-			 *
-			 *    (o)          (n)
-			 *    / \          / \
-			 *  (x) (y)  ->  (x) (y)
-			 *      /            /
-			 *    (n)          (c)
-			 *      \
-			 *      (c)
-			 */
-			node->rb_right = tmp = old->rb_right;
-			parent->rb_left = child;
-			rb_set_parent(tmp, node);
-			if (child)
-				rb_set_parent(child, parent);
-		}
-	}
-	if (black)
-		__rb_erase_color(child, parent, root);
+void rb_erase(struct rb_node *node, struct rb_root *root)
+{
+	rb_erase_augmented(node, root, dummy_copy, dummy_propagate,
+			   dummy_rotate);
 }
 EXPORT_SYMBOL(rb_erase);
 
diff --git a/lib/rbtree_test.c b/lib/rbtree_test.c
index 7ccc1e0..278dbfc 100644
--- a/lib/rbtree_test.c
+++ b/lib/rbtree_test.c
@@ -1,5 +1,6 @@
 #include <linux/module.h>
 #include <linux/rbtree.h>
+#include <linux/rbtree_internal.h>
 #include <linux/random.h>
 #include <asm/timex.h>
 
@@ -61,12 +62,6 @@ static inline u32 augment_recompute(struct test_node *node)
 	return max;
 }
 
-static void augment_callback(struct rb_node *rb, void *unused)
-{
-	struct test_node *node = rb_entry(rb, struct test_node, rb);
-	node->augmented = augment_recompute(node);
-}
-
 static void augment_rotate(struct rb_node *rb_old, struct rb_node *rb_new)
 {
 	struct test_node *old = rb_entry(rb_old, struct test_node, rb);
@@ -100,11 +95,29 @@ static void insert_augmented(struct test_node *node, struct rb_root *root)
 	rb_insert_augmented(&node->rb, root, augment_rotate);
 }
 
+static inline void augment_copy(struct rb_node *rb_old, struct rb_node *rb_new)
+{
+	struct test_node *old = rb_entry(rb_old, struct test_node, rb);
+	struct test_node *new = rb_entry(rb_new, struct test_node, rb);
+	new->augmented = old->augmented;
+}
+
+static inline void augment_propagate(struct rb_node *rb, struct rb_node *stop)
+{
+	while (rb != stop) {
+		struct test_node *node = rb_entry(rb, struct test_node, rb);
+		u32 augmented = augment_recompute(node);
+		if (node->augmented == augmented)
+			break;
+		node->augmented = augmented;
+		rb = rb_parent(&node->rb);
+	}
+}
+
 static void erase_augmented(struct test_node *node, struct rb_root *root)
 {
-	struct rb_node *deepest = rb_augment_erase_begin(&node->rb);
-	rb_erase(&node->rb, root);
-	rb_augment_erase_end(deepest, augment_callback, NULL);
+	rb_erase_augmented(&node->rb, root, augment_copy, augment_propagate,
+			   augment_rotate);
 }
 
 static void init(void)
-- 
1.7.7.3

Re: [PATCH 6/6] rbtree: remove prior augmented rbtree implementation

[Michel Lespinasse]

convert arch/x86/mm/pat_rbtree.c to the proposed augmented rbtree api
and remove the old augmented rbtree implementation.

Signed-off-by: Michel Lespinasse <walken@google.com>
---
 arch/x86/mm/pat_rbtree.c |   63 ++++++++++++++++++++++++++++------------
 include/linux/rbtree.h   |    8 -----
 lib/rbtree.c             |   71 ----------------------------------------------
 lib/rbtree_test.c        |   11 ++-----
 4 files changed, 47 insertions(+), 106 deletions(-)

diff --git a/arch/x86/mm/pat_rbtree.c b/arch/x86/mm/pat_rbtree.c
index 8acaddd..5b8c8b2 100644
--- a/arch/x86/mm/pat_rbtree.c
+++ b/arch/x86/mm/pat_rbtree.c
@@ -13,6 +13,7 @@
 #include <linux/kernel.h>
 #include <linux/module.h>
 #include <linux/rbtree.h>
+#include <linux/rbtree_internal.h>
 #include <linux/sched.h>
 #include <linux/gfp.h>
 
@@ -54,27 +55,51 @@ static u64 get_subtree_max_end(struct rb_node *node)
 	return ret;
 }
 
-/* Update 'subtree_max_end' for a node, based on node and its children */
-static void memtype_rb_augment_cb(struct rb_node *node, void *__unused)
+static u64 compute_subtree_max_end(struct memtype *data)
 {
-	struct memtype *data;
-	u64 max_end, child_max_end;
-
-	if (!node)
-		return;
-
-	data = container_of(node, struct memtype, rb);
-	max_end = data->end;
+	u64 max_end = data->end, child_max_end;
 
-	child_max_end = get_subtree_max_end(node->rb_right);
+	child_max_end = get_subtree_max_end(data->rb.rb_right);
 	if (child_max_end > max_end)
 		max_end = child_max_end;
 
-	child_max_end = get_subtree_max_end(node->rb_left);
+	child_max_end = get_subtree_max_end(data->rb.rb_left);
 	if (child_max_end > max_end)
 		max_end = child_max_end;
 
-	data->subtree_max_end = max_end;
+	return max_end;
+}
+
+/* Update 'subtree_max_end' after tree rotation. old and new are the
+ * former and current subtree roots */
+static void memtype_rb_rotate_cb(struct rb_node *old, struct rb_node *new)
+{
+	struct memtype *old_data = container_of(old, struct memtype, rb);
+	struct memtype *new_data = container_of(new, struct memtype, rb);
+
+	new_data->subtree_max_end = old_data->subtree_max_end;
+	old_data->subtree_max_end = compute_subtree_max_end(old_data);
+}
+
+static void memtype_rb_copy_cb(struct rb_node *old, struct rb_node *new)
+{
+	struct memtype *old_data = container_of(old, struct memtype, rb);
+	struct memtype *new_data = container_of(new, struct memtype, rb);
+
+	new_data->subtree_max_end = old_data->subtree_max_end;
+}
+
+/* Update 'subtree_max_end' for node and its parents */
+static void memtype_rb_propagate_cb(struct rb_node *node, struct rb_node *stop)
+{
+	while (node != stop) {
+		struct memtype *data = container_of(node, struct memtype, rb);
+		u64 subtree_max_end = compute_subtree_max_end(data);
+		if (data->subtree_max_end == subtree_max_end)
+			break;
+		data->subtree_max_end = subtree_max_end;
+		node = rb_parent(&data->rb);
+	}
 }
 
 /* Find the first (lowest start addr) overlapping range from rb tree */
@@ -179,15 +204,17 @@ static void memtype_rb_insert(struct rb_root *root, struct memtype *newdata)
 		struct memtype *data = container_of(*node, struct memtype, rb);
 
 		parent = *node;
+		if (data->subtree_max_end < newdata->end)
+			data->subtree_max_end = newdata->end;
 		if (newdata->start <= data->start)
 			node = &((*node)->rb_left);
 		else if (newdata->start > data->start)
 			node = &((*node)->rb_right);
 	}
 
+	newdata->subtree_max_end = newdata->end;
 	rb_link_node(&newdata->rb, parent, node);
-	rb_insert_color(&newdata->rb, root);
-	rb_augment_insert(&newdata->rb, memtype_rb_augment_cb, NULL);
+	rb_insert_augmented(&newdata->rb, root, memtype_rb_rotate_cb);
 }
 
 int rbt_memtype_check_insert(struct memtype *new, unsigned long *ret_type)
@@ -209,16 +236,14 @@ int rbt_memtype_check_insert(struct memtype *new, unsigned long *ret_type)
 
 struct memtype *rbt_memtype_erase(u64 start, u64 end)
 {
-	struct rb_node *deepest;
 	struct memtype *data;
 
 	data = memtype_rb_exact_match(&memtype_rbroot, start, end);
 	if (!data)
 		goto out;
 
-	deepest = rb_augment_erase_begin(&data->rb);
-	rb_erase(&data->rb, &memtype_rbroot);
-	rb_augment_erase_end(deepest, memtype_rb_augment_cb, NULL);
+	rb_erase_augmented(&data->rb, &memtype_rbroot, memtype_rb_copy_cb,
+			   memtype_rb_propagate_cb, memtype_rb_rotate_cb);
 out:
 	return data;
 }
diff --git a/include/linux/rbtree.h b/include/linux/rbtree.h
index bf836a2..487f00b 100644
--- a/include/linux/rbtree.h
+++ b/include/linux/rbtree.h
@@ -61,14 +61,6 @@ struct rb_root {
 extern void rb_insert_color(struct rb_node *, struct rb_root *);
 extern void rb_erase(struct rb_node *, struct rb_root *);
 
-typedef void (*rb_augment_f)(struct rb_node *node, void *data);
-
-extern void rb_augment_insert(struct rb_node *node,
-			      rb_augment_f func, void *data);
-extern struct rb_node *rb_augment_erase_begin(struct rb_node *node);
-extern void rb_augment_erase_end(struct rb_node *node,
-				 rb_augment_f func, void *data);
-
 /* Find logical next and previous nodes in a tree */
 extern struct rb_node *rb_next(const struct rb_node *);
 extern struct rb_node *rb_prev(const struct rb_node *);
diff --git a/lib/rbtree.c b/lib/rbtree.c
index 90349ec..62080ce 100644
--- a/lib/rbtree.c
+++ b/lib/rbtree.c
@@ -366,77 +366,6 @@ void rb_erase(struct rb_node *node, struct rb_root *root)
 }
 EXPORT_SYMBOL(rb_erase);
 
-static void rb_augment_path(struct rb_node *node, rb_augment_f func, void *data)
-{
-	struct rb_node *parent;
-
-up:
-	func(node, data);
-	parent = rb_parent(node);
-	if (!parent)
-		return;
-
-	if (node == parent->rb_left && parent->rb_right)
-		func(parent->rb_right, data);
-	else if (parent->rb_left)
-		func(parent->rb_left, data);
-
-	node = parent;
-	goto up;
-}
-
-/*
- * after inserting @node into the tree, update the tree to account for
- * both the new entry and any damage done by rebalance
- */
-void rb_augment_insert(struct rb_node *node, rb_augment_f func, void *data)
-{
-	if (node->rb_left)
-		node = node->rb_left;
-	else if (node->rb_right)
-		node = node->rb_right;
-
-	rb_augment_path(node, func, data);
-}
-EXPORT_SYMBOL(rb_augment_insert);
-
-/*
- * before removing the node, find the deepest node on the rebalance path
- * that will still be there after @node gets removed
- */
-struct rb_node *rb_augment_erase_begin(struct rb_node *node)
-{
-	struct rb_node *deepest;
-
-	if (!node->rb_right && !node->rb_left)
-		deepest = rb_parent(node);
-	else if (!node->rb_right)
-		deepest = node->rb_left;
-	else if (!node->rb_left)
-		deepest = node->rb_right;
-	else {
-		deepest = rb_next(node);
-		if (deepest->rb_right)
-			deepest = deepest->rb_right;
-		else if (rb_parent(deepest) != node)
-			deepest = rb_parent(deepest);
-	}
-
-	return deepest;
-}
-EXPORT_SYMBOL(rb_augment_erase_begin);
-
-/*
- * after removal, update the tree to account for the removed entry
- * and any rebalance damage.
- */
-void rb_augment_erase_end(struct rb_node *node, rb_augment_f func, void *data)
-{
-	if (node)
-		rb_augment_path(node, func, data);
-}
-EXPORT_SYMBOL(rb_augment_erase_end);
-
 /*
  * This function returns the first node (in sort order) of the tree.
  */
diff --git a/lib/rbtree_test.c b/lib/rbtree_test.c
index 278dbfc..140e173 100644
--- a/lib/rbtree_test.c
+++ b/lib/rbtree_test.c
@@ -129,16 +129,11 @@ static void init(void)
 	}
 }
 
-static bool is_red(struct rb_node *rb)
-{
-	return !(rb->__rb_parent_color & 1);
-}
-
 static int black_path_count(struct rb_node *rb)
 {
 	int count;
 	for (count = 0; rb; rb = rb_parent(rb))
-		count += !is_red(rb);
+		count += rb_is_black(rb);
 	return count;
 }
 
@@ -152,8 +147,8 @@ static void check(int nr_nodes)
 	for (rb = rb_first(&root); rb; rb = rb_next(rb)) {
 		struct test_node *node = rb_entry(rb, struct test_node, rb);
 		WARN_ON_ONCE(node->key < prev_key);
-		WARN_ON_ONCE(is_red(rb) &&
-			     (!rb_parent(rb) || is_red(rb_parent(rb))));
+		WARN_ON_ONCE(rb_is_red(rb) &&
+			     (!rb_parent(rb) || rb_is_red(rb_parent(rb))));
 		if (!count)
 			blacks = black_path_count(rb);
 		else
-- 
1.7.7.3

Re: [PATCH 1/6] rbtree: rb_erase updates and comments

[Rik van Riel]

On 07/20/2012 08:31 AM, Michel Lespinasse wrote:
> Minor updates to the rb_erase() function:
> - Reorder code to put simplest / common case (no more than 1 child) first.
> - Fetch the parent first, since it ends up being required in all 3 cases.
> - Add a few comments to illustrate case 2 (node to remove has 2 childs,
>    but one of them is the successor) and case 3 (node to remove has 2 childs,
>    successor is a left-descendant of the right child).
>
> Signed-off-by: Michel Lespinasse<walken@google.com>

Acked-by: Rik van Riel <riel@redhat.com>

-- 
All rights reversed

Re: [PATCH 2/6] rbtree: optimize fetching of sibling node

[Rik van Riel]

On 07/20/2012 08:31 AM, Michel Lespinasse wrote:
> When looking to fetch a node's sibling, we went through a sequence of:
> - check if node is the parent's left child
> - if it is, then fetch the parent's right child
>
> This can be replaced with:
> - fetch the parent's right child as an assumed sibling
> - check that node is NOT the fetched child
>
> This avoids fetching the parent's left child when node is actually
> that child. Saves a bit on code size, though it doesn't seem to make
> a large difference in speed.
>
> Signed-off-by: Michel Lespinasse<walken@google.com>

Acked-by: Rik van Riel <riel@redhat.com>


-- 
All rights reversed

Re: [PATCH 3/6] augmented rbtree test

[Rik van Riel]

On 07/20/2012 08:31 AM, Michel Lespinasse wrote:
> Signed-off-by: Michel Lespinasse <walken@google.com>

Could use a changelog.  Other than that:

Acked-by: Rik van Riel <riel@redhat.com>

Re: [PATCH 4/6] rbtree: faster augmented insert

[Rik van Riel]

On 07/20/2012 08:31 AM, Michel Lespinasse wrote:

> +++ b/lib/rbtree.c
> @@ -88,7 +88,8 @@ __rb_rotate_set_parents(struct rb_node *old, struct rb_node *new,
>   		root->rb_node = new;
>   }
>
> -void rb_insert_color(struct rb_node *node, struct rb_root *root)
> +inline void rb_insert_augmented(struct rb_node *node, struct rb_root *root,
> +				rb_augment_rotate *augment)
>   {
>   	struct rb_node *parent = rb_red_parent(node), *gparent, *tmp;
>
> @@ -152,6 +153,7 @@ void rb_insert_color(struct rb_node *node, struct rb_root *root)
>   					rb_set_parent_color(tmp, parent,
>   							    RB_BLACK);
>   				rb_set_parent_color(parent, node, RB_RED);
> +				augment(parent, node);

> +static inline void dummy(struct rb_node *old, struct rb_node *new) {}
> +
> +void rb_insert_color(struct rb_node *node, struct rb_root *root) {
> +	rb_insert_augmented(node, root, dummy);
> +}
>   EXPORT_SYMBOL(rb_insert_color);

While the above is what I would have done, the
question remains "what if the compiler decides
to not inline the function after all, and does
not remove the call to the dummy function in
rb_insert_color as a result?

Do we have some way to force inlining, so the
compiler is more likely to optimize out the
dummy call?

>   static void __rb_erase_color(struct rb_node *node, struct rb_node *parent,
> diff --git a/lib/rbtree_test.c b/lib/rbtree_test.c
> index 2dfafe4..5ace332 100644
> --- a/lib/rbtree_test.c
> +++ b/lib/rbtree_test.c
> @@ -67,22 +67,37 @@ static void augment_callback(struct rb_node *rb, void *unused)
>   	node->augmented = augment_recompute(node);
>   }
>
> +static void augment_rotate(struct rb_node *rb_old, struct rb_node *rb_new)
> +{
> +	struct test_node *old = rb_entry(rb_old, struct test_node, rb);
> +	struct test_node *new = rb_entry(rb_new, struct test_node, rb);
> +
> +	/* Rotation doesn't change subtree's augmented value */
> +	new->augmented = old->augmented;
> +	old->augmented = augment_recompute(old);
> +}

Is it worth documenting that rb_old is always the
parent of rb_new (at least, it seems to be in this
patch) ?

Re: [PATCH 5/6] rbtree: faster augmented erase

[Rik van Riel]

On 07/23/2012 09:54 PM, Michel Lespinasse wrote:
> Add an augmented tree rotation callback to __rb_erase_color(), so that
> augmented tree information can be maintained while rebalancing.
>
> Also introduce rb_erase_augmented(), which is a version of rb_erase()
> with augmented tree callbacks. We need three callbacks here: one to
> copy the subtree's augmented value after stitching in a new node as
> the subtree root (rb_erase_augmented cases 2 and 3), one to propagate
> the augmented values up after removing a node, and one to pass up to
> __rb_erase_color() to handle rebalancing.
>
> Things are set up so that rb_erase() uses dummy do-nothing callbacks,
> which get inlined and eliminated by the compiler, and also inlines the
> __rb_erase_color() call so as to generate similar code than before
> (once again, the compiler somehow generates smaller code than before
> with all that inlining, but the speed seems to be on par). For the
> augmented version rb_erase_augmented(), however, we use partial
> inlining: we want rb_erase_augmented() and its augmented copy and
> propagation callbacks to get inlined together, but we still call into
> a generic __rb_erase_color() (passing a non-inlined callback function)
> for the rebalancing work. This is intended to strike a reasonable
> compromise between speed and compiled code size.

I guess moving the inlined function to the include file
takes care of my concerns for patch 4/6...

> Signed-off-by: Michel Lespinasse <walken@google.com>

Acked-by: Rik van Riel <riel@redhat.com>

Re: [PATCH 4/6] rbtree: faster augmented insert

[Rik van Riel]

On 07/20/2012 08:31 AM, Michel Lespinasse wrote:
> Introduce rb_insert_augmented(), which is a version of rb_insert_color()
> with an added callback on tree rotations. This can be used for insertion
> into an augmented tree: the handcoded search phase must be updated to
> maintain the augmented information on insertion, and then the rbtree
> coloring/rebalancing algorithms keep it up to date.
>
> rb_insert_color() is now a special case of rb_insert_augmented() with
> a do-nothing callback. I used inlining to optimize out the callback,
> with the intent that this would generate the same code as previously
> for rb_insert_augmented(). This didn't fully work, as my compiler output
> is now *smaller* than before for that function. Speed wise, they seem
> comparable though.
>
> Signed-off-by: Michel Lespinasse <walken@google.com>

The second version of patch 5/6 takes care of my
concerns about this patch.

Acked-by: Rik van Riel <riel@redhat.com>

Re: [PATCH 6/6] rbtree: remove prior augmented rbtree implementation

[Rik van Riel]

On 07/23/2012 09:55 PM, Michel Lespinasse wrote:
> convert arch/x86/mm/pat_rbtree.c to the proposed augmented rbtree api
> and remove the old augmented rbtree implementation.
>
> Signed-off-by: Michel Lespinasse <walken@google.com>

Acked-by: Rik van Riel <riel@redhat.com>


I'm looking forward to using your new augmented rbtree
code for the rbtree based arch_get_unmapped_area code.
It should provide a nice speedup on munmap.

Re: [PATCH 4/6] rbtree: faster augmented insert

[Peter Zijlstra]

On Fri, 2012-07-20 at 05:31 -0700, Michel Lespinasse wrote:
> --- a/lib/rbtree.c
> +++ b/lib/rbtree.c
> @@ -88,7 +88,8 @@ __rb_rotate_set_parents(struct rb_node *old, struct rb_node *new,
>                 root->rb_node = new;
>  }
>  
> -void rb_insert_color(struct rb_node *node, struct rb_root *root)
> +inline void rb_insert_augmented(struct rb_node *node, struct rb_root *root,
> +                               rb_augment_rotate *augment)

Daniel probably knows best, but I would have expected something like:

__always_inline void 
__rb_insert(struct rb_node *node, struct rb_root *root,
	    const rb_augment_rotate *augment)

Where you force inline and use a const function pointer since GCC is
better with inlining them -- iirc, Daniel?

>  {
>         struct rb_node *parent = rb_red_parent(node), *gparent, *tmp;
>  
> @@ -152,6 +153,7 @@ void rb_insert_color(struct rb_node *node, struct rb_root *root)
>                                         rb_set_parent_color(tmp, parent,
>                                                             RB_BLACK);
>                                 rb_set_parent_color(parent, node, RB_RED);
> +                               augment(parent, node);

And possibly:
		if (augment)
			augment(parent, node);

>                                 parent = node;
>                                 tmp = node->rb_right;
>                         }


> +static inline void dummy(struct rb_node *old, struct rb_node *new) {}

That would obviate the need for the dummy..

> +void rb_insert_color(struct rb_node *node, struct rb_root *root) {

placed your { wrong..

> +       rb_insert_augmented(node, root, dummy);
> +}
>  EXPORT_SYMBOL(rb_insert_color); 

And use Daniel's __flatten here, like:

void rb_insert_color(struct rb_node *node, struct rb_root *root)
__flatten
{
	__rb_insert(node, root, NULL);
}
EXPORT_SYMBOL(rb_insert_color);

void rb_insert_augmented(struct rb_node *node, struct rb_root *root,
			 const rb_augment_rotate *augment) __flatten
{
	__rb_insert(node, root, augment);
}
EXPORT_SYMBOL(rb_insert_augmented);

Re: [PATCH 4/6] rbtree: faster augmented insert

[Peter Zijlstra]

On Fri, 2012-07-20 at 05:31 -0700, Michel Lespinasse wrote:
> 
> rb_insert_color() is now a special case of rb_insert_augmented() with
> a do-nothing callback. I used inlining to optimize out the callback,
> with the intent that this would generate the same code as previously
> for rb_insert_augmented(). This didn't fully work, as my compiler output
> is now *smaller* than before for that function. Speed wise, they seem
> comparable though. 

It might be good to mention which particular GCC you're using.

Re: [PATCH 5/6] rbtree: faster augmented erase

[Peter Zijlstra]

On Fri, 2012-07-20 at 05:31 -0700, Michel Lespinasse wrote:
> --- a/lib/rbtree_test.c
> +++ b/lib/rbtree_test.c
> @@ -1,5 +1,6 @@
>  #include <linux/module.h>
>  #include <linux/rbtree.h>
> +#include <linux/rbtree_internal.h> 

This confuses me.. either its internal to the rb-tree implementation and
users don't need to see it, or its not in which case huh?

Re: [PATCH 5/6] rbtree: faster augmented erase

[Peter Zijlstra]

On Fri, 2012-07-20 at 05:31 -0700, Michel Lespinasse wrote:
> +static inline void
> +rb_erase_augmented(struct rb_node *node, struct rb_root *root,
> +                  rb_augment_propagate *augment_propagate,
> +                  rb_augment_rotate *augment_rotate) 

So why put all this in a static inline in a header? As it stands
rb_erase() isn't inlined and its rather big, why would you want to
inline it for augmented callers? 

You could at least pull out the initial erase stuff into a separate
function, that way the rb_erase_augmented thing would shrink to
something like:

rb_erase_augmented(node, root)
{
	struct rb_node *parent, *child;
	bool black;

	__rb_erase(node, root, &parent, &child, &black);
	augmented_propagate(parent);
	if (black)
		__rb_erase_color(child, parent, root, augment_rotate);
}

Re: [PATCH 4/6] rbtree: faster augmented insert

[Peter Zijlstra]

On Fri, 2012-07-20 at 05:31 -0700, Michel Lespinasse wrote:
> +static void augment_rotate(struct rb_node *rb_old, struct rb_node *rb_new)
> +{
> +       struct test_node *old = rb_entry(rb_old, struct test_node, rb);
> +       struct test_node *new = rb_entry(rb_new, struct test_node, rb);
> +
> +       /* Rotation doesn't change subtree's augmented value */
> +       new->augmented = old->augmented;
> +       old->augmented = augment_recompute(old);
> +} 

> +static inline void augment_propagate(struct rb_node *rb)
> +{
> +       while (rb) {
> +               struct test_node *node = rb_entry(rb, struct test_node, rb);
> +               node->augmented = augment_recompute(node);
> +               rb = rb_parent(&node->rb);
> +       }
> +}

So why do we have to introduce these two new function pointers to pass
along when they can both be trivially expressed in the old single
augment function?

Re: [PATCH 4/6] rbtree: faster augmented insert

[Michel Lespinasse]

On Fri, Jul 27, 2012 at 12:26 PM, Peter Zijlstra <peterz@infradead.org> wrote:
> On Fri, 2012-07-20 at 05:31 -0700, Michel Lespinasse wrote:
>> --- a/lib/rbtree.c
>> +++ b/lib/rbtree.c
>> @@ -88,7 +88,8 @@ __rb_rotate_set_parents(struct rb_node *old, struct rb_node *new,
>>                 root->rb_node = new;
>>  }
>>
>> -void rb_insert_color(struct rb_node *node, struct rb_root *root)
>> +inline void rb_insert_augmented(struct rb_node *node, struct rb_root *root,
>> +                               rb_augment_rotate *augment)
>
> Daniel probably knows best, but I would have expected something like:
>
> __always_inline void
> __rb_insert(struct rb_node *node, struct rb_root *root,
>             const rb_augment_rotate *augment)
>
> Where you force inline and use a const function pointer since GCC is
> better with inlining them -- iirc, Daniel?

This hasn't been necessary with my compiler, but I can see how this
might help with older gcc versions. I really haven't investigated that
much and would be open to daniel's suggestions there.

To answer your question in the next email, we're using a gcc 4.6
variant with some local patches. TBH I don't know precisely what's in
there, however I think our compiler team makes a good job of working
with upstream so whatever changes they have are probably coming to a
future gcc version :)

>>  {
>>         struct rb_node *parent = rb_red_parent(node), *gparent, *tmp;
>>
>> @@ -152,6 +153,7 @@ void rb_insert_color(struct rb_node *node, struct rb_root *root)
>>                                         rb_set_parent_color(tmp, parent,
>>                                                             RB_BLACK);
>>                                 rb_set_parent_color(parent, node, RB_RED);
>> +                               augment(parent, node);
>
> And possibly:
>                 if (augment)
>                         augment(parent, node);
>
> That would obviate the need for the dummy..

__rb_insert() gets instanciated two times, one as rb_insert_color()
with dummy callbacks, and one as rb_insert_augmented() itself with
user-passed callbacks. Using NULL instead of dummy callbacks would
generate the same code in the rb_insert_color() instance, but not in
the version that takes user-passed callbacks (i.e. there would be an
extra check for NULL there, which we don't want).

>> +void rb_insert_color(struct rb_node *node, struct rb_root *root) {
>
> placed your { wrong..

Oops (caught myself a few times doing that, missed this one
apparently... thanks for noticing)

>> +       rb_insert_augmented(node, root, dummy);
>> +}
>>  EXPORT_SYMBOL(rb_insert_color);
>
> And use Daniel's __flatten here, like:
>
> void rb_insert_color(struct rb_node *node, struct rb_root *root)
> __flatten
> {
>         __rb_insert(node, root, NULL);
> }
> EXPORT_SYMBOL(rb_insert_color);
>
> void rb_insert_augmented(struct rb_node *node, struct rb_root *root,
>                          const rb_augment_rotate *augment) __flatten
> {
>         __rb_insert(node, root, augment);
> }
> EXPORT_SYMBOL(rb_insert_augmented);

Looks good, I'll try that and resubmit.

-- 
Michel "Walken" Lespinasse
A program is never fully debugged until the last user dies.

Re: [PATCH 4/6] rbtree: faster augmented insert

[Michel Lespinasse]

On Fri, Jul 27, 2012 at 1:04 PM, Peter Zijlstra <peterz@infradead.org> wrote:
> On Fri, 2012-07-20 at 05:31 -0700, Michel Lespinasse wrote:
>> +static void augment_rotate(struct rb_node *rb_old, struct rb_node *rb_new)
>> +{
>> +       struct test_node *old = rb_entry(rb_old, struct test_node, rb);
>> +       struct test_node *new = rb_entry(rb_new, struct test_node, rb);
>> +
>> +       /* Rotation doesn't change subtree's augmented value */
>> +       new->augmented = old->augmented;
>> +       old->augmented = augment_recompute(old);
>> +}
>
>> +static inline void augment_propagate(struct rb_node *rb)
>> +{
>> +       while (rb) {
>> +               struct test_node *node = rb_entry(rb, struct test_node, rb);
>> +               node->augmented = augment_recompute(node);
>> +               rb = rb_parent(&node->rb);
>> +       }
>> +}
>
> So why do we have to introduce these two new function pointers to pass
> along when they can both be trivially expressed in the old single
> augment function?

Its because augment_rotate() needs to be a static function that we can
take the address of and pass along as a callback to the tree
rebalancing functions, while augment_propagate() needs to be an inline
function that gets compiled within an __rb_erase() variant for a given
type of augmented rbtree.

-- 
Michel "Walken" Lespinasse
A program is never fully debugged until the last user dies.

Re: [PATCH 5/6] rbtree: faster augmented erase

[Michel Lespinasse]

On Fri, Jul 27, 2012 at 1:02 PM, Peter Zijlstra <peterz@infradead.org> wrote:
>On Fri, 2012-07-20 at 05:31 -0700, Michel Lespinasse wrote:
>> --- a/lib/rbtree_test.c
>> +++ b/lib/rbtree_test.c
>> @@ -1,5 +1,6 @@
>>  #include <linux/module.h>
>>  #include <linux/rbtree.h>
>> +#include <linux/rbtree_internal.h>
>This confuses me.. either its internal to the rb-tree implementation and
>users don't need to see it, or its not in which case huh?

So, I'm not 100% happy with this either.

What's going on is that I think it's best for users not to know about
these implementation details, and that's why I had moved these away
from include/linux/rbtree.h. However, I haven't been successful in
hiding these details from augmented rbtree users, so with my proposal,
if you want to implement some new feature using augmented rbtrees, you
do get exposed to some rbtree implementation details. This is
unfortunate but at least this exposure doesn't leak to your users -
you'd have to include linux/rbtree_internal.h only in your feature's C
file, so your users will never have to know about rbtree
implementation details.

>> +static inline void
>> +rb_erase_augmented(struct rb_node *node, struct rb_root *root,
>> +                  rb_augment_propagate *augment_propagate,
>> +                  rb_augment_rotate *augment_rotate)
>
> So why put all this in a static inline in a header? As it stands
> rb_erase() isn't inlined and its rather big, why would you want to
> inline it for augmented callers?

Just as the non-augmented rb_erase() is generated (as a non-inline
function) by merging together the rb_erase_augmented() inline function
and its dummy callbacks, I want each library that uses augmented
rbtrees to generate their own rb_erase() equivalent using their own
callbacks. The inline function in rbtree_internal.h is only to be used
as a template for generating one non-inline instance for each data
structure that uses augmented rbtrees.

> You could at least pull out the initial erase stuff into a separate
> function, that way the rb_erase_augmented thing would shrink to
> something like:
>
> rb_erase_augmented(node, root)
> {
>         struct rb_node *parent, *child;
>         bool black;
>
>         __rb_erase(node, root, &parent, &child, &black);
>         augmented_propagate(parent);
>         if (black)
>                 __rb_erase_color(child, parent, root, augment_rotate);
> }

I see that you looked at the first version of patch 5, where
augmented_propagate() still always updated all the way to the root. I
have since sent an updated version of patch 5 that does more limited
updates; however this makes it harder to do what you propose as the
callbacks now need to happen in more places than just before
__rb_erase_color().

-- 
Michel "Walken" Lespinasse
A program is never fully debugged until the last user dies.

Re: [PATCH 5/6] rbtree: faster augmented erase

[Michel Lespinasse]

On Fri, Jul 27, 2012 at 5:44 PM, Michel Lespinasse <walken@google.com> wrote:
> On Fri, Jul 27, 2012 at 1:02 PM, Peter Zijlstra <peterz@infradead.org> wrote:
>> As it stands rb_erase() isn't inlined and its rather big,
>> why would you want to inline it for augmented callers?
>
> Just as the non-augmented rb_erase() is generated (as a non-inline
> function) by merging together the rb_erase_augmented() inline function
> and its dummy callbacks, I want each library that uses augmented
> rbtrees to generate their own rb_erase() equivalent using their own
> callbacks. The inline function in rbtree_internal.h is only to be used
> as a template for generating one non-inline instance for each data
> structure that uses augmented rbtrees.

One more thing while we're talking about compiled code size. As you
noted, the non-augmented rb_erase() is pretty big. However, that size
includes the inlined rebalancing code. For the augmented erase
functions, my proposal is to the rebalancing part (rb_erase_color with
the rotate callback) will not be inlined, so as to limit the size of
the erase functions for each augmented rbtree data structure.

-- 
Michel "Walken" Lespinasse
A program is never fully debugged until the last user dies.