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

diff --git a/gee/treemap.vala b/gee/treemap.vala
index 156ec8f..cc61d36 100644 (file)
--- 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 (file)
--- 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 (file)
--- 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));
+               }
+       }
 }