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

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

5 files changed, 200 insertions, 11 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().