Fix overlay bug with multiple field patches

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

@@ -14,9 +14,12 @@
 
 package com.google.firebase.firestore.local;
 
+import static com.google.firebase.firestore.model.DocumentCollections.emptyDocumentMap;
+
 import com.google.firebase.database.collection.ImmutableSortedMap;
 import com.google.firebase.firestore.model.Document;
 import com.google.firebase.firestore.model.DocumentKey;
+import java.util.Map;
 
 /**
  * Represents a set of document along with their mutation batch ID.
@@ -33,6 +36,16 @@ public final class LocalDocumentsResult {
     this.documents = documents;
   }
 
+  public static LocalDocumentsResult fromOverlayedDocuments(
+      int batchId, Map<DocumentKey, OverlayedDocument> overlays) {
+    ImmutableSortedMap<DocumentKey, Document> documents = emptyDocumentMap();
+    for (Map.Entry<DocumentKey, OverlayedDocument> entry : overlays.entrySet()) {
+      documents = documents.insert(entry.getKey(), entry.getValue().getDocument());
+    }
+
+    return new LocalDocumentsResult(batchId, documents);
+  }
+
   public int getBatchId() {
     return batchId;
   }

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

@@ -119,29 +119,36 @@ ImmutableSortedMap<DocumentKey, Document> getLocalViewOfDocuments(
       Map<DocumentKey, MutableDocument> docs, Set<DocumentKey> existenceStateChanged) {
     Map<DocumentKey, Overlay> overlays = new HashMap<>();
     populateOverlays(overlays, docs.keySet());
-    return computeViews(docs, overlays, existenceStateChanged);
+    ImmutableSortedMap<DocumentKey, Document> result = emptyDocumentMap();
+    for (Map.Entry<DocumentKey, OverlayedDocument> entry :
+        computeViews(docs, overlays, existenceStateChanged).entrySet()) {
+      result = result.insert(entry.getKey(), entry.getValue().getDocument());
+    }
+
+    return result;
   }
 
   /**
-   * Similar to {@link #getDocuments}, but creates the local view from the given {@code baseDocs}
-   * without retrieving documents from the local store.
+   * Gets the overlayed documents for the given document map, which will include the local view of
+   * those documents and a {@code FieldMask} indicating which fields are mutated locally, null if
+   * overlay is a Set or Delete mutation.
    *
    * @param docs The documents to apply local mutations to get the local views.
    */
