Simpliufy

Author: schmidt-sebastian

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

@@ -31,6 +31,7 @@
 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 com.google.firebase.firestore.util.Function;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.HashMap;
@@ -111,7 +112,7 @@ Document getDocument(DocumentKey key) {
    */
   ImmutableSortedMap<DocumentKey, Document> getDocuments(Iterable<DocumentKey> keys) {
     Map<DocumentKey, MutableDocument> docs = remoteDocumentCache.getAll(keys);
-    return getLocalViewOfDocuments(docs, new HashSet<>());
+    return getLocalViewOfDocuments(docs, documentOverlayCache::getOverlay, new HashSet<>());
   }
 
   /**
@@ -123,11 +124,13 @@ ImmutableSortedMap<DocumentKey, Document> getDocuments(Iterable<DocumentKey> key
    *     useful to determine if some documents overlay needs to be recalculated.
    */
   ImmutableSortedMap<DocumentKey, Document> getLocalViewOfDocuments(
-      Map<DocumentKey, MutableDocument> docs, Set<DocumentKey> existenceStateChanged) {
+      Map<DocumentKey, MutableDocument> docs,
+      Function<DocumentKey, Overlay> overlays,
+      Set<DocumentKey> existenceStateChanged) {
     ImmutableSortedMap<DocumentKey, Document> results = emptyDocumentMap();
     Map<DocumentKey, MutableDocument> recalculateDocuments = new HashMap<>();
     for (Map.Entry<DocumentKey, MutableDocument> entry : docs.entrySet()) {
-      Overlay overlay = documentOverlayCache.getOverlay(entry.getKey());
+      Overlay overlay = overlays.apply(entry.getKey());
       // Recalculate an overlay if the document's existence state is changed due to a remote
       // event *and* the overlay is a PatchMutation. This is because document existence state
       // can change if some patch mutation's preconditions are met.
@@ -262,18 +265,18 @@ private ImmutableSortedMap<DocumentKey, Document> getDocumentsMatchingCollection
    * @param offset The offset to index into.
    * @param count The number of documents to return
    * @return A LocalDocumentsRResult with the documents that follow the provided offset and the last
-   * processed batch id.
+   *     processed batch id.
    */
   LocalDocumentsResult getNextDocuments(String collectionGroup, IndexOffset offset, int count) {
     Map<DocumentKey, MutableDocument> docs =
         remoteDocumentCache.getAll(collectionGroup, offset, count);
-    Map<DocumentKey, Overlay> overlays = Collections.emptyMap();
+    final Map<DocumentKey, Overlay> overlays = new HashMap<>();
 
-    int largestBatchId = 0;
+    int largestBatchId = -1;
     if (count - docs.size() > 0) {
-      overlays =
+      overlays.putAll(
           documentOverlayCache.getOverlays(
-              collectionGroup, offset.getLargestBatchId(), count - docs.size());
+              collectionGroup, offset.getLargestBatchId(), count - docs.size()));
       List<DocumentKey> mutatedKeys = new ArrayList<>();
       for (Map.Entry<DocumentKey, Overlay> entry : overlays.entrySet()) {
         if (!docs.containsKey(entry.getKey())) {
@@ -284,17 +287,14 @@ LocalDocumentsResult getNextDocuments(String collectionGroup, IndexOffset offset
       docs.putAll(remoteDocumentCache.getAll(mutatedKeys));
     }
 
-    ImmutableSortedMap<DocumentKey, Document> localDocs = emptyDocumentMap();
-    for (MutableDocument doc : docs.values()) {
-      Overlay overlay =
-          overlays.containsKey(doc.getKey())
-              ? overlays.get(doc.getKey())
-              : documentOverlayCache.getOverlay(doc.getKey());
-      if (overlay != null) {
-        overlay.getMutation().applyToLocalView(doc, null, Timestamp.now());
-      }
-      localDocs = localDocs.insert(doc.getKey(), doc);
-    }
+    ImmutableSortedMap<DocumentKey, Document> localDocs =
+        getLocalViewOfDocuments(
+            docs,
+            key ->
+                overlays.containsKey(key)
+                    ? overlays.get(key)
+                    : documentOverlayCache.getOverlay(key),
+            Collections.emptySet());
     return new LocalDocumentsResult(largestBatchId, localDocs);
   }