rbtree: handle 1-child recoloring in rb_erase() instead of rb_erase_color()

An interesting observation for rb_erase() is that when a node has exactly one child, the node must be black and the child must be red. An interesting consequence is that removing such a node can be done by simply replacing it with its child and making the child black, which we can do efficiently in rb_erase(). __rb_erase_color() then only needs to handle the no-childs case and can be modified accordingly. Signed-off-by: Michel Lespinasse <walken@google.com> Acked-by: Rik van Riel <riel@redhat.com> Cc: Peter Zijlstra <a.p.zijlstra@chello.nl> Cc: Andrea Arcangeli <aarcange@redhat.com> Cc: David Woodhouse <dwmw2@infradead.org> Signed-off-by: Andrew Morton <akpm@linux-foundation.org> Signed-off-by: Linus Torvalds <torvalds@linux-foundation.org>

rbtree: handle 1-child recoloring in rb_erase() instead of rb_erase_color()
An interesting observation for rb_erase() is that when a node has exactly one child, the node must be black and the child must be red. An interesting consequence is that removing such a node can be done by simply replacing it with its child and making the child black, which we can do efficiently in rb_erase(). __rb_erase_color() then only needs to handle the no-childs case and can be modified accordingly. Signed-off-by: Michel Lespinasse <walken@google.com> Acked-by: Rik van Riel <riel@redhat.com> Cc: Peter Zijlstra <a.p.zijlstra@chello.nl> Cc: Andrea Arcangeli <aarcange@redhat.com> Cc: David Woodhouse <dwmw2@infradead.org> Signed-off-by: Andrew Morton <akpm@linux-foundation.org> Signed-off-by: Linus Torvalds <torvalds@linux-foundation.org>
Michel Lespinasse · Linus Torvalds
1 parent 60670b8034
Showing 1 changed file with 62 additions and 43 deletions Side-by-side Diff
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
@@ -50,6 +51,11 @@
 #define rb_is_red(r)   (!rb_color(r))
 #define rb_is_black(r) rb_color(r)
  
+static inline void rb_set_black(struct rb_node *rb)
+{
+	rb->__rb_parent_color |= RB_BLACK;
+}
+
 static inline void rb_set_parent(struct rb_node *rb, struct rb_node *p)
 {
 	rb->__rb_parent_color = rb_color(rb) | (unsigned long)p;
  
  
  
@@ -214,27 +220,18 @@
 }
 EXPORT_SYMBOL(rb_insert_color);
  
-static void __rb_erase_color(struct rb_node *node, struct rb_node *parent,
-			     struct rb_root *root)
+static void __rb_erase_color(struct rb_node *parent, struct rb_root *root)
 {
-	struct rb_node *sibling, *tmp1, *tmp2;
+	struct rb_node *node = NULL, *sibling, *tmp1, *tmp2;
  
 	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.
+		 * Loop invariants:
+		 * - node is black (or NULL on first iteration)
+		 * - node is not the root (parent is not NULL)
+		 * - All leaf paths going through parent and node have a
+		 *   black node count that is 1 lower than other leaf paths.
 		 */
-		if (node && rb_is_red(node)) {
-			rb_set_parent_color(node, parent, RB_BLACK);
-			break;
-		} else if (!parent) {
-			break;
-		}
 		sibling = parent->rb_right;
 		if (node != sibling) {	/* node == parent->rb_left */
 			if (rb_is_red(sibling)) {
  
@@ -268,17 +265,22 @@
 					 *      / \           / \
 					 *     Sl  Sr        Sl  Sr
 					 *
-					 * This leaves us violating 5), so
-					 * recurse at p. If p is red, the
-					 * recursion will just flip it to black
-					 * and exit. If coming from Case 1,
-					 * p is known to be red.
+					 * This leaves us violating 5) which
+					 * can be fixed by flipping p to black
+					 * if it was red, or by recursing at p.
+					 * p is red when coming from Case 1.
 					 */
 					rb_set_parent_color(sibling, parent,
 							    RB_RED);
-					node = parent;
-					parent = rb_parent(node);
-					continue;
+					if (rb_is_red(parent))
+						rb_set_black(parent);
+					else {
+						node = parent;
+						parent = rb_parent(node);
+						if (parent)
+							continue;
+					}
+					break;
 				}
 				/*
 				 * Case 3 - right rotate at sibling
@@ -339,9 +341,15 @@
 					/* Case 2 - sibling color flip */
 					rb_set_parent_color(sibling, parent,
 							    RB_RED);
-					node = parent;
-					parent = rb_parent(node);
-					continue;
+					if (rb_is_red(parent))
+						rb_set_black(parent);
+					else {
+						node = parent;
+						parent = rb_parent(node);
+						if (parent)
+							continue;
+					}
+					break;
 				}
 				/* Case 3 - right rotate at sibling */
 				sibling->rb_right = tmp1 = tmp2->rb_left;
  
  
  
  
