Fix overlay bug leas to patch mutation optimization not being applied

firebase-firestore/CHANGELOG.md

@@ -1,4 +1,5 @@
 # Unreleased
+* [fixed] Fix an issue that stops some performance optimization being applied.
 
 # 24.4.1
 * [fixed] Fix `FAILED_PRECONDITION` when writing to a deleted document in a

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

@@ -160,6 +160,7 @@ private Map<DocumentKey, OverlayedDocument> computeViews(
     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
       // event *and* the overlay is a PatchMutation. This is because document existence state
       // can change if some patch mutation's preconditions are met.
@@ -174,6 +175,9 @@ private Map<DocumentKey, OverlayedDocument> computeViews(
         overlay
             .getMutation()
             .applyToLocalView(doc, overlay.getMutation().getFieldMask(), Timestamp.now());
+      } else { // overlay == null
+        // Using EMPTY to indicate there is no overlay for the document.
+        mutatedFields.put(doc.getKey(), FieldMask.EMPTY);
       }
     }
 

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

@@ -21,7 +21,7 @@
 /** Represents a local view (overlay) of a document, and the fields that are locally mutated. */
 public class OverlayedDocument {
   private Document overlayedDocument;
-  private FieldMask mutatedFields;
+  @Nullable private FieldMask mutatedFields;
 
   OverlayedDocument(Document overlayedDocument, FieldMask mutatedFields) {
     this.overlayedDocument = overlayedDocument;
@@ -33,9 +33,14 @@ public Document getDocument() {
   }
 
   /**
-   * The fields that are locally mutated by patch mutations. If the overlayed document is from set
-   * or delete mutations, this returns null.
+   * The fields that are locally mutated by patch mutations.
+   *
+   * <p>If the overlayed document is from set or delete mutations, returns null.
+   *
+   * <p>If there is no overlay (mutation) for the document, returns FieldMask.EMPTY.
    */
+  // TODO(b/262245989): This screams for a proper sum type (Tagged Union) which does not exist in
+  // Java (yet).
   public @Nullable FieldMask getMutatedFields() {
     return mutatedFields;
   }

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

@@ -142,7 +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).getDocument();
+      OverlayedDocument overlayedDocument = documentMap.get(key);
+      MutableDocument document = (MutableDocument) overlayedDocument.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.
@@ -151,6 +152,7 @@ public Map<DocumentKey, Mutation> applyToLocalDocumentSet(
       if (overlay != null) {
         overlays.put(key, overlay);
       }
+
       if (!document.isValidDocument()) {
         document.convertToNoDocument(SnapshotVersion.NONE);
       }

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

@@ -24,9 +24,14 @@
 import com.google.firebase.firestore.model.MutableDocument;
 import com.google.firebase.firestore.model.ResourcePath;
 import com.google.firebase.firestore.model.SnapshotVersion;
+import com.google.firebase.firestore.model.mutation.DeleteMutation;
 import com.google.firebase.firestore.model.mutation.Mutation;
 import com.google.firebase.firestore.model.mutation.Overlay;
+import com.google.firebase.firestore.model.mutation.PatchMutation;
+import com.google.firebase.firestore.model.mutation.SetMutation;
 import java.util.Collection;
+import java.util.Collections;
+import java.util.HashMap;
 import java.util.Map;
 import java.util.SortedSet;
 
@@ -35,10 +40,17 @@
  * mutations read.
  */
 class CountingQueryEngine extends QueryEngine {
+  enum OverlayType {
+    Patch,
+    Set,
+    Delete
+  }
+
   private final QueryEngine queryEngine;
 
   private final int[] overlaysReadByCollection = new int[] {0};
   private final int[] overlaysReadByKey = new int[] {0};
+  private final Map<DocumentKey, OverlayType> overlayTypes = new HashMap();
   private final int[] documentsReadByCollection = new int[] {0};
   private final int[] documentsReadByKey = new int[] {0};
 
@@ -49,6 +61,7 @@ class CountingQueryEngine extends QueryEngine {
   void resetCounts() {
     overlaysReadByCollection[0] = 0;
     overlaysReadByKey[0] = 0;
+    overlayTypes.clear();
     documentsReadByCollection[0] = 0;
     documentsReadByKey[0] = 0;
   }
@@ -104,6 +117,14 @@ int getOverlaysReadByKey() {
     return overlaysReadByKey[0];
   }
 
+  /**
+   * Returns the types of overlay returned by the OverlayCahce's `getOverlays()` API (since the last
+   * call to `resetCounts()`)
+   */
+  Map<DocumentKey, OverlayType> getOverlayTypes() {
+    return Collections.unmodifiableMap(overlayTypes);
+  }
+
   private RemoteDocumentCache wrapRemoteDocumentCache(RemoteDocumentCache subject) {
     return new RemoteDocumentCache() {
       @Override
@@ -160,12 +181,19 @@ private DocumentOverlayCache wrapOverlayCache(DocumentOverlayCache subject) {
       @Override
       public Overlay getOverlay(DocumentKey key) {
         ++overlaysReadByKey[0];
-        return subject.getOverlay(key);
+        Overlay overlay = subject.getOverlay(key);
+        overlayTypes.put(key, getOverlayType(overlay));
+        return overlay;
       }
 
       public Map<DocumentKey, Overlay> getOverlays(SortedSet<DocumentKey> keys) {
         overlaysReadByKey[0] += keys.size();
-        return subject.getOverlays(keys);
+        Map<DocumentKey, Overlay> overlays = subject.getOverlays(keys);
+        for (Map.Entry<DocumentKey, Overlay> entry : overlays.entrySet()) {
+          overlayTypes.put(entry.getKey(), getOverlayType(entry.getValue()));
+        }
+
+        return overlays;
       }
 
       @Override
@@ -182,6 +210,9 @@ public void removeOverlaysForBatchId(int batchId) {
       public Map<DocumentKey, Overlay> getOverlays(ResourcePath collection, int sinceBatchId) {
         Map<DocumentKey, Overlay> result = subject.getOverlays(collection, sinceBatchId);
         overlaysReadByCollection[0] += result.size();
+        for (Map.Entry<DocumentKey, Overlay> entry : result.entrySet()) {
+          overlayTypes.put(entry.getKey(), getOverlayType(entry.getValue()));
+        }
         return result;
       }
 
@@ -191,8 +222,23 @@ public Map<DocumentKey, Overlay> getOverlays(
         Map<DocumentKey, Overlay> result =
             subject.getOverlays(collectionGroup, sinceBatchId, count);
         overlaysReadByCollection[0] += result.size();
+        for (Map.Entry<DocumentKey, Overlay> entry : result.entrySet()) {
+          overlayTypes.put(entry.getKey(), getOverlayType(entry.getValue()));
+        }
         return result;
       }
+
+      private OverlayType getOverlayType(Overlay overlay) {
+        if (overlay.getMutation() instanceof SetMutation) {
+          return OverlayType.Set;
+        } else if (overlay.getMutation() instanceof PatchMutation) {
+          return OverlayType.Patch;
+        } else if (overlay.getMutation() instanceof DeleteMutation) {
+          return OverlayType.Delete;
+        } else {
+          throw new IllegalStateException("Overlay is a unrecognizable mutation.");
+        }
+      }
     };
   }
 }

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

@@ -24,6 +24,7 @@
 import static com.google.firebase.firestore.testutil.TestUtil.existenceFilterEvent;
 import static com.google.firebase.firestore.testutil.TestUtil.filter;
 import static com.google.firebase.firestore.testutil.TestUtil.key;
+import static com.google.firebase.firestore.testutil.TestUtil.keyMap;
 import static com.google.firebase.firestore.testutil.TestUtil.keySet;
 import static com.google.firebase.firestore.testutil.TestUtil.keys;
 import static com.google.firebase.firestore.testutil.TestUtil.map;
@@ -82,6 +83,7 @@
 import java.util.Collection;
 import java.util.Collections;
 import java.util.List;
+import java.util.Map;
 import java.util.Map.Entry;
 import java.util.Set;
 import java.util.stream.Collectors;
@@ -332,6 +334,11 @@ protected void assertOverlaysRead(int byKey, int byCollection) {
         "Overlays read (by collection)", byCollection, queryEngine.getOverlaysReadByCollection());
   }
 
+  /** Asserts the expected overlay types. */
+  protected void assertOverlayTypes(Map<DocumentKey, CountingQueryEngine.OverlayType> expected) {
+    assertEquals("Overlays types", expected, queryEngine.getOverlayTypes());
+  }
+
   /**
    * Asserts the expected numbers of documents read by the RemoteDocumentCache since the last call
    * to `resetPersistenceStats()`.
@@ -975,6 +982,7 @@ public void testReadsAllDocumentsForInitialCollectionQueries() {
     localStore.executeQuery(query, /* usePreviousResults= */ true);
     assertRemoteDocumentsRead(/* byKey= */ 0, /* byCollection= */ 2);
     assertOverlaysRead(/* byKey= */ 0, /* byCollection= */ 1);
+    assertOverlayTypes(keyMap("foo/bonk", CountingQueryEngine.OverlayType.Set));
   }
 
   @Test
@@ -1689,4 +1697,21 @@ public void testMultipleFieldPatchesOnLocalDocs() {
     assertChanged(doc("foo/bar", 0, map("likes", 1, "stars", 2)).setHasLocalMutations());
     assertContains(doc("foo/bar", 0, map("likes", 1, "stars", 2)).setHasLocalMutations());
   }
+
+  @Test
+  public void testUpdateOnRemoteDocLeadsToUpdateOverlay() {
+    Query query = query("foo");
+    allocateQuery(query);
+
+    applyRemoteEvent(updateRemoteEvent(doc("foo/baz", 10, map("a", 1)), asList(2), emptyList()));
+    applyRemoteEvent(updateRemoteEvent(doc("foo/bar", 20, map()), asList(2), emptyList()));
+    writeMutation(patchMutation("foo/baz", map("b", 2)));
+
+    resetPersistenceStats();
+
+    localStore.executeQuery(query, /* usePreviousResults= */ true);
+    assertRemoteDocumentsRead(/* byKey= */ 0, /* byCollection= */ 2);
+    assertOverlaysRead(/* byKey= */ 0, /* byCollection= */ 1);
+    assertOverlayTypes(keyMap("foo/baz", CountingQueryEngine.OverlayType.Patch));
+  }
 }

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

@@ -22,6 +22,7 @@
 import static com.google.firebase.firestore.testutil.TestUtil.fieldIndex;
 import static com.google.firebase.firestore.testutil.TestUtil.filter;
 import static com.google.firebase.firestore.testutil.TestUtil.key;
+import static com.google.firebase.firestore.testutil.TestUtil.keyMap;
 import static com.google.firebase.firestore.testutil.TestUtil.map;
 import static com.google.firebase.firestore.testutil.TestUtil.orderBy;
 import static com.google.firebase.firestore.testutil.TestUtil.query;
@@ -209,6 +210,12 @@ public void testUsesPartiallyIndexedOverlaysWhenAvailable() {
     Query query = query("coll").filter(filter("matches", "==", true));
     executeQuery(query);
     assertOverlaysRead(/* byKey= */ 1, /* byCollection= */ 1);
+    assertOverlayTypes(
+        keyMap(
+            "coll/a",
+            CountingQueryEngine.OverlayType.Set,
+            "coll/b",
+            CountingQueryEngine.OverlayType.Set));
     assertQueryReturned("coll/a", "coll/b");
   }
 
@@ -238,6 +245,7 @@ public void testDoesNotUseLimitWhenIndexIsOutdated() {
     // The query engine first reads the documents by key and then re-runs the query without limit.
     assertRemoteDocumentsRead(/* byKey= */ 5, /* byCollection= */ 0);
     assertOverlaysRead(/* byKey= */ 5, /* byCollection= */ 1);
+    assertOverlayTypes(keyMap("coll/b", CountingQueryEngine.OverlayType.Delete));
     assertQueryReturned("coll/a", "coll/c");
   }
 
@@ -279,6 +287,7 @@ public void testIndexesServerTimestamps() {
     Query query = query("coll").orderBy(orderBy("time", "asc"));
     executeQuery(query);
     assertOverlaysRead(/* byKey= */ 1, /* byCollection= */ 0);
+    assertOverlayTypes(keyMap("coll/a", CountingQueryEngine.OverlayType.Set));
     assertQueryReturned("coll/a");
   }
 }

firebase-firestore/src/testUtil/java/com/google/firebase/firestore/testutil/TestUtil.java

@@ -254,6 +254,14 @@ public static ImmutableSortedSet<DocumentKey> keySet(DocumentKey... keys) {
     return keySet;
   }
 
+  public static <T> Map<DocumentKey, T> keyMap(Object... entries) {
+    Map<DocumentKey, T> res = new LinkedHashMap<>();
+    for (int i = 0; i < entries.length; i += 2) {
+      res.put(DocumentKey.fromPathString((String) entries[i]), (T) entries[i + 1]);
+    }
+    return res;
+  }
+
   public static FieldFilter filter(String key, String operator, Object value) {
     return FieldFilter.create(field(key), operatorFromString(operator), wrap(value));
   }