1 files changed, 50 insertions, 12 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 487a25a..db53969 100644
--- a/Firestore/core/src/firebase/firestore/immutable/array_sorted_map.h
+++ b/Firestore/core/src/firebase/firestore/immutable/array_sorted_map.h
@@ -24,10 +24,12 @@
 #include <memory>
 #include <utility>
 
+#include "Firestore/core/src/firebase/firestore/immutable/keys_view.h"
 #include "Firestore/core/src/firebase/firestore/immutable/map_entry.h"
 #include "Firestore/core/src/firebase/firestore/immutable/sorted_map_base.h"
 #include "Firestore/core/src/firebase/firestore/util/comparison.h"
 #include "Firestore/core/src/firebase/firestore/util/firebase_assert.h"
+#include "Firestore/core/src/firebase/firestore/util/range.h"
 
 namespace firebase {
 namespace firestore {
@@ -133,6 +135,7 @@ class ArraySortedMap : public SortedMapBase {
    */
   using array_type = FixedArray<value_type>;
   using const_iterator = typename array_type::const_iterator;
+  using const_key_iterator = util::iterator_first<const_iterator>;
 
   using array_pointer = std::shared_ptr<const array_type>;
 
@@ -162,8 +165,8 @@ class ArraySortedMap : public SortedMapBase {
     return array_->size();
   }
 
-  const key_comparator_type& comparator() const {
-    return key_comparator_;
+  const C& comparator() const {
+    return key_comparator_.comparator();
   }
 
   /**
@@ -176,12 +179,12 @@ class ArraySortedMap : public SortedMapBase {
    */
   ArraySortedMap insert(const K& key, const V& value) const {
     const_iterator current_end = end();
-    const_iterator pos = LowerBound(key);
+    const_iterator pos = lower_bound(key);
     bool replacing_entry = false;
 
     if (pos != current_end) {
-      // LowerBound found an entry where pos->first >= pair.first. Reversing the
-      // argument order here tests pair.first < pos->first.
+      // lower_bound found an entry where pos->first >= pair.first. Reversing
+      // the argument order here tests pair.first < pos->first.
       replacing_entry = !key_comparator_(key, *pos);
       if (replacing_entry && value == pos->second) {
         return *this;
@@ -239,9 +242,9 @@ class ArraySortedMap : public SortedMapBase {
    */
   const_iterator find(const K& key) const {
     const_iterator not_found = end();
-    const_iterator lower_bound = LowerBound(key);
-    if (lower_bound != not_found && !key_comparator_(key, *lower_bound)) {
-      return lower_bound;
+    const_iterator bound = lower_bound(key);
+    if (bound != not_found && !key_comparator_(key, *bound)) {
+      return bound;
     } else {
       return not_found;
     }
@@ -258,6 +261,19 @@ class ArraySortedMap : public SortedMapBase {
     return found == end() ? npos : static_cast<size_type>(found - begin());
   }
 
+  /**
+   * Finds the first entry in the map containing a key greater than or equal
+   * to the given key.
+   *
+   * @param key The key to look up.
+   * @return An iterator pointing to the entry containing the key or the next
+   *     largest key. Can return end() if all keys in the map are less than the
+   *     requested key.
+   */
+  const_iterator lower_bound(const K& key) const {
+    return std::lower_bound(begin(), end(), key, key_comparator_);
+  }
+
   const_iterator min() const {
     return begin();
   }
@@ -285,6 +301,32 @@ class ArraySortedMap : public SortedMapBase {
     return array_->end();
   }
 
+  /**
+   * Returns a view of this SortedMap containing just the keys that have been
+   * inserted.
+   */
+  const util::range<const_key_iterator> keys() const {
+    return KeysView(*this);
+  }
+
+  /**
+   * Returns a view of this SortedMap containing just the keys that have been
+   * inserted that are greater than or equal to the given key.
+   */
+  const util::range<const_key_iterator> keys_from(const K& key) const {
+    return KeysViewFrom(*this, key);
+  }
+
+  /**
+   * Returns a view of this SortedMap containing just the keys that have been
+   * inserted that are greater than or equal to the given start_key and less
+   * than the given end_key.
+   */
+  const util::range<const_key_iterator> keys_in(const K& start_key,
+                                                const K& end_key) const {
+    return impl::KeysViewIn(*this, start_key, end_key, comparator());
+  }
+
  private:
   static array_pointer EmptyArray() {
     static const array_pointer kEmptyArray =
@@ -301,10 +343,6 @@ class ArraySortedMap : public SortedMapBase {
     return ArraySortedMap{array, key_comparator_};
   }
 
-  const_iterator LowerBound(const K& key) const {
-    return std::lower_bound(begin(), end(), key, key_comparator_);
-  }
-
   array_pointer array_;
   key_comparator_type key_comparator_;
 };