Minimize rereading of batches in FSTLocalDocumentsView documentsMatchingCollectionQuery

Firestore/Source/Local/FSTLocalDocumentsView.mm

@@ -70,8 +70,12 @@ - (nullable FSTMaybeDocument *)documentForKey:(const DocumentKey &)key {
 // Internal version of documentForKey: which allows reusing `batches`.
 - (nullable FSTMaybeDocument *)documentForKey:(const DocumentKey &)key
                                     inBatches:(NSArray<FSTMutationBatch *> *)batches {
-  FSTMaybeDocument *_Nullable remoteDoc = [self.remoteDocumentCache entryForKey:key];
-  return [self localDocument:remoteDoc key:key inBatches:batches];
+  FSTMaybeDocument *_Nullable document = [self.remoteDocumentCache entryForKey:key];
+  for (FSTMutationBatch *batch in batches) {
+    document = [batch applyTo:document documentKey:key];
+  }
+
+  return document;
 }
 
 - (FSTMaybeDocumentDictionary *)documentsForKeys:(const DocumentKeySet &)keys {
@@ -110,37 +114,34 @@ - (FSTDocumentDictionary *)documentsMatchingDocumentQuery:(const ResourcePath &)
 }
 
 - (FSTDocumentDictionary *)documentsMatchingCollectionQuery:(FSTQuery *)query {
-  // Query the remote documents and overlay mutations.
-  // TODO(mikelehen): There may be significant overlap between the mutations affecting these
-  // remote documents and the allMutationBatchesAffectingQuery mutations. Consider optimizing.
   __block FSTDocumentDictionary *results = [self.remoteDocumentCache documentsMatchingQuery:query];
-  results = [self localDocuments:results];
-
-  // Now use the mutation queue to discover any other documents that may match the query after
-  // applying mutations.
-  DocumentKeySet matchingKeys;
-  NSArray<FSTMutationBatch *> *matchingMutationBatches =
+  // Get locally persisted mutation batches.
+  NSArray<FSTMutationBatch *> *matchingBatches =
       [self.mutationQueue allMutationBatchesAffectingQuery:query];
-  for (FSTMutationBatch *batch in matchingMutationBatches) {
+
+  for (FSTMutationBatch *batch in matchingBatches) {
     for (FSTMutation *mutation in batch.mutations) {
-      // TODO(mikelehen): PERF: Check if this mutation actually affects the query to reduce work.
+      // Only process documents belonging to the collection.
+      if (!query.path.IsImmediateParentOf(mutation.key.path())) {
+        continue;
+      }
 
-      // If the key is already in the results, we can skip it.
-      if (![results containsKey:mutation.key]) {
-        matchingKeys = matchingKeys.insert(mutation.key);
+      FSTDocumentKey *key = static_cast<FSTDocumentKey *>(mutation.key);
+      // baseDoc may be nil for the documents that weren't yet written to the backend.
+      FSTMaybeDocument *baseDoc = results[key];
+      FSTMaybeDocument *mutatedDoc =
+          [mutation applyTo:baseDoc baseDocument:baseDoc localWriteTime:batch.localWriteTime];
+
+      if ([mutatedDoc isKindOfClass:[FSTDeletedDocument class]]) {
+        results = [results dictionaryByRemovingObjectForKey:key];
+      } else if ([mutatedDoc isKindOfClass:[FSTDocument class]]) {
+        results = [results dictionaryBySettingObject:(FSTDocument *)mutatedDoc forKey:key];
+      } else {
+        HARD_FAIL("Unknown document: %s", mutatedDoc);
       }
     }
   }
 
-  // Now add in results for the matchingKeys.
-  FSTMaybeDocumentDictionary *matchingKeysDocs = [self documentsForKeys:matchingKeys];
-  [matchingKeysDocs
-      enumerateKeysAndObjectsUsingBlock:^(FSTDocumentKey *key, FSTMaybeDocument *doc, BOOL *stop) {
-        if ([doc isKindOfClass:[FSTDocument class]]) {
-          results = [results dictionaryBySettingObject:(FSTDocument *)doc forKey:key];
-        }
-      }];
-
   // Finally, filter out any documents that don't actually match the query. Note that the extra
   // reference here prevents ARC from deallocating the initial unfiltered results while we're
   // enumerating them.
@@ -155,54 +156,6 @@ - (FSTDocumentDictionary *)documentsMatchingCollectionQuery:(FSTQuery *)query {
   return results;
 }
 
-/**
- * Takes a remote document and applies local mutations to generate the local view of the
- * document.
- *
- * @param document The base remote document to apply mutations to.
- * @param documentKey The key of the document (necessary when remoteDocument is nil).
- */
-- (nullable FSTMaybeDocument *)localDocument:(nullable FSTMaybeDocument *)document
-                                         key:(const DocumentKey &)documentKey
-                                   inBatches:(NSArray<FSTMutationBatch *> *)batches {
-  for (FSTMutationBatch *batch in batches) {
-    document = [batch applyTo:document documentKey:documentKey];
-  }
-
-  return document;
-}
-
-/**
- * Takes a set of remote documents and applies local mutations to generate the local view of
- * the documents.
- *
- * @param documents The base remote documents to apply mutations to.
- * @return The local view of the documents.
- */
-- (FSTDocumentDictionary *)localDocuments:(FSTDocumentDictionary *)documents {
-  __block DocumentKeySet keySet;
-  [documents
-      enumerateKeysAndObjectsUsingBlock:^(FSTDocumentKey *key, FSTDocument *doc, BOOL *stop) {
-        keySet = keySet.insert(doc.key);
-      }];
-  NSArray<FSTMutationBatch *> *batches =
-      [self.mutationQueue allMutationBatchesAffectingDocumentKeys:keySet];
-
-  __block FSTDocumentDictionary *result = documents;
-  [documents enumerateKeysAndObjectsUsingBlock:^(FSTDocumentKey *key, FSTDocument *remoteDocument,
-                                                 BOOL *stop) {
-    FSTMaybeDocument *mutatedDoc = [self localDocument:remoteDocument key:key inBatches:batches];
-    if ([mutatedDoc isKindOfClass:[FSTDeletedDocument class]]) {
-      result = [result dictionaryByRemovingObjectForKey:key];
-    } else if ([mutatedDoc isKindOfClass:[FSTDocument class]]) {
-      result = [result dictionaryBySettingObject:(FSTDocument *)mutatedDoc forKey:key];
-    } else {
-      HARD_FAIL("Unknown document: %s", mutatedDoc);
-    }
-  }];
-  return result;
-}
-
 @end
 
 NS_ASSUME_NONNULL_END

Firestore/core/src/firebase/firestore/model/base_path.h

@@ -139,6 +139,16 @@ class BasePath {
     return size() <= rhs.size() && std::equal(begin(), end(), rhs.begin());
   }
 
+  /**
+   * Returns true if the given argument is a direct child of this path.
+   *
+   * Empty path is a parent of any path that consists of a single segment.
+   */
+  bool IsImmediateParentOf(const T& potential_child) const {
+    return size() + 1 == potential_child.size() &&
+           std::equal(begin(), end(), potential_child.begin());
+  }
+
   bool operator==(const BasePath& rhs) const {
     return segments_ == rhs.segments_;
   }