@@ -369,23 +377,31 @@
 void rb_erase(struct rb_node *node, struct rb_root *root)
 {
 	struct rb_node *child = node->rb_right, *tmp = node->rb_left;
-	struct rb_node *parent;
-	int color;
+	struct rb_node *parent, *rebalance;
  
 	if (!tmp) {
-	case1:
-		/* Case 1: node to erase has no more than 1 child (easy!) */
+		/*
+		 * Case 1: node to erase has no more than 1 child (easy!)
+		 *
+		 * Note that if there is one child it must be red due to 5)
+		 * and node must be black due to 4). We adjust colors locally
+		 * so as to bypass __rb_erase_color() later on.
+		 */
  
 		parent = rb_parent(node);
-		color = rb_color(node);
-
-		if (child)
-			rb_set_parent(child, parent);
 		__rb_change_child(node, child, parent, root);
+		if (child) {
+			rb_set_parent_color(child, parent, RB_BLACK);
+			rebalance = NULL;
+		} else {
+			rebalance = rb_is_black(node) ? parent : NULL;
+		}
 	} else if (!child) {
 		/* Still case 1, but this time the child is node->rb_left */
-		child = tmp;
-		goto case1;
+		parent = rb_parent(node);
+		__rb_change_child(node, tmp, parent, root);
+		rb_set_parent_color(tmp, parent, RB_BLACK);
+		rebalance = NULL;
 	} else {
 		struct rb_node *old = node, *left;
  
  
  
  
@@ -397,26 +413,29 @@
  
 		child = node->rb_right;
 		parent = rb_parent(node);
-		color = rb_color(node);
  
 		if (parent == old) {
 			parent = node;
 		} else {
-			if (child)
-				rb_set_parent(child, parent);
 			parent->rb_left = child;
  
 			node->rb_right = old->rb_right;
 			rb_set_parent(old->rb_right, node);
 		}
  
+		if (child) {
+			rb_set_parent_color(child, parent, RB_BLACK);
+			rebalance = NULL;
+		} else {
+			rebalance = rb_is_black(node) ? parent : NULL;
+		}
 		node->__rb_parent_color = old->__rb_parent_color;
 		node->rb_left = old->rb_left;
 		rb_set_parent(old->rb_left, node);
 	}
  
-	if (color == RB_BLACK)
-		__rb_erase_color(child, parent, root);
+	if (rebalance)
+		__rb_erase_color(rebalance, root);
 }
 EXPORT_SYMBOL(rb_erase);
...	...	@@ -2,7 +2,8 @@
2	2	Red Black Trees
3	3	(C) 1999 Andrea Arcangeli <andrea@suse.de>
4	4	(C) 2002 David Woodhouse <dwmw2@infradead.org>
5		-
	5	+ (C) 2012 Michel Lespinasse <walken@google.com>
	6	+
6	7	This program is free software; you can redistribute it and/or modify
7	8	it under the terms of the GNU General Public License as published by
8	9	the Free Software Foundation; either version 2 of the License, or
...	...	@@ -50,6 +51,11 @@
50	51	#define rb_is_red(r) (!rb_color(r))
51	52	#define rb_is_black(r) rb_color(r)
52	53
	54	+static inline void rb_set_black(struct rb_node *rb)
	55	+{
	56	+ rb->__rb_parent_color \|= RB_BLACK;
	57	+}
	58	+
53	59	static inline void rb_set_parent(struct rb_node rb, struct rb_node p)
54	60	{
55	61	rb->__rb_parent_color = rb_color(rb) \| (unsigned long)p;
56	62
57	63
58	64
...	...	@@ -214,27 +220,18 @@
214	220	}
215	221	EXPORT_SYMBOL(rb_insert_color);
216	222
217		-static void __rb_erase_color(struct rb_node node, struct rb_node parent,
218		- struct rb_root *root)
	223	+static void __rb_erase_color(struct rb_node parent, struct rb_root root)
