Simplify

Author: schmidt-sebastian

firebase-firestore/src/main/java/com/google/firebase/firestore/local/IndexBackfiller.java

@@ -137,15 +137,15 @@ private int writeEntriesForCollectionGroup(
     Collection<FieldIndex> fieldIndexes = indexManager.getFieldIndexes(collectionGroup);
     IndexOffset existingOffset = getExistingOffset(fieldIndexes);
 
-    LocalDocumentsResult nextDocuments =
+    LocalDocumentsResult nextBatch =
         localDocumentsView.getNextDocuments(
             collectionGroup, existingOffset, documentsRemainingUnderCap);
+    indexManager.updateIndexEntries(nextBatch.getDocuments());
 
-    indexManager.updateIndexEntries(nextDocuments.getDocuments());
-    IndexOffset newOffset = getNewOffset(existingOffset, nextDocuments);
+    IndexOffset newOffset = getNewOffset(existingOffset, nextBatch);
     indexManager.updateCollectionGroup(collectionGroup, newOffset);
 
-    return nextDocuments.getDocuments().size();
+    return nextBatch.getDocuments().size();
   }
 
   private IndexOffset getNewOffset(IndexOffset existingOffset, LocalDocumentsResult lookupResult) {

firebase-firestore/src/main/java/com/google/firebase/firestore/local/LocalDocumentsView.java

@@ -17,6 +17,7 @@
 import static com.google.firebase.firestore.model.DocumentCollections.emptyDocumentMap;
 import static com.google.firebase.firestore.util.Assert.hardAssert;
 
+import androidx.annotation.Nullable;
 import androidx.annotation.VisibleForTesting;
 import com.google.firebase.Timestamp;
 import com.google.firebase.database.collection.ImmutableSortedMap;
@@ -31,7 +32,6 @@
 import com.google.firebase.firestore.model.mutation.MutationBatch;
 import com.google.firebase.firestore.model.mutation.Overlay;
 import com.google.firebase.firestore.model.mutation.PatchMutation;
-import java.util.ArrayList;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.HashSet;
@@ -86,15 +86,10 @@ DocumentOverlayCache getDocumentOverlayCache() {
    */
   Document getDocument(DocumentKey key) {
     Overlay overlay = documentOverlayCache.getOverlay(key);
-    // Only read from remote document cache if overlay is a patch.
-    MutableDocument document =
-        (overlay == null || overlay.getMutation() instanceof PatchMutation)
-            ? remoteDocumentCache.get(key)
-            : MutableDocument.newInvalidDocument(key);
+    MutableDocument document = getBaseDocument(key, overlay);
     if (overlay != null) {
       overlay.getMutation().applyToLocalView(document, null, Timestamp.now());
     }
-
     return document;
   }
 
@@ -279,16 +274,13 @@ LocalDocumentsResult getNextDocuments(String collectionGroup, IndexOffset offset
             : Collections.emptyMap();
 
     int largestBatchId = -1;
-    List<DocumentKey> mutatedKeys = new ArrayList<>();
-    for (Overlay overlay : overlays.values()) {
-      if (!docs.containsKey(overlay.getKey())) {
-        mutatedKeys.add(overlay.getKey());
+    if (count - docs.size() > 0) {
+      for (Overlay overlay : overlays.values()) {
+        if (!docs.containsKey(overlay.getKey())) {
+          docs.put(overlay.getKey(), getBaseDocument(overlay.getKey(), overlay));
+        }
+        largestBatchId = Math.max(largestBatchId, overlay.getLargestBatchId());
       }
-      largestBatchId = Math.max(largestBatchId, overlay.getLargestBatchId());
-    }
-
-    if (!mutatedKeys.isEmpty()) {
-      docs.putAll(remoteDocumentCache.getAll(mutatedKeys));
     }
 
     ImmutableSortedMap<DocumentKey, Document> localDocs =
@@ -325,4 +317,11 @@ private ImmutableSortedMap<DocumentKey, Document> getDocumentsMatchingCollection
     }
     return results;
   }
+
+  /** Provides a base document to apply `overlay`. */
+  private MutableDocument getBaseDocument(DocumentKey key, @Nullable Overlay overlay) {
+    return (overlay == null || overlay.getMutation() instanceof PatchMutation)
+        ? remoteDocumentCache.get(key)
+        : MutableDocument.newInvalidDocument(key);
+  }
 }