-  ImmutableSortedMap<DocumentKey, Document> getLocalViewOfDocuments(
+  Map<DocumentKey, OverlayedDocument> getOverlayedDocuments(
       Map<DocumentKey, MutableDocument> docs) {
     Map<DocumentKey, Overlay> overlays = new HashMap<>();
     populateOverlays(overlays, docs.keySet());
     return computeViews(docs, overlays, new HashSet<>());
   }
 
   /*Computes the local view for doc */
-  private ImmutableSortedMap<DocumentKey, Document> computeViews(
+  private Map<DocumentKey, OverlayedDocument> computeViews(
       Map<DocumentKey, MutableDocument> docs,
       Map<DocumentKey, Overlay> overlays,
       Set<DocumentKey> existenceStateChanged) {
-    ImmutableSortedMap<DocumentKey, Document> results = emptyDocumentMap();
     Map<DocumentKey, MutableDocument> recalculateDocuments = new HashMap<>();
+    Map<DocumentKey, FieldMask> mutatedFields = new HashMap<>();
     for (MutableDocument doc : docs.values()) {
       Overlay overlay = overlays.get(doc.getKey());
       // Recalculate an overlay if the document's existence state is changed due to a remote
@@ -154,19 +161,26 @@ private ImmutableSortedMap<DocumentKey, Document> computeViews(
           && (overlay == null || overlay.getMutation() instanceof PatchMutation)) {
         recalculateDocuments.put(doc.getKey(), doc);
       } else if (overlay != null) {
+        mutatedFields.put(doc.getKey(), overlay.getFieldMask());
         overlay.getMutation().applyToLocalView(doc, null, Timestamp.now());
       }
     }
 
-    recalculateAndSaveOverlays(recalculateDocuments);
+    Map<DocumentKey, FieldMask> recalculatedFields =
+        recalculateAndSaveOverlays(recalculateDocuments);
+    mutatedFields.putAll(recalculatedFields);
 
+    Map<DocumentKey, OverlayedDocument> result = new HashMap<>();
     for (Map.Entry<DocumentKey, MutableDocument> entry : docs.entrySet()) {
-      results = results.insert(entry.getKey(), entry.getValue());
+      result.put(
+          entry.getKey(),
+          new OverlayedDocument(entry.getValue(), mutatedFields.get(entry.getKey())));
     }
-    return results;
+    return result;
   }
 
-  private void recalculateAndSaveOverlays(Map<DocumentKey, MutableDocument> docs) {
+  private Map<DocumentKey, FieldMask> recalculateAndSaveOverlays(
+      Map<DocumentKey, MutableDocument> docs) {
     List<MutationBatch> batches =
         mutationQueue.getAllMutationBatchesAffectingDocumentKeys(docs.keySet());
 
@@ -211,6 +225,8 @@ private void recalculateAndSaveOverlays(Map<DocumentKey, MutableDocument> docs)
       }
       documentOverlayCache.saveOverlays(entry.getKey(), overlays);
     }
+
+    return masks;
   }
 
   /**
@@ -311,9 +327,9 @@ LocalDocumentsResult getNextDocuments(String collectionGroup, IndexOffset offset
     }
 
     populateOverlays(overlays, docs.keySet());
-    ImmutableSortedMap<DocumentKey, Document> localDocs =
+    Map<DocumentKey, OverlayedDocument> localDocs =
         computeViews(docs, overlays, Collections.emptySet());
-    return new LocalDocumentsResult(largestBatchId, localDocs);
+    return LocalDocumentsResult.fromOverlayedDocuments(largestBatchId, localDocs);
   }
 
   /**

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

@@ -258,8 +258,8 @@ public LocalDocumentsResult writeLocally(List<Mutation> mutations) {
           }
           // Load and apply all existing mutations. This lets us compute the current base state for
           // all non-idempotent transforms before applying any additional user-provided writes.
-          ImmutableSortedMap<DocumentKey, Document> documents =
-              localDocuments.getLocalViewOfDocuments(remoteDocs);
+          Map<DocumentKey, OverlayedDocument> overlayedDocuments =
+              localDocuments.getOverlayedDocuments(remoteDocs);
 
           // For non-idempotent mutations (such as `FieldValue.increment()`), we record the base
           // state in a separate patch mutation. This is later used to guarantee consistent values
@@ -268,7 +268,8 @@ public LocalDocumentsResult writeLocally(List<Mutation> mutations) {
           List<Mutation> baseMutations = new ArrayList<>();
           for (Mutation mutation : mutations) {
             ObjectValue baseValue =
-                mutation.extractTransformBaseValue(documents.get(mutation.getKey()));
+                mutation.extractTransformBaseValue(
+                    overlayedDocuments.get(mutation.getKey()).getDocument());
             if (baseValue != null) {
               // NOTE: The base state should only be applied if there's some existing
               // document to override, so use a Precondition of exists=true
@@ -284,9 +285,10 @@ public LocalDocumentsResult writeLocally(List<Mutation> mutations) {
           MutationBatch batch =
               mutationQueue.addMutationBatch(localWriteTime, baseMutations, mutations);
           Map<DocumentKey, Mutation> overlays =
-              batch.applyToLocalDocumentSet(documents, docsWithoutRemoteVersion);
+              batch.applyToLocalDocumentSet(overlayedDocuments, docsWithoutRemoteVersion);
           documentOverlayCache.saveOverlays(batch.getBatchId(), overlays);
-          return new LocalDocumentsResult(batch.getBatchId(), documents);
+          return LocalDocumentsResult.fromOverlayedDocuments(
+              batch.getBatchId(), overlayedDocuments);
         });
   }
 

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

@@ -0,0 +1,42 @@
+// Copyright 2022 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package com.google.firebase.firestore.local;
+
+import com.google.firebase.firestore.model.Document;
+import com.google.firebase.firestore.model.mutation.FieldMask;
+import javax.annotation.Nullable;
+
+/** Represents a local view (overlay) of a document, and the fields that are locally mutated. */
+public class OverlayedDocument {
+  private Document overlay;
+  private FieldMask mutatedFields;
+
+  OverlayedDocument(Document overlay, FieldMask mutatedFields) {
+    this.overlay = overlay;
+    this.mutatedFields = mutatedFields;
+  }
+
+  public Document getDocument() {
+    return overlay;
+  }
+
+  /**
+   * The fields that are locally mutated by patch mutations. If the overlayed document is from set
+   * or delete mutations, this returns null.
+   */
+  public @Nullable FieldMask getMutatedFields() {
+    return mutatedFields;
+  }
+}

firebase-firestore/src/main/java/com/google/firebase/firestore/model/mutation/MutationBatch.java

@@ -18,8 +18,7 @@
 
 import androidx.annotation.Nullable;
 import com.google.firebase.Timestamp;
-import com.google.firebase.database.collection.ImmutableSortedMap;
-import com.google.firebase.firestore.model.Document;
+import com.google.firebase.firestore.local.OverlayedDocument;
 import com.google.firebase.firestore.model.DocumentKey;
 import com.google.firebase.firestore.model.MutableDocument;
 import com.google.firebase.firestore.model.SnapshotVersion;
@@ -136,7 +135,7 @@ public FieldMask applyToLocalView(MutableDocument document, @Nullable FieldMask
    * applications.
    */
   public Map<DocumentKey, Mutation> applyToLocalDocumentSet(
-      ImmutableSortedMap<DocumentKey, Document> documentMap,
+      Map<DocumentKey, OverlayedDocument> documentMap,
       Set<DocumentKey> documentsWithoutRemoteVersion) {
     // TODO(mrschmidt): This implementation is O(n^2). If we iterate through the mutations first
     // (as done in `applyToLocalView(MutableDocument d)`), we can reduce the complexity to
@@ -145,8 +144,8 @@ public Map<DocumentKey, Mutation> applyToLocalDocumentSet(
     for (DocumentKey key : getKeys()) {
       // TODO(mutabledocuments): This method should take a map of MutableDocuments and we should
       // remove this cast.
-      MutableDocument document = (MutableDocument) documentMap.get(key);
-      FieldMask mutatedFields = applyToLocalView(document);
+      MutableDocument document = (MutableDocument) documentMap.get(key).getDocument();
+      FieldMask mutatedFields = applyToLocalView(document, documentMap.get(key).getMutatedFields());
       // Set mutationFields to null if the document is only from local mutations, this creates
       // a Set(or Delete) mutation, instead of trying to create a patch mutation as the overlay.
       mutatedFields = documentsWithoutRemoteVersion.contains(key) ? null : mutatedFields;

firebase-firestore/src/main/java/com/google/firebase/firestore/model/mutation/Overlay.java

@@ -36,4 +36,10 @@ public static Overlay create(int largestBatchId, Mutation mutation) {
   public DocumentKey getKey() {
     return getMutation().getKey();
   }
+
+  public FieldMask getFieldMask() {
+    return getMutation() instanceof PatchMutation
+        ? ((PatchMutation) getMutation()).getMask()
+        : null;
+  }
 }

firebase-firestore/src/test/java/com/google/firebase/firestore/local/LocalStoreTestCase.java

@@ -1616,4 +1616,69 @@ public void testCanHandleBatchAckWhenPendingBatchesHaveOtherDocs() {
     assertContains(doc("foo/bar", 0, map("foo", "bar-set")).setHasLocalMutations());
     assertContains(doc("foo/another", 0, map("foo", "another")).setHasLocalMutations());
   }
+
+  @Test
+  public void testMultipleFieldPatchesOnRemoteDocs() {
+    Query query = query("foo");
+    allocateQuery(query);
+    assertTargetId(2);
+
+    applyRemoteEvent(
+        addedRemoteEvent(doc("foo/bar", 1, map("likes", 0, "stars", 0)), asList(2), emptyList()));
+    assertChanged(doc("foo/bar", 1, map("likes", 0, "stars", 0)));
+    assertContains(doc("foo/bar", 1, map("likes", 0, "stars", 0)));
+
+    writeMutation(patchMutation("foo/bar", map("likes", 1)));
+    assertChanged(doc("foo/bar", 1, map("likes", 1, "stars", 0)).setHasLocalMutations());
+    assertContains(doc("foo/bar", 1, map("likes", 1, "stars", 0)).setHasLocalMutations());
+
+    writeMutation(patchMutation("foo/bar", map("stars", 1)));
+    assertChanged(doc("foo/bar", 1, map("likes", 1, "stars", 1)).setHasLocalMutations());
+    assertContains(doc("foo/bar", 1, map("likes", 1, "stars", 1)).setHasLocalMutations());
+
+    writeMutation(patchMutation("foo/bar", map("stars", 2)));
+    assertChanged(doc("foo/bar", 1, map("likes", 1, "stars", 2)).setHasLocalMutations());
+    assertContains(doc("foo/bar", 1, map("likes", 1, "stars", 2)).setHasLocalMutations());
+  }
+
+  @Test
+  public void testMultipleFieldPatchesInOneBatchOnRemoteDocs() {
+    Query query = query("foo");
+    allocateQuery(query);
+    assertTargetId(2);
+
+    applyRemoteEvent(
+        addedRemoteEvent(doc("foo/bar", 1, map("likes", 0, "stars", 0)), asList(2), emptyList()));
+    assertChanged(doc("foo/bar", 1, map("likes", 0, "stars", 0)));
+    assertContains(doc("foo/bar", 1, map("likes", 0, "stars", 0)));
+
+    writeMutations(
+        Lists.newArrayList(
+            patchMutation("foo/bar", map("likes", 1)), patchMutation("foo/bar", map("stars", 1))));
+    assertChanged(doc("foo/bar", 1, map("likes", 1, "stars", 1)).setHasLocalMutations());
+    assertContains(doc("foo/bar", 1, map("likes", 1, "stars", 1)).setHasLocalMutations());
+
+    writeMutation(patchMutation("foo/bar", map("stars", 2)));
+    assertChanged(doc("foo/bar", 1, map("likes", 1, "stars", 2)).setHasLocalMutations());
+    assertContains(doc("foo/bar", 1, map("likes", 1, "stars", 2)).setHasLocalMutations());
+  }
+
+  @Test
+  public void testMultipleFieldPatchesOnLocalDocs() {
+    writeMutation(setMutation("foo/bar", map("likes", 0, "stars", 0)));
+    assertChanged(doc("foo/bar", 0, map("likes", 0, "stars", 0)).setHasLocalMutations());
+    assertContains(doc("foo/bar", 0, map("likes", 0, "stars", 0)).setHasLocalMutations());
+
+    writeMutation(patchMutation("foo/bar", map("likes", 1)));
+    assertChanged(doc("foo/bar", 0, map("likes", 1, "stars", 0)).setHasLocalMutations());
+    assertContains(doc("foo/bar", 0, map("likes", 1, "stars", 0)).setHasLocalMutations());
+
+    writeMutation(patchMutation("foo/bar", map("stars", 1)));
+    assertChanged(doc("foo/bar", 0, map("likes", 1, "stars", 1)).setHasLocalMutations());
+    assertContains(doc("foo/bar", 0, map("likes", 1, "stars", 1)).setHasLocalMutations());
+
+    writeMutation(patchMutation("foo/bar", map("stars", 2)));
+    assertChanged(doc("foo/bar", 0, map("likes", 1, "stars", 2)).setHasLocalMutations());
+    assertContains(doc("foo/bar", 0, map("likes", 1, "stars", 2)).setHasLocalMutations());
+  }
 }