Implement find-related methods on C++ immutable maps (#1145)

* Standardize method ordering across sorted maps
* Add SortedMap::find
* Add SortedMap::find_index
* Add SortedMap::contains

11 files changed, 260 insertions, 80 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 6808297..6ffe017 100644
--- a/Firestore/core/src/firebase/firestore/immutable/array_sorted_map.h
+++ b/Firestore/core/src/firebase/firestore/immutable/array_sorted_map.h
@@ -152,6 +152,20 @@ class ArraySortedMap : public SortedMapBase {
         key_comparator_{comparator} {
   }
 
+  /** Returns true if the map contains no elements. */
+  bool empty() const {
+    return size() == 0;
+  }
+
+  /** Returns the number of items in this map. */
+  size_type size() const {
+    return array_->size();
+  }
+
+  const key_comparator_type& comparator() const {
+    return key_comparator_;
+  }
+
   /**
    * Creates a new map identical to this one, but with a key-value pair added or
    * updated.
@@ -212,6 +226,10 @@ class ArraySortedMap : public SortedMapBase {
     }
   }
 
+  bool contains(const K& key) const {
+    return find(key) != end();
+  }
+
   /**
    * Finds a value in the map.
    *
@@ -229,18 +247,15 @@ class ArraySortedMap : public SortedMapBase {
     }
   }
 
-  const key_comparator_type& comparator() const {
-    return key_comparator_;
-  }
-
-  /** Returns true if the map contains no elements. */
-  bool empty() const {
-    return size() == 0;
-  }
-
-  /** Returns the number of items in this map. */
-  size_type size() const {
-    return array_->size();
+  /**
+   * Finds the index of the given key in the map.
+   *
+   * @param key The key to look up.
+   * @return The index of the entry containing the key, or npos if not found.
+   */
+  size_type find_index(const K& key) const {
+    auto found = find(key);
+    return found == end() ? npos : static_cast<size_type>(found - begin());
   }
 
   /**
diff --git a/Firestore/core/src/firebase/firestore/immutable/llrb_node.h b/Firestore/core/src/firebase/firestore/immutable/llrb_node.h
index 3dc4030..d2a2227 100644
--- a/Firestore/core/src/firebase/firestore/immutable/llrb_node.h
+++ b/Firestore/core/src/firebase/firestore/immutable/llrb_node.h
@@ -57,16 +57,16 @@ class LlrbNode : public SortedMapBase {
   LlrbNode() : LlrbNode{EmptyRep()} {
   }
 
-  /** Returns the number of elements at this node or beneath it in the tree. */
-  size_type size() const {
-    return rep_->size_;
-  }
-
   /** Returns true if this is an empty node--a leaf node in the tree. */
   bool empty() const {
     return size() == 0;
   }
 
+  /** Returns the number of elements at this node or beneath it in the tree. */
+  size_type size() const {
+    return rep_->size_;
+  }
+
   /** Returns true if this node is red (as opposed to black). */
   bool red() const {
     return static_cast<bool>(rep_->color_);
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 d113d34..f1377a2 100644
--- a/Firestore/core/src/firebase/firestore/immutable/llrb_node_iterator.h
+++ b/Firestore/core/src/firebase/firestore/immutable/llrb_node_iterator.h
@@ -100,6 +100,41 @@ class LlrbNodeIterator {
   }
 
   /**
+   * Constructs an iterator pointing to the first node whose key is not less
+   * than the given key. If the key is in the tree then the lower bound will be
+   * the node containing the key. If the key is not in the tree, the lower bound
+   * will the first node greater than the key. If all nodes in the tree are less
+   * than the given key, returns an equivalent to `End()`.
+   */
+  template <typename C>
+  static LlrbNodeIterator LowerBound(const node_type* root,
+                                     const key_type& key,
+                                     const C& comparator) {
+    stack_type stack;
+
+    const node_type* node = root;
+    while (!node->empty()) {
+      util::ComparisonResult cmp = util::Compare(key, node->key(), comparator);
+      if (cmp == util::ComparisonResult::Same) {
+        // Found exactly what we're looking for so we're done.
+        stack.push(node);
+        return LlrbNodeIterator{std::move(stack)};
+
+      } else if (cmp == util::ComparisonResult::Ascending) {
+        // key < node.key (for the forward direction)
+        stack.push(node);
+        node = &node->left();
+      } else {
+        // key > node.key (for the forward direction). Don't put this in the
+        // stack because we don't need to revisit it in the iteration order.
+        node = &node->right();
+      }
+    }
+
+    return LlrbNodeIterator{std::move(stack)};
+  }
+
+  /**
    * Returns true if this iterator points at the end of the iteration sequence.
    */
   bool is_end() const {
diff --git a/Firestore/core/src/firebase/firestore/immutable/sorted_map.h b/Firestore/core/src/firebase/firestore/immutable/sorted_map.h
index 7c8c832..24eb5bf 100644
--- a/Firestore/core/src/firebase/firestore/immutable/sorted_map.h
+++ b/Firestore/core/src/firebase/firestore/immutable/sorted_map.h
@@ -133,6 +133,28 @@ class SortedMap : public impl::SortedMapBase {
     return *this;
   }
 
+  /** Returns true if the map contains no elements. */
+  bool empty() const {
+    switch (tag_) {
+      case Tag::Array:
+        return array_.empty();
+      case Tag::Tree:
+        return tree_.empty();
+    }
+    FIREBASE_UNREACHABLE();
+  }
+
+  /** Returns the number of items in this map. */
+  size_type size() const {
+    switch (tag_) {
+      case Tag::Array:
+        return array_.size();
+      case Tag::Tree:
+        return tree_.size();
+    }
+    FIREBASE_UNREACHABLE();
+  }
+
   /**
    * Creates a new map identical to this one, but with a key-value pair added or
    * updated.
@@ -179,24 +201,45 @@ class SortedMap : public impl::SortedMapBase {
     FIREBASE_UNREACHABLE();
   }
 
-  /** Returns true if the map contains no elements. */
-  bool empty() const {
+  bool contains(const K& key) const {
     switch (tag_) {
       case Tag::Array:
-        return array_.empty();
+        return array_.contains(key);
       case Tag::Tree:
-        return tree_.empty();
+        return tree_.contains(key);
     }
     FIREBASE_UNREACHABLE();
   }
 
-  /** Returns the number of items in this map. */
-  size_type size() const {
+  /**
+   * Finds a value in the map.
+   *
+   * @param key The key to look up.
+   * @return An iterator pointing to the entry containing the key, or end() if
+   *     not found.
+   */
+  const_iterator find(const K& key) const {
     switch (tag_) {
       case Tag::Array:
-        return array_.size();
+        return const_iterator(array_.find(key));
       case Tag::Tree:
-        return tree_.size();
+        return const_iterator{tree_.find(key)};
+    }
+    FIREBASE_UNREACHABLE();
+  }
+
+  /**
+   * Finds the index of the given key in the map.
+   *
+   * @param key The key to look up.
+   * @return The index of the entry containing the key, or npos if not found.
+   */
+  size_type find_index(const K& key) const {
+    switch (tag_) {
+      case Tag::Array:
+        return array_.find_index(key);
+      case Tag::Tree:
+        return tree_.find_index(key);
     }
     FIREBASE_UNREACHABLE();
   }
diff --git a/Firestore/core/src/firebase/firestore/immutable/sorted_map_base.cc b/Firestore/core/src/firebase/firestore/immutable/sorted_map_base.cc
index f2faa33..954bdb9 100644
--- a/Firestore/core/src/firebase/firestore/immutable/sorted_map_base.cc
+++ b/Firestore/core/src/firebase/firestore/immutable/sorted_map_base.cc
@@ -23,6 +23,7 @@ namespace impl {
 
 // Define external storage for constants:
 constexpr SortedMapBase::size_type SortedMapBase::kFixedSize;
+constexpr SortedMapBase::size_type SortedMapBase::npos;
 
 }  // namespace impl
 }  // namespace immutable
diff --git a/Firestore/core/src/firebase/firestore/immutable/sorted_map_base.h b/Firestore/core/src/firebase/firestore/immutable/sorted_map_base.h
index cfb19c1..a19bd77 100644
--- a/Firestore/core/src/firebase/firestore/immutable/sorted_map_base.h
+++ b/Firestore/core/src/firebase/firestore/immutable/sorted_map_base.h
@@ -54,6 +54,12 @@ class SortedMapBase {
    */
   // TODO(wilhuff): actually use this for switching implementations.
   static constexpr size_type kFixedSize = 25;
+
+  /**
+   * A sentinel return value that indicates not found. Functionally similar to
+   * std::string::npos.
+   */
+  static constexpr size_type npos = static_cast<size_type>(-1);
 };
 
 }  // namespace impl
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 c5eddc2..dac07b4 100644
--- a/Firestore/core/src/firebase/firestore/immutable/tree_sorted_map.h
+++ b/Firestore/core/src/firebase/firestore/immutable/tree_sorted_map.h
@@ -72,6 +72,20 @@ class TreeSortedMap : public SortedMapBase, private util::ComparatorHolder<C> {
     return TreeSortedMap{std::move(node), comparator};
   }
 
+  /** Returns true if the map contains no elements. */
+  bool empty() const {
+    return root_.empty();
+  }
+
+  /** Returns the number of items in this map. */
+  size_type size() const {
+    return root_.size();
+  }
+
+  const node_type& root() const {
+    return root_;
+  }
+
   /**
    * Creates a new map identical to this one, but with a key-value pair added or
    * updated.
@@ -97,18 +111,65 @@ class TreeSortedMap : public SortedMapBase, private util::ComparatorHolder<C> {
     return TreeSortedMap{this->comparator()};
   }
 
-  /** Returns true if the map contains no elements. */
-  bool empty() const {
-    return root_.empty();
+  bool contains(const K& key) const {
+    // Inline the tree traversal here to avoid building up the stack required
+    // to construct a full iterator.
+    const C& comparator = this->comparator();
+    const node_type* node = &root();
+    while (!node->empty()) {
+      util::ComparisonResult cmp = util::Compare(key, node->key(), comparator);
+      if (cmp == util::ComparisonResult::Same) {
+        return true;
+      } else if (cmp == util::ComparisonResult::Ascending) {
+        node = &node->left();
+      } else {
+        node = &node->right();
+      }
+    }
+    return false;
   }
 
-  /** Returns the number of items in this map. */
-  size_type size() const {
-    return root_.size();
+  /**
+   * Finds a value in the map.
+   *
+   * @param key The key to look up.
+   * @return An iterator pointing to the entry containing the key, or end() if
+   *     not found.
+   */
+  const_iterator find(const K& key) const {
+    const_iterator found = LowerBound(key);
+    if (!found.is_end() && !this->comparator()(key, found->first)) {
+      return found;
+    } else {
+      return end();
+    }
   }
 
-  const node_type& root() const {
-    return root_;
+  /**
+   * Finds the index of the given key in the map.
+   *
+   * @param key The key to look up.
+   * @return The index of the entry containing the key, or npos if not found.
+   */
+  size_type find_index(const K& key) const {
+    const C& comparator = this->comparator();
+
+    size_type pruned_nodes = 0;
+    const node_type* node = &root_;
+    while (!node->empty()) {
+      util::ComparisonResult cmp = util::Compare(key, node->key(), comparator);
+      if (cmp == util::ComparisonResult::Same) {
+        return pruned_nodes + node->left().size();
+
+      } else if (cmp == util::ComparisonResult::Ascending) {
+        node = &node->left();
+
+      } else if (cmp == util::ComparisonResult::Descending) {
+        pruned_nodes += node->left().size() + 1;
+        node = &node->right();
+      }
+    }
+    return npos;
   }
 
   /**
@@ -133,6 +194,10 @@ class TreeSortedMap : public SortedMapBase, private util::ComparatorHolder<C> {
       : util::ComparatorHolder<C>{comparator}, root_{std::move(root)} {
   }
 
+  const_iterator LowerBound(const K& key) const noexcept {
+    return const_iterator::LowerBound(&root_, key, this->comparator());
+  }
+
   TreeSortedMap Wrap(node_type&& root) noexcept {
     return TreeSortedMap{std::move(root), this->comparator()};
   }
diff --git a/Firestore/core/src/firebase/firestore/util/comparison.h b/Firestore/core/src/firebase/firestore/util/comparison.h
index 23207f5..d7f4dfd 100644
--- a/Firestore/core/src/firebase/firestore/util/comparison.h
+++ b/Firestore/core/src/firebase/firestore/util/comparison.h
@@ -114,9 +114,10 @@ struct Comparator<std::vector<uint8_t>>
  * Perform a three-way comparison between the left and right values using
  * the appropriate Comparator for the values based on their type.
  */
-template <typename T>
-ComparisonResult Compare(const T& left, const T& right) {
-  Comparator<T> less_than;
+template <typename T, typename C = Comparator<T>>
+ComparisonResult Compare(const T& left,
+                         const T& right,
+                         const C& less_than = C()) {
   if (less_than(left, right)) {
     return ComparisonResult::Ascending;
   } else if (less_than(right, left)) {
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 9f18f2d..ea8ae6e 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,14 +32,6 @@ namespace impl {
 using IntMap = ArraySortedMap<int, int>;
 constexpr IntMap::size_type kFixedSize = IntMap::kFixedSize;
 
-TEST(ArraySortedMap, SearchForSpecificKey) {
-  IntMap map{{1, 3}, {2, 4}};
-
-  ASSERT_TRUE(Found(map, 1, 3));
-  ASSERT_TRUE(Found(map, 2, 4));
-  ASSERT_TRUE(NotFound(map, 3));
-}
-
 TEST(ArraySortedMap, RemoveKeyValuePair) {
   IntMap map{{1, 3}, {2, 4}};
 
@@ -100,13 +92,6 @@ TEST(ArraySortedMap, RemovesMiddle) {
   ASSERT_TRUE(Found(s1, 3, 3));
 }
 
-TEST(ArraySortedMap, Override) {
-  IntMap map = IntMap{}.insert(10, 10).insert(10, 8);
-
-  ASSERT_TRUE(Found(map, 10, 8));
-  ASSERT_FALSE(Found(map, 10, 10));
-}
-
 TEST(ArraySortedMap, ChecksSize) {
   std::vector<int> to_insert = Sequence(kFixedSize);
   IntMap map = ToMap<IntMap>(to_insert);
@@ -118,11 +103,6 @@ TEST(ArraySortedMap, ChecksSize) {
   ASSERT_ANY_THROW(map.insert(next, next));
 }
 
-TEST(ArraySortedMap, EmptyGet) {
-  IntMap map;
-  EXPECT_TRUE(NotFound(map, 10));
-}
-
 TEST(ArraySortedMap, EmptyRemoval) {
   IntMap map;
   IntMap new_map = map.erase(1);
@@ -158,26 +138,6 @@ TEST(ArraySortedMap, BalanceProblem) {
   ASSERT_SEQ_EQ(Pairs(Sorted(to_insert)), map);
 }
 
-// TODO(wilhuff): Iterators
-
-// TODO(wilhuff): IndexOf
-
-TEST(ArraySortedMap, AvoidsCopying) {
-  IntMap map = IntMap{}.insert(10, 20);
-  auto found = map.find(10);
-  ASSERT_NE(found, map.end());
-  EXPECT_EQ(20, found->second);
-
-  // Verify that inserting something with equal keys and values just returns
-  // the same underlying array.
-  IntMap duped = map.insert(10, 20);
-  auto duped_found = duped.find(10);
-
-  // If everything worked correctly, the backing array should not have been
-  // copied and the pointer to the entry with 10 as key should be the same.
-  EXPECT_EQ(found, duped_found);
-}
-
 }  // 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 f6d00eb..0c5b2b9 100644
--- a/Firestore/core/test/firebase/firestore/immutable/sorted_map_test.cc
+++ b/Firestore/core/test/firebase/firestore/immutable/sorted_map_test.cc
@@ -77,9 +77,8 @@ TYPED_TEST(SortedMapTest, Empty) {
   EXPECT_TRUE(map.empty());
   EXPECT_EQ(0u, map.size());
 
-  // TODO(wilhuff): re-add find
-  // EXPECT_TRUE(NotFound(map, 1));
-  // EXPECT_TRUE(NotFound(map, 10));
+  EXPECT_TRUE(NotFound(map, 1));
+  EXPECT_TRUE(NotFound(map, 10));
 }
 
 TYPED_TEST(SortedMapTest, Size) {
@@ -116,6 +115,50 @@ TYPED_TEST(SortedMapTest, Increasing) {
   ASSERT_EQ(Pairs(empty), Collect(map));
 }
 
+TYPED_TEST(SortedMapTest, Overwrite) {
+  TypeParam map = TypeParam().insert(10, 10).insert(10, 8);
+
+  ASSERT_TRUE(Found(map, 10, 8));
+  ASSERT_FALSE(Found(map, 10, 10));
+}
+
+TYPED_TEST(SortedMapTest, BalanceProblem) {
+  std::vector<int> to_insert{1, 7, 8, 5, 2, 6, 4, 0, 3};
+
+  TypeParam map = ToMap<TypeParam>(to_insert);
+  ASSERT_SEQ_EQ(Pairs(Sorted(to_insert)), map);
+}
+
+TYPED_TEST(SortedMapTest, FindEmpty) {
+  TypeParam map;
+  EXPECT_TRUE(NotFound(map, 10));
+}
+
+TYPED_TEST(SortedMapTest, FindSpecificKey) {
+  TypeParam map = TypeParam{}.insert(1, 3).insert(2, 4);
+
+  ASSERT_TRUE(Found(map, 1, 3));
+  ASSERT_TRUE(Found(map, 2, 4));
+  ASSERT_TRUE(NotFound(map, 3));
+}
+
+TYPED_TEST(SortedMapTest, FindIndex) {
+  std::vector<int> to_insert{1, 3, 4, 7, 9, 50};
+  TypeParam map = ToMap<TypeParam>(to_insert);
+
+  ASSERT_EQ(TypeParam::npos, map.find_index(0));
+  ASSERT_EQ(0u, map.find_index(1));
+  ASSERT_EQ(TypeParam::npos, map.find_index(2));
+  ASSERT_EQ(1u, map.find_index(3));
+  ASSERT_EQ(2u, map.find_index(4));
+  ASSERT_EQ(TypeParam::npos, map.find_index(5));
+  ASSERT_EQ(TypeParam::npos, map.find_index(6));
+  ASSERT_EQ(3u, map.find_index(7));
+  ASSERT_EQ(TypeParam::npos, map.find_index(8));
+  ASSERT_EQ(4u, map.find_index(9));
+  ASSERT_EQ(5u, map.find_index(50));
+}
+
 TYPED_TEST(SortedMapTest, IteratorsAreDefaultConstructible) {
   // If this compiles the test has succeeded
   typename TypeParam::const_iterator iter;
diff --git a/Firestore/core/test/firebase/firestore/immutable/testing.h b/Firestore/core/test/firebase/firestore/immutable/testing.h
index 0b25b66..9e839c6 100644
--- a/Firestore/core/test/firebase/firestore/immutable/testing.h
+++ b/Firestore/core/test/firebase/firestore/immutable/testing.h
@@ -30,6 +30,11 @@ namespace immutable {
 
 template <typename Container, typename K>
 testing::AssertionResult NotFound(const Container& map, const K& key) {
+  if (map.contains(key)) {
+    return testing::AssertionFailure()
+           << "Should not have found " << key << " using contains()";
+  }
+
   auto found = map.find(key);
   if (found == map.end()) {
     return testing::AssertionSuccess();
@@ -44,9 +49,15 @@ template <typename Container, typename K, typename V>
 testing::AssertionResult Found(const Container& map,
                                const K& key,
                                const V& expected) {
+  if (!map.contains(key)) {
+    return testing::AssertionFailure()
+           << "Did not find key " << key << " using contains()";
+  }
+
   auto found = map.find(key);
   if (found == map.end()) {
-    return testing::AssertionFailure() << "Did not find key " << key;
+    return testing::AssertionFailure()
+           << "Did not find key " << key << " using find()";
   }
   if (found->second == expected) {
     return testing::AssertionSuccess();