Fix overlay bug with multiple field patches (#3623)

Author: wu-hui

firebase-firestore/CHANGELOG.md

@@ -3,6 +3,9 @@ by opting into a release at
 [go/firebase-android-release](http:go/firebase-android-release) (Googlers only).
 
 # Unreleased
+- [fixed] Fixed an issue where patching multiple fields shadows each other (#3528).
+
+# 24.1.1
 - [fixed] Fixed an issue in the experimental index engine that might have
   caused Firestore to exclude document results for limit queries with local
   modifications.

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,45 @@ 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(
+  /**
+   * Computes the local view for the given documents.
+   *
+   * @param docs The documents to compute views for. It also has the base version of the documents.
+   * @param overlays The overlays that need to be applied to the given base version of the
+   *     documents.
+   * @param existenceStateChanged A set of documents whose existence states might have changed. This
+   *     is used to determine if we need to re-calculate overlays from mutation queues.
+   * @return A map represents the local documents view.
+   */
+  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 +170,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.getMutation().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 +234,8 @@ private void recalculateAndSaveOverlays(Map<DocumentKey, MutableDocument> docs)
       }
       documentOverlayCache.saveOverlays(entry.getKey(), overlays);
     }
+
+    return masks;
   }
 
   /**
@@ -311,9 +336,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 overlayedDocument;
+  private FieldMask mutatedFields;
+
+  OverlayedDocument(Document overlayedDocument, FieldMask mutatedFields) {
+    this.overlayedDocument = overlayedDocument;
+    this.mutatedFields = mutatedFields;
+  }
+
+  public Document getDocument() {
+    return overlayedDocument;
+  }
+
+  /**
+   * 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/DeleteMutation.java

@@ -79,4 +79,9 @@ public void applyToRemoteDocument(MutableDocument document, MutationResult mutat
 
     return previousMask;
   }
+
+  @Override
+  public @Nullable FieldMask getFieldMask() {
+    return null;
+  }
 }

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

@@ -164,6 +164,12 @@ public abstract void applyToRemoteDocument(
   public abstract @Nullable FieldMask applyToLocalView(
       MutableDocument document, @Nullable FieldMask previousMask, Timestamp localWriteTime);
 
+  /**
+   * Returns a {@code FieldMask} representing the fields that will be changed by applying this
+   * mutation. Returns {@code null} if the mutation will overwrite the entire document.
+   */
+  public abstract @Nullable FieldMask getFieldMask();
+
   /** Helper for derived classes to implement .equals(). */
   boolean hasSameKeyAndPrecondition(Mutation other) {
     return key.equals(other.key) && precondition.equals(other.precondition);

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;
@@ -101,14 +100,12 @@ public void applyToRemoteDocument(MutableDocument document, MutationBatchResult
   }
 
   /**
-   * Computes the local view of a document given all the mutations in this batch. Returns an {@link
-   * FieldMask} representing all the fields that are mutated.
+   * Computes the local view of a document given all the mutations in this batch.
+   *
+   * @param mutatedFields The document to be mutated.
+   * @param mutatedFields Fields that are already mutated before applying the current one.
+   * @return An {@link FieldMask} representing all the fields that are mutated.
    */
-  public FieldMask applyToLocalView(MutableDocument document) {
-    FieldMask mutatedFields = FieldMask.fromSet(new HashSet<>());
-    return applyToLocalView(document, mutatedFields);
-  }
-
   public FieldMask applyToLocalView(MutableDocument document, @Nullable FieldMask mutatedFields) {
     // First, apply the base state. This allows us to apply non-idempotent transform against a
     // consistent set of values.
@@ -136,7 +133,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 +142,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/PatchMutation.java

@@ -101,7 +101,8 @@ public ObjectValue getValue() {
    * Returns the mask to apply to {@link #getValue}, where only fields that are in both the
    * fieldMask and the value will be updated.
    */
-  public FieldMask getMask() {
+  @Override
+  public FieldMask getFieldMask() {
     return mask;
   }
 

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

@@ -103,6 +103,11 @@ public FieldMask applyToLocalView(
     return null;
   }
 
+  @Override
+  public @Nullable FieldMask getFieldMask() {
+    return null;
+  }
+
   /** Returns the object value to use when setting the document. */
   public ObjectValue getValue() {
     return value;

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

@@ -66,4 +66,9 @@ public FieldMask applyToLocalView(
       MutableDocument document, @Nullable FieldMask previousMask, Timestamp localWriteTime) {
     throw Assert.fail("VerifyMutation should only be used in Transactions.");
   }
+
+  @Override
+  public @Nullable FieldMask getFieldMask() {
+    return null;
+  }
 }

firebase-firestore/src/main/java/com/google/firebase/firestore/remote/RemoteSerializer.java

@@ -269,7 +269,7 @@ public com.google.firestore.v1.Write encodeMutation(Mutation mutation) {
       builder.setUpdate(encodeDocument(mutation.getKey(), ((SetMutation) mutation).getValue()));
     } else if (mutation instanceof PatchMutation) {
       builder.setUpdate(encodeDocument(mutation.getKey(), ((PatchMutation) mutation).getValue()));
-      builder.setUpdateMask(encodeDocumentMask(((PatchMutation) mutation).getMask()));
+      builder.setUpdateMask(encodeDocumentMask((mutation.getFieldMask())));
     } else if (mutation instanceof DeleteMutation) {
       builder.setDelete(encodeKey(mutation.getKey()));
     } else if (mutation instanceof VerifyMutation) {