diff --git a/Firestore/Source/Local/FSTLocalDocumentsView.mm b/Firestore/Source/Local/FSTLocalDocumentsView.mm
index 43378790a81..7c6ad3834bb 100644
--- a/Firestore/Source/Local/FSTLocalDocumentsView.mm
+++ b/Firestore/Source/Local/FSTLocalDocumentsView.mm
@@ -111,17 +111,13 @@ - (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];
+  NSArray<FSTMutationBatch *> *matchingBatches =
+      [self.mutationQueue allMutationBatchesAffectingQuery:query];
+  results = [self localDocuments:results inBatches:matchingBatches];
 
-  // Now use the mutation queue to discover any other documents that may match the query after
-  // applying mutations.
   DocumentKeySet matchingKeys;
-  NSArray<FSTMutationBatch *> *matchingMutationBatches =
-      [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.
 
@@ -133,17 +129,16 @@ - (FSTDocumentDictionary *)documentsMatchingCollectionQuery:(FSTQuery *)query {
   }
 
   // Now add in results for the matchingKeys.
-  FSTMaybeDocumentDictionary *matchingKeysDocs = [self documentsForKeys:matchingKeys];
+  FSTMaybeDocumentDictionary *matchingKeysDocs =
+    [self documentsForKeys:matchingKeys inBatches:matchingBatches];
   [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.
+  // Note that the extra reference here prevents ARC from deallocating the initial unfiltered
+  // results while we're enumerating them.
   FSTDocumentDictionary *unfiltered = results;
   [unfiltered
       enumerateKeysAndObjectsUsingBlock:^(FSTDocumentKey *key, FSTDocument *doc, BOOL *stop) {
@@ -151,7 +146,6 @@ - (FSTDocumentDictionary *)documentsMatchingCollectionQuery:(FSTQuery *)query {
           results = [results dictionaryByRemovingObjectForKey:key];
         }
       }];
-
   return results;
 }