1 files changed, 73 insertions, 8 deletions

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()};
   }