Ugly fix fixes

Author: wu-hui

firebase-firestore/CHANGELOG.md

@@ -16,7 +16,7 @@ by opting into a release at
 - [fixed] Fixed an AppCheck issue that caused Firestore listeners to stop
   working and receive a "Permission Denied" error. This issue only occurred for
   AppCheck users that set their expiration time to under an hour.
-- [fixed] Fixed a potential problem during Firestore's shutdown that prevented 
+- [fixed] Fixed a potential problem during Firestore's shutdown that prevented
   the shutdown from proceeding if a network connection was opened right before.
 - [changed] Queries are now send to the backend before the SDK starts local 
   processing, which reduces overall Query latency.

firebase-firestore/src/androidTest/java/com/google/firebase/firestore/QueryTest.java

@@ -886,6 +886,7 @@ public void testCollectionGroupQueriesWithWhereFiltersOnArbitraryDocumentIds() {
   }
 
   // See: https://github.com/firebase/firebase-android-sdk/issues/3528
+  // TODO(Overlay): These two tests should be part of local store tests instead.
   @Test
   public void testAddThenUpdatesWhileOffline() {
     CollectionReference collection = testCollection();

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

@@ -122,6 +122,19 @@ ImmutableSortedMap<DocumentKey, Document> getLocalViewOfDocuments(
     return computeViews(docs, overlays, existenceStateChanged);
   }
 
+  /**
+   * Similar to {@link #getDocuments}, but creates the local view from the given {@code baseDocs}
+   * without retrieving documents from the local store.
+   *
+   * @param docs The documents to apply local mutations to get the local views.
+   */
+  ImmutableSortedMap<DocumentKey, Document> getLocalViewOfDocuments(
+      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(
       Map<DocumentKey, MutableDocument> docs,

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

@@ -244,9 +244,22 @@ public LocalDocumentsResult writeLocally(List<Mutation> mutations) {
     return persistence.runTransaction(
         "Locally write mutations",
         () -> {
+          // Figure out which keys do not have a remote version in the cache, this is needed to
+          // create the right overlay mutation: if no remote version presents, we do not need to
+          // create overlays as patch mutations.
+          // TODO(Overlay): Is there a better way to determine this? Document version does not work
+          // because local mutations set them back to 0.
+          Map<DocumentKey, MutableDocument> remoteDocs = remoteDocuments.getAll(keys);
+          Set<DocumentKey> docsWithoutRemoteVersion = new HashSet<>();
+          for (Map.Entry<DocumentKey, MutableDocument> entry : remoteDocs.entrySet()) {
+            if (!entry.getValue().isValidDocument()) {
+              docsWithoutRemoteVersion.add(entry.getKey());
+            }
+          }
           // 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.getDocuments(keys);
+          ImmutableSortedMap<DocumentKey, Document> documents =
+              localDocuments.getLocalViewOfDocuments(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
@@ -270,7 +283,8 @@ public LocalDocumentsResult writeLocally(List<Mutation> mutations) {
 
           MutationBatch batch =
               mutationQueue.addMutationBatch(localWriteTime, baseMutations, mutations);
-          Map<DocumentKey, Mutation> overlays = batch.applyToLocalDocumentSet(documents);
+          Map<DocumentKey, Mutation> overlays =
+              batch.applyToLocalDocumentSet(documents, docsWithoutRemoteVersion);
           documentOverlayCache.saveOverlays(batch.getBatchId(), overlays);
           return new LocalDocumentsResult(batch.getBatchId(), documents);
         });

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

@@ -23,7 +23,6 @@
 import com.google.firebase.firestore.model.FieldPath;
 import com.google.firebase.firestore.model.MutableDocument;
 import com.google.firebase.firestore.model.ObjectValue;
-import com.google.firebase.firestore.model.SnapshotVersion;
 import com.google.firestore.v1.Value;
 import java.util.ArrayList;
 import java.util.HashMap;
@@ -94,7 +93,7 @@ public static Mutation calculateOverlayMutation(MutableDocument doc, @Nullable F
     }
 
     // mask == null when there are Set or Delete being applied to get to the current document.
-    if (mask == null || doc.getVersion().equals(SnapshotVersion.NONE)) {
+    if (mask == null) {
       if (doc.isNoDocument()) {
         return new DeleteMutation(doc.getKey(), Precondition.NONE);
       } else {

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

@@ -136,7 +136,8 @@ public FieldMask applyToLocalView(MutableDocument document, @Nullable FieldMask
    * applications.
    */
   public Map<DocumentKey, Mutation> applyToLocalDocumentSet(
-      ImmutableSortedMap<DocumentKey, Document> documentMap) {
+      ImmutableSortedMap<DocumentKey, Document> 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
     // O(n).
@@ -146,6 +147,9 @@ public Map<DocumentKey, Mutation> applyToLocalDocumentSet(
       // remove this cast.
       MutableDocument document = (MutableDocument) documentMap.get(key);
       FieldMask mutatedFields = applyToLocalView(document);
+      // 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;
       Mutation overlay = Mutation.calculateOverlayMutation(document, mutatedFields);
       overlays.put(key, overlay);
       if (!document.isValidDocument()) {

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

@@ -759,7 +759,7 @@ public void testHandlesDeleteMutationThenPatchMutationThenAckThenAck() {
 
     acknowledgeMutation(2); // delete mutation
     assertRemoved("foo/bar");
-    assertContains(deletedDoc("foo/bar", 2).setHasCommittedMutations());
+    assertContains(deletedDoc("foo/bar", 0).setHasLocalMutations());
 
     acknowledgeMutation(3); // patch mutation
     assertChanged(unknownDoc("foo/bar", 3));