219	224	{
220		- struct rb_node sibling, tmp1, *tmp2;
	225	+ struct rb_node node = NULL, sibling, tmp1, tmp2;
221	226
222	227	while (true) {
223	228	/*
224		- * Loop invariant: all leaf paths going through node have a
225		- * black node count that is 1 lower than other leaf paths.
226		- *
227		- * If node is red, we can flip it to black to adjust.
228		- * If node is the root, all leaf paths go through it.
229		- * Otherwise, we need to adjust the tree through color flips
230		- * and tree rotations as per one of the 4 cases below.
	229	+ * Loop invariants:
	230	+ * - node is black (or NULL on first iteration)
	231	+ * - node is not the root (parent is not NULL)
	232	+ * - All leaf paths going through parent and node have a
	233	+ * black node count that is 1 lower than other leaf paths.
231	234	*/
232		- if (node && rb_is_red(node)) {
233		- rb_set_parent_color(node, parent, RB_BLACK);
234		- break;
235		- } else if (!parent) {
236		- break;
237		- }
238	235	sibling = parent->rb_right;
239	236	if (node != sibling) { /* node == parent->rb_left */
240	237	if (rb_is_red(sibling)) {
241	238
...	...	@@ -268,17 +265,22 @@
268	265	* / \ / \
269	266	* Sl Sr Sl Sr
270	267	*
271		- * This leaves us violating 5), so
272		- * recurse at p. If p is red, the
273		- * recursion will just flip it to black
274		- * and exit. If coming from Case 1,
275		- * p is known to be red.
	268	+ * This leaves us violating 5) which
	269	+ * can be fixed by flipping p to black
	270	+ * if it was red, or by recursing at p.
	271	+ * p is red when coming from Case 1.
276	272	*/
277	273	rb_set_parent_color(sibling, parent,
278	274	RB_RED);
279		- node = parent;
280		- parent = rb_parent(node);
281		- continue;
	275	+ if (rb_is_red(parent))
	276	+ rb_set_black(parent);
	277	+ else {
	278	+ node = parent;
	279	+ parent = rb_parent(node);
	280	+ if (parent)
	281	+ continue;
	282	+ }
	283	+ break;
282	284	}
283	285	/*
284	286	* Case 3 - right rotate at sibling
...	...	@@ -339,9 +341,15 @@
339	341	/* Case 2 - sibling color flip */
340	342	rb_set_parent_color(sibling, parent,
341	343	RB_RED);
342		- node = parent;
343		- parent = rb_parent(node);
344		- continue;
	344	+ if (rb_is_red(parent))
	345	+ rb_set_black(parent);
	346	+ else {
	347	+ node = parent;
	348	+ parent = rb_parent(node);
	349	+ if (parent)
	350	+ continue;
	351	+ }
	352	+ break;
345	353	}
346	354	/* Case 3 - right rotate at sibling */
347	355	sibling->rb_right = tmp1 = tmp2->rb_left;
348	356
349	357
350	358
351	359
...	...	@@ -369,23 +377,31 @@
369	377	void rb_erase(struct rb_node node, struct rb_root root)
370	378	{
371	379	struct rb_node child = node->rb_right, tmp = node->rb_left;
372		- struct rb_node *parent;
373		- int color;
	380	+ struct rb_node parent, rebalance;
374	381
375	382	if (!tmp) {
376		- case1:
377		- /* Case 1: node to erase has no more than 1 child (easy!) */
	383	+ /*
	384	+ * Case 1: node to erase has no more than 1 child (easy!)
	385	+ *
	386	+ * Note that if there is one child it must be red due to 5)
	387	+ * and node must be black due to 4). We adjust colors locally
	388	+ * so as to bypass __rb_erase_color() later on.
	389	+ */
378	390
379	391	parent = rb_parent(node);
380		- color = rb_color(node);
381		-
382		- if (child)
383		- rb_set_parent(child, parent);
384	392	__rb_change_child(node, child, parent, root);
	393	+ if (child) {
	394	+ rb_set_parent_color(child, parent, RB_BLACK);
	395	+ rebalance = NULL;
	396	+ } else {
	397	+ rebalance = rb_is_black(node) ? parent : NULL;
	398	+ }
385	399	} else if (!child) {
386	400	/* Still case 1, but this time the child is node->rb_left */
387		- child = tmp;
388		- goto case1;
	401	+ parent = rb_parent(node);
	402	+ __rb_change_child(node, tmp, parent, root);
	403	+ rb_set_parent_color(tmp, parent, RB_BLACK);
	404	+ rebalance = NULL;
389	405	} else {
390	406	struct rb_node old = node, left;
391	407
392	408
393	409
394	410
...	...	@@ -397,26 +413,29 @@
397	413
398	414	child = node->rb_right;
399	415	parent = rb_parent(node);
400		- color = rb_color(node);
401	416
402	417	if (parent == old) {
403	418	parent = node;
404	419	} else {
405		- if (child)
406		- rb_set_parent(child, parent);
407	420	parent->rb_left = child;
408	421
409	422	node->rb_right = old->rb_right;
410	423	rb_set_parent(old->rb_right, node);
411	424	}
412	425
	426	+ if (child) {
	427	+ rb_set_parent_color(child, parent, RB_BLACK);
	428	+ rebalance = NULL;
	429	+ } else {
	430	+ rebalance = rb_is_black(node) ? parent : NULL;
	431	+ }
413	432	node->__rb_parent_color = old->__rb_parent_color;
414	433	node->rb_left = old->rb_left;
415	434	rb_set_parent(old->rb_left, node);
416	435	}
417	436
418		- if (color == RB_BLACK)
419		- __rb_erase_color(child, parent, root);
	437	+ if (rebalance)
	438	+ __rb_erase_color(rebalance, root);
420	439	}
421	440	EXPORT_SYMBOL(rb_erase);
422	441