[libgee] Fix move_red_right and move_red_left in the Tree implementations

[Didier Villevalois <dvillevalois src gnome org>]

commit 0eab459be015030d9edd6cf4c2a556e405e56bd7
Author: Maciej Piechotka <uzytkownik2 gmail com>
Date:   Sat Sep 19 22:59:33 2009 +0200

    Fix move_red_right and move_red_left in the Tree implementations
    
    Fixes bug 595703.
    
    Also:
     - For simplification move to 2-3 LLRB tree.
     - Add CONSTSTENCY_CHECK and DEBUG options

 gee/treemap.vala       |   21 ++++-----
 gee/treeset.vala       |  111 ++++++++++++++++++++++++++++++++++++++++--------
 tests/testtreeset.vala |   49 +++++++++++++++++++++-
 3 files changed, 151 insertions(+), 30 deletions(-)
---
diff --git a/gee/treemap.vala b/gee/treemap.vala
index 156ec8f..cc61d36 100644
--- a/gee/treemap.vala
+++ b/gee/treemap.vala
@@ -200,10 +200,6 @@ public class Gee.TreeMap<K,V> : Gee.AbstractMap<K,V> {
 			_size++;
 		}
 
-		if (is_red (node.left) && is_red (node.right)) {
-			node.flip ();
-		}
-
 		int cmp = key_compare_func (key, node.key);
 		if (cmp == 0) {
 			node.value = value;
@@ -226,7 +222,7 @@ public class Gee.TreeMap<K,V> : Gee.AbstractMap<K,V> {
 
 	private void move_red_left (ref Node<K, V> root) {
 		root.flip ();
-		if (root.right != null && is_red (root.right.left)) {
+		if (is_red (root.right.left)) {
 			rotate_right (ref root.right);
 			rotate_left (ref root);
 			root.flip ();
@@ -235,8 +231,8 @@ public class Gee.TreeMap<K,V> : Gee.AbstractMap<K,V> {
 
 	private void move_red_right (ref Node<K, V> root) {
 		root.flip ();
-		if (root.left != null && is_red (root.left.left)) {
-			rotate_right (ref root.right);
+		if (is_red (root.left.left)) {
+			rotate_right (ref root);
 			root.flip ();
 		}
 	}
@@ -284,7 +280,7 @@ public class Gee.TreeMap<K,V> : Gee.AbstractMap<K,V> {
 			if (left == null) {
 				return false;
 			}
-			if (is_black (left) && is_black (left.left)) {
+			if (node.left != null && is_black (left) && is_black (left.left)) {
 				move_red_left (ref node);
 			}
 			bool r = remove_from_node (ref node.left, key, out value);
@@ -295,12 +291,12 @@ public class Gee.TreeMap<K,V> : Gee.AbstractMap<K,V> {
 				rotate_right (ref node);
 			}
 
-			weak Node<K,V> r = node.right;
+			weak Node<K,V>? r = node.right;
 			if (key_compare_func (key, node.key) == 0 && r == null) {
 				fix_removal (ref node, null, out value);
 				return true;
 			}
-			if (r == null || (is_black (r) && is_black (r.left))) {
+			if (is_black (r) && r != null && is_black (r.left)) {
 				move_red_right (ref node);
 			}
 			if (key_compare_func (key, node.key) == 0) {
@@ -323,6 +319,9 @@ public class Gee.TreeMap<K,V> : Gee.AbstractMap<K,V> {
 		if (is_red (node.left) && is_red (node.left.left)) {
 			rotate_right (ref node);
 		}
+		if (is_red (node.left) && is_red (node.right)) {
+			node.flip ();
+		}
 	}
 
 	/**
@@ -389,7 +388,7 @@ public class Gee.TreeMap<K,V> : Gee.AbstractMap<K,V> {
 		}
 
 		public void flip () {
-			color.flip ();
+			color = color.flip ();
 			if (left != null) {
 				left.color = left.color.flip ();
 			}
diff --git a/gee/treeset.vala b/gee/treeset.vala
index f8836af..3b22504 100644
--- a/gee/treeset.vala
+++ b/gee/treeset.vala
@@ -78,42 +78,61 @@ public class Gee.TreeSet<G> : AbstractSet<G> {
 		return false;
 	}
 
-	private void rotate_right (ref Node<G> root) {
+	private inline void rotate_right (ref Node<G> root) {
 		Node<G> pivot = (owned) root.left;
 		pivot.color = root.color;
 		root.color = Node.Color.RED;
 		root.left = (owned) pivot.right;
 		pivot.right = (owned) root;
 		root = (owned) pivot;
+#if DEBUG
+		stdout.printf (dump ("after rotate right on %s".printf ((string)root.right.key)));
+#endif
 	}
 
-	private void rotate_left (ref Node<G> root) {
+	private inline void rotate_left (ref Node<G> root) {
 		Node<G> pivot = (owned) root.right;
 		pivot.color = root.color;
 		root.color = Node.Color.RED;
 		root.right = (owned) pivot.left;
 		pivot.left = (owned) root;
 		root = (owned) pivot;
+#if DEBUG
+		stdout.printf (dump ("after rotate left on %s".printf ((string)root.left.key)));
+#endif
 	}
 
-	private bool is_red (Node<G>? n) {
+	private inline bool is_red (Node<G>? n) {
 		return n != null && n.color == Node.Color.RED;
 	}
 
-	private bool is_black (Node<G>? n) {
+	private inline bool is_black (Node<G>? n) {
 		return n == null || n.color == Node.Color.BLACK;
 	}
 
-	private void fix_up (ref Node<G> node) {
+	private inline void fix_up (ref Node<G> node) {
+#if DEBUG
+		var n = (string)node.key;
+#endif
 		if (is_black (node.left) && is_red (node.right)) {
 			rotate_left (ref node);
 		}
 		if (is_red (node.left) && is_red (node.left.left)) {
 			rotate_right (ref node);
 		}
+		if (is_red (node.left) && is_red (node.right)) {
+			node.flip ();
+		}
+#if DEBUG
+		stdout.printf (dump ("after fix up on %s".printf (n)));
+#endif
 	}
 
 	private bool add_to_node (ref Node<G>? node, owned G item, Node<G>? prev, Node<G>? next) {
+#if DEBUG
+		if (node != null)
+			stdout.printf ("Adding %s to %s\n".printf ((string) item, (string) node.key));
+#endif
 		if (node == null) {
 			node = new Node<G> ((owned) item, prev, next);
 			if (prev == null) {
@@ -126,10 +145,6 @@ public class Gee.TreeSet<G> : AbstractSet<G> {
 			return true;
 		}
 
-		if (is_red (node.left) && is_red (node.right)) {
-			node.flip ();
-		}
-
 		int cmp = compare_func (item, node.key);
 		if (cmp == 0) {
 			fix_up (ref node);
@@ -151,30 +166,48 @@ public class Gee.TreeSet<G> : AbstractSet<G> {
 	 * If the element already exists in the set it will not be added twice.
 	 */
 	public override bool add (G item) {
+#if CONSISTENCY_CHECKS
+		check ();
+#endif
 		bool r = add_to_node (ref root, item, null, null);
 		root.color = Node.Color.BLACK;
+#if CONSISTENCY_CHECKS
+		check ();
+#endif
 		stamp++;
 		return r;
 	}
 
-	private void move_red_left (ref Node<G> root) {
+	private inline void move_red_left (ref Node<G> root) {
+#if DEBUG
+		var n = (string)root.key;
+#endif
 		root.flip ();
-		if (root.right != null && is_red (root.right.left)) {
+		if (is_red (root.right.left)) {
 			rotate_right (ref root.right);
 			rotate_left (ref root);
 			root.flip ();
 		}
+#if DEBUG
+		stdout.printf (dump ("after red left on %s".printf (n)));
+#endif
 	}
 
-	private void move_red_right (ref Node<G> root) {
+	private inline void move_red_right (ref Node<G> root) {
+#if DEBUG
+		var n = (string)root.key;
+#endif
 		root.flip ();
-		if (root.left != null && is_red (root.left.left)) {
-			rotate_right (ref root.right);
+		if (is_red (root.left.left)) {
+			rotate_right (ref root);
 			root.flip ();
 		}
+#if DEBUG
+		stdout.printf (dump ("after red right on %s".printf (n)));
+#endif
 	}
 
-	private void fix_removal (ref Node<G> node, out G? key = null) {
+	private inline void fix_removal (ref Node<G> node, out G? key = null) {
 		Node<G> n = (owned) node;
 		if (&key != null)
 			key = (owned) n.key;
@@ -208,6 +241,9 @@ public class Gee.TreeSet<G> : AbstractSet<G> {
 	}
 
 	private bool remove_from_node (ref Node<G>? node, G item) {
+#if DEBUG
+		stdout.printf ("Removing %s from %s\n", (string)item, node != null ? (string)node.key : null);
+#endif
 		if (node == null) {
 			return false;
 		} else if (compare_func (item, node.key) < 0) {
@@ -226,12 +262,12 @@ public class Gee.TreeSet<G> : AbstractSet<G> {
 				rotate_right (ref node);
 			}
 
-			weak Node<G> r = node.right;
+			weak Node<G>? r = node.right;
 			if (compare_func (item, node.key) == 0 && r == null) {
 				fix_removal (ref node, null);
 				return true;
 			}
-			if (r == null || (is_black (r) && is_black (r.left))) {
+			if (is_black (r) && r != null && is_black (r.left)) {
 				move_red_right (ref node);
 			}
 			if (compare_func (item, node.key) == 0) {
@@ -250,10 +286,16 @@ public class Gee.TreeSet<G> : AbstractSet<G> {
 	 * @inheritDoc
 	 */
 	public override bool remove (G item) {
+#if CONSISTENCY_CHECKS
+		check ();
+#endif
 		bool b = remove_from_node (ref root, item);
 		if (root != null) {
 			root.color = Node.Color.BLACK;
 		}
+#if CONSISTENCY_CHECKS
+		check ();
+#endif
 		stamp++;
 		return b;
 	}
@@ -281,6 +323,39 @@ public class Gee.TreeSet<G> : AbstractSet<G> {
 		return new Iterator<G> (this);
 	}
 
+#if CONSISTENCY_CHECKS
+	public inline void check () {
+		check_subtree (root);
+		stdout.printf ("%s\n", dump ());
+	}
+
+	private inline uint check_subtree (Node<G>? node) {
+		if (node == null)
+			return 0;
+		assert (!(is_black (node.left) && is_red (node.right))); // Check left-leaning
+		assert (!(is_red (node) && is_red (node.left))); // Check red property
+		uint l = check_subtree (node.left);
+		uint r = check_subtree (node.right);
+		assert (l == r);
+		return l + (node.color == Node.Color.BLACK ? 1 : 0);
+	}
+#endif
+#if DEBUG
+	public string dump (string? when = null) {
+		return "TreeSet dump%s:\n%s".printf (when == null ? "" : (" " + when), dump_node (root));
+	}
+
+	private inline string dump_node (Node<G>? node, uint depth = 0) {
+		if (node != null)
+			return dump_node (node.left, depth + 1) +
+			       "%s%s%p(%s)\033[0m\n".printf (string.nfill (depth, ' '),
+			                                   node.color == Node.Color.RED ? "\033[0;31m" : "",
+			                                   node, (string)node.key) +
+			       dump_node (node.right, depth + 1);
+		return "";
+	}
+#endif
+
 	[Compact]
 	private class Node<G> {
 		public enum Color {
@@ -310,7 +385,7 @@ public class Gee.TreeSet<G> : AbstractSet<G> {
 		}
 
 		public void flip () {
-			color.flip ();
+			color = color.flip ();
 			if (left != null) {
 				left.color = left.color.flip ();
 			}
diff --git a/tests/testtreeset.vala b/tests/testtreeset.vala
index 038b341..bf0a67a 100644
--- a/tests/testtreeset.vala
+++ b/tests/testtreeset.vala
@@ -30,6 +30,7 @@ public class TreeSetTests : SetTests {
 		add_test ("[TreeSet] selected functions", test_selected_functions);
 		add_test ("[TreeSet] GObject properties", test_gobject_properties);
 		add_test ("[TreeSet] ordering", test_ordering);
+		add_test ("[TreeSet] add and remove", test_add_remove);
 	}
 
 	public override void set_up () {
@@ -50,7 +51,7 @@ public class TreeSetTests : SetTests {
 		assert (test_set.compare_func == (CompareFunc) strcmp);
 	}
 
-	public new void test_gobject_properties() {
+	public new void test_gobject_properties () {
 		var test_set = test_collection as TreeSet<string>;
 
 		// Check the set exists
@@ -109,4 +110,50 @@ public class TreeSetTests : SetTests {
 		assert (iterator.get () == "two");
 		assert (iterator.next () == false);
 	}
+
+	public new void test_add_remove () {
+		var test_set = test_collection as TreeSet<string>;
+
+		// Check the set exists
+		assert (test_set != null);
+
+		base.test_remove_all ();
+
+		var to_add = new string[] {
+			"3", "10", "5", "6", "13", "8", "12", "11", "1", "0", "9", "2", "14", "7", "15", "4"
+		};
+		var to_remove = new string[] {
+			"11", "13", "1", "12", "4", "0", "2", "5", "6", "3", "14", "10", "7", "15", "8", "9"
+		};
+
+		foreach (var s in to_add) {
+			assert (!test_set.contains (s));
+			assert (test_set.add (s));
+			assert (test_set.contains (s));
+		}
+		foreach (var s in to_remove) {
+			assert (test_set.contains (s));
+			assert (test_set.remove (s));
+
+			assert (!test_set.contains (s));
+		}
+
+		to_add = new string[] {
+			"2", "9", "13", "7", "12", "14", "8", "1", "5", "6", "11", "15", "3", "10", "0", "4"
+		};
+		to_remove = new string[] {
+			"11", "10", "14", "6", "13", "4", "3", "15", "8", "5", "7", "0", "12", "2", "9", "1"
+		};
+
+		foreach (var s in to_add) {
+			assert (!test_set.contains (s));
+			assert (test_set.add (s));
+			assert (test_set.contains (s));
+		}
+		foreach (var s in to_remove) {
+			assert (test_set.contains (s));
+			assert (test_set.remove (s));
+			assert (!test_set.contains (s));
+		}
+	}
 }