Implement erase in C++ immutable maps (#1158)

* Add SortedMap::min
* Add SortedMap::erase

7 files changed, 338 insertions, 107 deletions

diff --git a/Firestore/core/src/firebase/firestore/immutable/array_sorted_map.h b/Firestore/core/src/firebase/firestore/immutable/array_sorted_map.h
index 6ffe017..487a25a 100644
--- a/Firestore/core/src/firebase/firestore/immutable/array_sorted_map.h
+++ b/Firestore/core/src/firebase/firestore/immutable/array_sorted_map.h
@@ -258,6 +258,18 @@ class ArraySortedMap : public SortedMapBase {
     return found == end() ? npos : static_cast<size_type>(found - begin());
   }
 
+  const_iterator min() const {
+    return begin();
+  }
+
+  const_iterator max() const {
+    if (empty()) {
+      return end();
+    }
+
+    return end() - 1;
+  }
+
   /**
    * Returns an iterator pointing to the first entry in the map. If there are
    * no entries in the map, begin() == end().
diff --git a/Firestore/core/src/firebase/firestore/immutable/llrb_node.h b/Firestore/core/src/firebase/firestore/immutable/llrb_node.h
index d2a2227..80c2d86 100644
--- a/Firestore/core/src/firebase/firestore/immutable/llrb_node.h
+++ b/Firestore/core/src/firebase/firestore/immutable/llrb_node.h
@@ -97,6 +97,25 @@ class LlrbNode : public SortedMapBase {
                   const V& value,
                   const Comparator& comparator) const;
 
+  template <typename Comparator>
+  LlrbNode erase(const K& key, const Comparator& comparator) const;
+
+  const LlrbNode& min() const {
+    const LlrbNode* node = this;
+    while (!node->left().empty()) {
+      node = &node->left();
+    }
+    return *node;
+  }
+
+  const LlrbNode& max() const {
+    const LlrbNode* node = this;
+    while (!node->right().empty()) {
+      node = &node->right();
+    }
+    return *node;
+  }
+
  private:
   struct Rep {
     Rep(value_type&& entry,
@@ -198,12 +217,20 @@ class LlrbNode : public SortedMapBase {
                        const V& value,
                        const Comparator& comparator) const;
 
+  template <typename Comparator>
+  LlrbNode InnerErase(const K& key, const Comparator& comparator) const;
+
   void FixUp();
+  void FixRootColor();
 
   void RotateLeft();
   void RotateRight();
   void FlipColor();
 
+  void RemoveMin();
+  void MoveRedLeft();
+  void MoveRedRight();
+
   size_type OppositeColor() const noexcept {
     return rep_->color_ == Color::Red ? Color::Black : Color::Red;
   }
@@ -217,10 +244,7 @@ LlrbNode<K, V> LlrbNode<K, V>::insert(const K& key,
                                       const V& value,
                                       const Comparator& comparator) const {
   LlrbNode root = InnerInsert(key, value, comparator);
-  // The root must always be black
-  if (root.red()) {
-    root.rep_->color_ = Color::Black;
-  }
+  root.FixRootColor();
   return root;
 }
 
@@ -259,6 +283,64 @@ LlrbNode<K, V> LlrbNode<K, V>::InnerInsert(const K& key,
 }
 
 template <typename K, typename V>
+template <typename Comparator>
+LlrbNode<K, V> LlrbNode<K, V>::erase(const K& key,
+                                     const Comparator& comparator) const {
+  LlrbNode root = InnerErase(key, comparator);
+  root.FixRootColor();
+  return root;
+}
+
+template <typename K, typename V>
+template <typename Comparator>
+LlrbNode<K, V> LlrbNode<K, V>::InnerErase(const K& key,
+                                          const Comparator& comparator) const {
+  if (empty()) {
+    // Empty node already frozen
+    return LlrbNode{};
+  }
+
+  LlrbNode n = Clone();
+
+  bool descending = comparator(key, n.key());
+  if (descending) {
+    if (!n.left().empty() && !n.left().red() && !n.left().left().red()) {
+      n.MoveRedLeft();
+    }
+    n.set_left(n.left().InnerErase(key, comparator));
+
+  } else {
+    if (n.left().red()) {
+      n.RotateRight();
+    }
+
+    if (!n.right().empty() && !n.right().red() && !n.right().left().red()) {
+      n.MoveRedRight();
+    }
+
+    if (util::Compare(key, n.key(), comparator) ==
+        util::ComparisonResult::Same) {
+      if (n.right().empty()) {
+        return LlrbNode{};
+
+      } else {
+        // Move the minimum node from the right subtree in place of this node.
+        LlrbNode smallest = n.right().min();
+        LlrbNode new_right = n.right().Clone();
+        new_right.RemoveMin();
+
+        n.set_entry(smallest.entry());
+        n.set_right(std::move(new_right));
+      }
+    } else {
+      n.set_right(n.right().InnerErase(key, comparator));
+    }
+  }
+  n.FixUp();
+  return n;
+}
+
+template <typename K, typename V>
 void LlrbNode<K, V>::FixUp() {
   set_size(left().size() + 1 + right().size());
 
@@ -273,6 +355,62 @@ void LlrbNode<K, V>::FixUp() {
   }
 }
 
+/**
+ * Fixes the root node so its color is always black.
+ *
+ * This change is safe because a red `root` must be a new root:
+ *   * If the key is not found, InnerErase returns an existing root, but
+ *     existing roots will already be black (by the invariant).
+ *   * If the key is found, InnerErase returns a new root, which is safe to
+ *     modify.
+ */
+template <typename K, typename V>
+void LlrbNode<K, V>::FixRootColor() {
+  if (red()) {
+    rep_->color_ = Color::Black;
+  }
+}
+
+template <typename K, typename V>
+void LlrbNode<K, V>::RemoveMin() {
+  // If the left node is empty then the right node must be empty (because the
+  // tree is left-leaning) and this node must be the minimum.
+  if (left().empty()) {
+    *this = LlrbNode{};
+    return;
+  }
+
+  if (!left().red() && !left().left().red()) {
+    MoveRedLeft();
+  }
+
+  LlrbNode new_left = left().Clone();
+  new_left.RemoveMin();
+  set_left(std::move(new_left));
+  FixUp();
+}
+
+template <typename K, typename V>
+void LlrbNode<K, V>::MoveRedLeft() {
+  FlipColor();
+  if (right().left().red()) {
+    LlrbNode new_right = right().Clone();
+    new_right.RotateRight();
+    set_right(std::move(new_right));
+    RotateLeft();
+    FlipColor();
+  }
+}
+
+template <typename K, typename V>
+void LlrbNode<K, V>::MoveRedRight() {
+  FlipColor();
+  if (left().left().red()) {
+    RotateRight();
+    FlipColor();
+  }
+}
+
 /* Rotates left:
  *
  *      X              R
diff --git a/Firestore/core/src/firebase/firestore/immutable/llrb_node_iterator.h b/Firestore/core/src/firebase/firestore/immutable/llrb_node_iterator.h
index f1377a2..5011947 100644
--- a/Firestore/core/src/firebase/firestore/immutable/llrb_node_iterator.h
+++ b/Firestore/core/src/firebase/firestore/immutable/llrb_node_iterator.h
@@ -76,6 +76,9 @@ class LlrbNodeIterator {
   using reference = typename node_type::value_type const&;
   using difference_type = std::ptrdiff_t;
 
+  explicit LlrbNodeIterator(stack_type&& stack) : stack_(std::move(stack)) {
+  }
+
   /**
    * Constructs an iterator starting at the first node in the iteration
    * sequence of the tree represented by the given root node (i.e. it points at
@@ -196,9 +199,6 @@ class LlrbNodeIterator {
   }
 
  private:
-  explicit LlrbNodeIterator(stack_type&& stack) : stack_(std::move(stack)) {
-  }
-
   static void AccumulateLeft(const node_type* node, stack_type* stack) {
     for (; !node->empty(); node = &node->left()) {
       stack->push(node);
diff --git a/Firestore/core/src/firebase/firestore/immutable/sorted_map.h b/Firestore/core/src/firebase/firestore/immutable/sorted_map.h
index 24eb5bf..6f4d91b 100644
--- a/Firestore/core/src/firebase/firestore/immutable/sorted_map.h
+++ b/Firestore/core/src/firebase/firestore/immutable/sorted_map.h
@@ -196,7 +196,12 @@ class SortedMap : public impl::SortedMapBase {
       case Tag::Array:
         return SortedMap{array_.erase(key)};
       case Tag::Tree:
-        return SortedMap{tree_.erase(key)};
+        tree_type result = tree_.erase(key);
+        if (result.empty()) {
+          // Flip back to the array representation for empty arrays.
+          return SortedMap{};
+        }
+        return SortedMap{std::move(result)};
     }
     FIREBASE_UNREACHABLE();
   }
@@ -244,6 +249,26 @@ class SortedMap : public impl::SortedMapBase {
     FIREBASE_UNREACHABLE();
   }
 
+  const_iterator min() const {
+    switch (tag_) {
+      case Tag::Array:
+        return const_iterator(array_.min());
+      case Tag::Tree:
+        return const_iterator{tree_.min()};
+    }
+    FIREBASE_UNREACHABLE();
+  }
+
+  const_iterator max() const {
+    switch (tag_) {
+      case Tag::Array:
+        return const_iterator(array_.max());
+      case Tag::Tree:
+        return const_iterator{tree_.max()};
+    }
+    FIREBASE_UNREACHABLE();
+  }
+
   /**
    * Returns an iterator pointing to the first entry in the map. If there are
    * no entries in the map, begin() == end().
diff --git a/Firestore/core/src/firebase/firestore/immutable/tree_sorted_map.h b/Firestore/core/src/firebase/firestore/immutable/tree_sorted_map.h
index dac07b4..609b9da 100644
--- a/Firestore/core/src/firebase/firestore/immutable/tree_sorted_map.h
+++ b/Firestore/core/src/firebase/firestore/immutable/tree_sorted_map.h
@@ -106,9 +106,8 @@ class TreeSortedMap : public SortedMapBase, private util::ComparatorHolder<C> {
    * @return A new map without that value.
    */
   TreeSortedMap erase(const K& key) const {
-    // TODO(wilhuff): Actually implement erase
-    (void)key;
-    return TreeSortedMap{this->comparator()};
+    const C& comparator = this->comparator();
+    return TreeSortedMap{root_.erase(key, comparator), comparator};
   }
 
   bool contains(const K& key) const {
@@ -172,6 +171,21 @@ class TreeSortedMap : public SortedMapBase, private util::ComparatorHolder<C> {
     return npos;
   }
 
+  const_iterator min() const {
+    return begin();
+  }
+
+  const_iterator max() const {
+    if (empty()) {
+      return end();
+    }
+
+    const node_type& max_node = root_.max();
+    typename const_iterator::stack_type stack;
+    stack.push(&max_node);
+    return const_iterator{std::move(stack)};
+  }
+
   /**
    * Returns a forward iterator pointing to the first entry in the map. If there
    * are no entries in the map, begin() == end().
diff --git a/Firestore/core/test/firebase/firestore/immutable/array_sorted_map_test.cc b/Firestore/core/test/firebase/firestore/immutable/array_sorted_map_test.cc
index ea8ae6e..9a23df5 100644
--- a/Firestore/core/test/firebase/firestore/immutable/array_sorted_map_test.cc
+++ b/Firestore/core/test/firebase/firestore/immutable/array_sorted_map_test.cc
@@ -32,66 +32,6 @@ namespace impl {
 using IntMap = ArraySortedMap<int, int>;
 constexpr IntMap::size_type kFixedSize = IntMap::kFixedSize;
 
-TEST(ArraySortedMap, RemoveKeyValuePair) {
-  IntMap map{{1, 3}, {2, 4}};
-
-  IntMap new_map = map.erase(1);
-  ASSERT_TRUE(Found(new_map, 2, 4));
-  ASSERT_TRUE(NotFound(new_map, 1));
-
-  // Make sure the original one is not mutated
-  ASSERT_TRUE(Found(map, 1, 3));
-  ASSERT_TRUE(Found(map, 2, 4));
-}
-
-TEST(ArraySortedMap, MoreRemovals) {
-  IntMap map = IntMap{}
-                   .insert(1, 1)
-                   .insert(50, 50)
-                   .insert(3, 3)
-                   .insert(4, 4)
-                   .insert(7, 7)
-                   .insert(9, 9)
-                   .insert(1, 20)
-                   .insert(18, 18)
-                   .insert(3, 2)
-                   .insert(4, 71)
-                   .insert(7, 42)
-                   .insert(9, 88);
-
-  ASSERT_TRUE(Found(map, 7, 42));
-  ASSERT_TRUE(Found(map, 3, 2));
-  ASSERT_TRUE(Found(map, 1, 20));
-
-  IntMap s1 = map.erase(7);
-  IntMap s2 = map.erase(3);
-  IntMap s3 = map.erase(1);
-
-  ASSERT_TRUE(NotFound(s1, 7));
-  ASSERT_TRUE(Found(s1, 3, 2));
-  ASSERT_TRUE(Found(s1, 1, 20));
-
-  ASSERT_TRUE(Found(s2, 7, 42));
-  ASSERT_TRUE(NotFound(s2, 3));
-  ASSERT_TRUE(Found(s2, 1, 20));
-
-  ASSERT_TRUE(Found(s3, 7, 42));
-  ASSERT_TRUE(Found(s3, 3, 2));
-  ASSERT_TRUE(NotFound(s3, 1));
-}
-
-TEST(ArraySortedMap, RemovesMiddle) {
-  IntMap map{{1, 1}, {2, 2}, {3, 3}};
-  ASSERT_TRUE(Found(map, 1, 1));
-  ASSERT_TRUE(Found(map, 2, 2));
-  ASSERT_TRUE(Found(map, 3, 3));
-
-  IntMap s1 = map.erase(2);
-  ASSERT_TRUE(Found(s1, 1, 1));
-  ASSERT_TRUE(NotFound(s1, 2));
-  ASSERT_TRUE(Found(s1, 3, 3));
-}
-
 TEST(ArraySortedMap, ChecksSize) {
   std::vector<int> to_insert = Sequence(kFixedSize);
   IntMap map = ToMap<IntMap>(to_insert);
@@ -103,41 +43,6 @@ TEST(ArraySortedMap, ChecksSize) {
   ASSERT_ANY_THROW(map.insert(next, next));
 }
 
-TEST(ArraySortedMap, EmptyRemoval) {
-  IntMap map;
-  IntMap new_map = map.erase(1);
-  EXPECT_TRUE(new_map.empty());
-  EXPECT_EQ(0u, new_map.size());
-  EXPECT_TRUE(NotFound(new_map, 1));
-}
-
-TEST(ArraySortedMap, InsertionAndRemovalOfMaxItems) {
-  auto expected_size = kFixedSize;
-  auto n = static_cast<int>(expected_size);
-  std::vector<int> to_insert = Shuffled(Sequence(n));
-  std::vector<int> to_remove = Shuffled(to_insert);
-
-  // Add them to the map
-  IntMap map = ToMap<IntMap>(to_insert);
-  ASSERT_EQ(expected_size, map.size())
-      << "Check if all N objects are in the map";
-
-  // check the order is correct
-  ASSERT_SEQ_EQ(Pairs(Sorted(to_insert)), map);
-
-  for (int i : to_remove) {
-    map = map.erase(i);
-  }
-  ASSERT_EQ(0u, map.size()) << "Check we removed all of the items";
-}
-
-TEST(ArraySortedMap, BalanceProblem) {
-  std::vector<int> to_insert{1, 7, 8, 5, 2, 6, 4, 0, 3};
-
-  IntMap map = ToMap<IntMap>(to_insert);
-  ASSERT_SEQ_EQ(Pairs(Sorted(to_insert)), map);
-}
-
 }  // namespace impl
 }  // namespace immutable
 }  // namespace firestore
diff --git a/Firestore/core/test/firebase/firestore/immutable/sorted_map_test.cc b/Firestore/core/test/firebase/firestore/immutable/sorted_map_test.cc
index 0c5b2b9..25332c2 100644
--- a/Firestore/core/test/firebase/firestore/immutable/sorted_map_test.cc
+++ b/Firestore/core/test/firebase/firestore/immutable/sorted_map_test.cc
@@ -102,12 +102,14 @@ TYPED_TEST(SortedMapTest, Size) {
 }
 
 TYPED_TEST(SortedMapTest, Increasing) {
-  std::vector<int> to_insert = Sequence(this->large_number());
+  int n = this->large_number();
+  std::vector<int> to_insert = Sequence(n);
   TypeParam map = ToMap<TypeParam>(to_insert);
   ASSERT_EQ(this->large_size(), map.size());
 
   for (int i : to_insert) {
     map = map.erase(i);
+    ASSERT_EQ(static_cast<size_t>(n - i - 1), map.size());
   }
   ASSERT_EQ(0u, map.size());
 
@@ -129,6 +131,107 @@ TYPED_TEST(SortedMapTest, BalanceProblem) {
   ASSERT_SEQ_EQ(Pairs(Sorted(to_insert)), map);
 }
 
+TYPED_TEST(SortedMapTest, EmptyRemoval) {
+  TypeParam map;
+  TypeParam new_map = map.erase(1);
+  EXPECT_TRUE(new_map.empty());
+  EXPECT_EQ(0u, new_map.size());
+  EXPECT_TRUE(NotFound(new_map, 1));
+}
+
+TYPED_TEST(SortedMapTest, RemoveKeyValuePair) {
+  TypeParam map = TypeParam{}.insert(1, 3).insert(2, 4);
+
+  TypeParam new_map = map.erase(1);
+  ASSERT_TRUE(Found(new_map, 2, 4));
+  ASSERT_TRUE(NotFound(new_map, 1));
+
+  // Make sure the original one is not mutated
+  ASSERT_TRUE(Found(map, 1, 3));
+  ASSERT_TRUE(Found(map, 2, 4));
+}
+
+TYPED_TEST(SortedMapTest, MoreRemovals) {
+  TypeParam map = TypeParam()
+                      .insert(1, 1)
+                      .insert(50, 50)
+                      .insert(3, 3)
+                      .insert(4, 4)
+                      .insert(7, 7)
+                      .insert(9, 9)
+                      .insert(1, 20)
+                      .insert(18, 18)
+                      .insert(3, 2)
+                      .insert(4, 71)
+                      .insert(7, 42)
+                      .insert(9, 88);
+
+  ASSERT_TRUE(Found(map, 7, 42));
+  ASSERT_TRUE(Found(map, 3, 2));
+  ASSERT_TRUE(Found(map, 1, 20));
+
+  TypeParam s1 = map.erase(7);
+  TypeParam s2 = map.erase(3);
+  TypeParam s3 = map.erase(1);
+
+  ASSERT_TRUE(NotFound(s1, 7));
+  ASSERT_TRUE(Found(s1, 3, 2));
+  ASSERT_TRUE(Found(s1, 1, 20));
+
+  ASSERT_TRUE(Found(s2, 7, 42));
+  ASSERT_TRUE(NotFound(s2, 3));
+  ASSERT_TRUE(Found(s2, 1, 20));
+
+  ASSERT_TRUE(Found(s3, 7, 42));
+  ASSERT_TRUE(Found(s3, 3, 2));
+  ASSERT_TRUE(NotFound(s3, 1));
+}
+
+TYPED_TEST(SortedMapTest, RemovesMiddle) {
+  TypeParam map = TypeParam{}.insert(1, 1).insert(2, 2).insert(3, 3);
+  ASSERT_TRUE(Found(map, 1, 1));
+  ASSERT_TRUE(Found(map, 2, 2));
+  ASSERT_TRUE(Found(map, 3, 3));
+
+  TypeParam s1 = map.erase(2);
+  ASSERT_TRUE(Found(s1, 1, 1));
+  ASSERT_TRUE(NotFound(s1, 2));
+  ASSERT_TRUE(Found(s1, 3, 3));
+}
+
+TYPED_TEST(SortedMapTest, InsertionAndRemovalOfMaxItems) {
+  auto expected_size = this->large_size();
+  int n = this->large_number();
+  std::vector<int> to_insert = Shuffled(Sequence(n));
+  std::vector<int> to_remove = Shuffled(to_insert);
+
+  // Add them to the map
+  TypeParam map = ToMap<TypeParam>(to_insert);
+  ASSERT_EQ(expected_size, map.size())
+      << "Check if all N objects are in the map";
+
+  // check the order is correct
+  ASSERT_SEQ_EQ(Pairs(Sorted(to_insert)), map);
+
+  for (int i : to_remove) {
+    map = map.erase(i);
+  }
+  ASSERT_EQ(0u, map.size()) << "Check we removed all of the items";
+}
+
+TYPED_TEST(SortedMapTest, EraseDoesNotInvalidateIterators) {
+  std::vector<int> keys = Sequence(1, 4);
+  TypeParam original = ToMap<TypeParam>(keys);
+
+  auto begin = original.begin();
+  auto end = original.end();
+  ASSERT_EQ(Collect(original), (std::vector<std::pair<int, int>>{begin, end}));
+
+  TypeParam erased = original.erase(2);
+  ASSERT_EQ(erased.size(), original.size() - 1);
+  ASSERT_EQ(Collect(original), (std::vector<std::pair<int, int>>{begin, end}));
+}
+
 TYPED_TEST(SortedMapTest, FindEmpty) {
   TypeParam map;
   EXPECT_TRUE(NotFound(map, 10));
@@ -159,6 +262,40 @@ TYPED_TEST(SortedMapTest, FindIndex) {
   ASSERT_EQ(5u, map.find_index(50));
 }
 
+TYPED_TEST(SortedMapTest, MinMax) {
+  TypeParam empty;
+  auto min = empty.min();
+  auto max = empty.max();
+  ASSERT_EQ(empty.end(), min);
+  ASSERT_EQ(empty.end(), max);
+  ASSERT_EQ(min, max);
+
+  TypeParam one = empty.insert(1, 1);
+  min = one.min();
+  max = one.max();
+  ASSERT_NE(one.end(), min);
+  ASSERT_NE(one.end(), max);
+  ASSERT_EQ(1, min->first);
+  ASSERT_EQ(1, max->first);
+
+  // Just a regular iterator
+  ++min;
+  ASSERT_EQ(one.end(), min);
+
+  TypeParam two = one.insert(2, 2);
+  min = two.min();
+  max = two.max();
+  ASSERT_EQ(1, min->first);
+  ASSERT_EQ(2, max->first);
+
+  std::vector<int> to_insert = Sequence(this->large_number());
+  TypeParam lots = ToMap<TypeParam>(to_insert);
+  min = lots.min();
+  max = lots.max();
+  ASSERT_EQ(to_insert.front(), min->first);
+  ASSERT_EQ(to_insert.back(), max->first);
+}
+
 TYPED_TEST(SortedMapTest, IteratorsAreDefaultConstructible) {
   // If this compiles the test has succeeded
   typename TypeParam::const_iterator iter;