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

8 files changed, 202 insertions, 36 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)) {