Use a prefix scan when fetching documents matching a query. (#488)

Minor optimization (which is already present in the ts code).

2 files changed, 11 insertions, 6 deletions

diff --git a/Firestore/Example/Tests/Local/FSTRemoteDocumentCacheTests.m b/Firestore/Example/Tests/Local/FSTRemoteDocumentCacheTests.m
index e6ec699..16fe3bf 100644
--- a/Firestore/Example/Tests/Local/FSTRemoteDocumentCacheTests.m
+++ b/Firestore/Example/Tests/Local/FSTRemoteDocumentCacheTests.m
@@ -105,6 +105,8 @@ static const int kVersion = 42;
 - (void)testDocumentsMatchingQuery {
   if (!self.remoteDocumentCache) return;
 
+  // TODO(rsgowman): This just verifies that we do a prefix scan against the
+  // query path. We'll need more tests once we add index support.
   [self setTestDocumentAtPath:@"a/1"];
   [self setTestDocumentAtPath:@"b/1"];
   [self setTestDocumentAtPath:@"b/2"];
@@ -114,12 +116,10 @@ static const int kVersion = 42;
   FSTDocumentDictionary *results = [self.remoteDocumentCache documentsMatchingQuery:query];
   NSArray *expected =
       @[ FSTTestDoc(@"b/1", kVersion, _kDocData, NO), FSTTestDoc(@"b/2", kVersion, _kDocData, NO) ];
+  XCTAssertEqual([results count], [expected count]);
   for (FSTDocument *doc in expected) {
     XCTAssertEqualObjects([results objectForKey:doc.key], doc);
   }
-
-  // TODO(mikelehen): Perhaps guard against extra documents in the result set once our
-  // implementations are smarter.
 }
 
 #pragma mark - Helpers
diff --git a/Firestore/Source/Local/FSTLevelDBRemoteDocumentCache.mm b/Firestore/Source/Local/FSTLevelDBRemoteDocumentCache.mm
index b2df67f..02f9f3e 100644
--- a/Firestore/Source/Local/FSTLevelDBRemoteDocumentCache.mm
+++ b/Firestore/Source/Local/FSTLevelDBRemoteDocumentCache.mm
@@ -21,6 +21,7 @@
 #include <string>
 
 #import "Firestore/Protos/objc/firestore/local/MaybeDocument.pbobjc.h"
+#import "Firestore/Source/Core/FSTQuery.h"
 #import "Firestore/Source/Local/FSTLevelDBKey.h"
 #import "Firestore/Source/Local/FSTLocalSerializer.h"
 #import "Firestore/Source/Local/FSTWriteGroup.h"
@@ -28,6 +29,7 @@
 #import "Firestore/Source/Model/FSTDocumentDictionary.h"
 #import "Firestore/Source/Model/FSTDocumentKey.h"
 #import "Firestore/Source/Model/FSTDocumentSet.h"
+#import "Firestore/Source/Model/FSTPath.h"
 #import "Firestore/Source/Util/FSTAssert.h"
 
 #include "Firestore/Port/ordered_code.h"
@@ -102,10 +104,11 @@ static ReadOptions StandardReadOptions() {
 }
 
 - (FSTDocumentDictionary *)documentsMatchingQuery:(FSTQuery *)query {
-  // TODO(mikelehen): PERF: At least filter to the documents that match the path of the query.
   FSTDocumentDictionary *results = [FSTDocumentDictionary documentDictionary];
 
-  std::string startKey = [FSTLevelDBRemoteDocumentKey keyPrefix];
+  // Documents are ordered by key, so we can use a prefix scan to narrow down
+  // the documents we need to match the query against.
+  std::string startKey = [FSTLevelDBRemoteDocumentKey keyPrefixWithResourcePath:query.path];
   std::unique_ptr<Iterator> it(_db->NewIterator(StandardReadOptions()));
   it->Seek(startKey);
 
@@ -113,7 +116,9 @@ static ReadOptions StandardReadOptions() {
   for (; it->Valid() && [currentKey decodeKey:it->key()]; it->Next()) {
     FSTMaybeDocument *maybeDoc =
         [self decodedMaybeDocument:it->value() withKey:currentKey.documentKey];
-    if ([maybeDoc isKindOfClass:[FSTDocument class]]) {
+    if (![query.path isPrefixOfPath:maybeDoc.key.path]) {
+      break;
+    } else if ([maybeDoc isKindOfClass:[FSTDocument class]]) {
       results = [results dictionaryBySettingObject:(FSTDocument *)maybeDoc forKey:maybeDoc.key];
     }
   }