Backport overlay bug fixes (Android PR 3691 3623)

packages/firestore/src/local/indexeddb_schema_converter.ts

@@ -18,8 +18,13 @@
 import { User } from '../auth/user';
 import { ListenSequence } from '../core/listen_sequence';
 import { SnapshotVersion } from '../core/snapshot_version';
-import { DocumentKeySet, documentKeySet } from '../model/collections';
+import {
+  DocumentKeyMap,
+  DocumentKeySet,
+  documentKeySet
+} from '../model/collections';
 import { DocumentKey } from '../model/document_key';
+import { FieldMask } from '../model/field_mask';
 import { ResourcePath } from '../model/path';
 import { debugAssert, fail, hardAssert } from '../util/assert';
 import { BATCHID_UNKNOWN } from '../util/types';
@@ -484,7 +489,9 @@ export class SchemaConverter implements SimpleDbSchemaConverter {
       this.serializer.remoteSerializer
     );
 
-    const promises: Array<PersistencePromise<void>> = [];
+    const promises: Array<
+      PersistencePromise<DocumentKeyMap<FieldMask | null>>
+    > = [];
     const userToDocumentSet = new Map<string, DocumentKeySet>();
 
     return mutationsStore

packages/firestore/src/local/local_documents_view.ts

@@ -33,7 +33,8 @@ import {
   newOverlayMap,
   documentMap,
   mutableDocumentMap,
-  documentKeySet
+  documentKeySet,
+  DocumentKeyMap
 } from '../model/collections';
 import { Document, MutableDocument } from '../model/document';
 import { DocumentKey } from '../model/document_key';
@@ -52,6 +53,7 @@ import { SortedMap } from '../util/sorted_map';
 import { DocumentOverlayCache } from './document_overlay_cache';
 import { IndexManager } from './index_manager';
 import { MutationQueue } from './mutation_queue';
+import { OverlayedDocument } from './overlayed_document';
 import { PersistencePromise } from './persistence_promise';
 import { PersistenceTransaction } from './persistence_transaction';
 import { RemoteDocumentCache } from './remote_document_cache';
@@ -92,7 +94,7 @@ export class LocalDocumentsView {
           mutationApplyToLocalView(
             overlay.mutation,
             document,
-            null,
+            FieldMask.empty(),
             Timestamp.now()
           );
         }
@@ -141,10 +143,34 @@ export class LocalDocumentsView {
         docs,
         overlays,
         existenceStateChanged
-      );
+      ).next(computeViewsResult => {
+        let result = documentMap();
+        computeViewsResult.forEach((documentKey, overlayedDocument) => {
+          result = result.insert(
+            documentKey,
+            overlayedDocument.overlayedDocument
+          );
+        });
+        return result;
+      });
     });
   }
 
+  /**
+   * Gets the overlayed documents for the given document map, which will include
+   * the local view of those documents and a `FieldMask` indicating which fields
+   * are mutated locally, `null` if overlay is a Set or Delete mutation.
+   */
+  getOverlayedDocuments(
+    transaction: PersistenceTransaction,
+    docs: MutableDocumentMap
+  ): PersistencePromise<DocumentKeyMap<OverlayedDocument>> {
+    const overlays = newOverlayMap();
+    return this.populateOverlays(transaction, overlays, docs).next(() =>
+      this.computeViews(transaction, docs, overlays, documentKeySet())
+    );
+  }
+
   /**
    * Fetches the overlays for {@code docs} and adds them to provided overlay map
    * if the map does not already contain an entry for the given document key.
@@ -170,17 +196,26 @@ export class LocalDocumentsView {
   }
 
   /**
-   * Computes the local view for documents, applying overlays from both
-   * `memoizedOverlays` and the overlay cache.
+   * 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.
    */
   computeViews(
     transaction: PersistenceTransaction,
     docs: MutableDocumentMap,
     overlays: OverlayMap,
     existenceStateChanged: DocumentKeySet
-  ): PersistencePromise<DocumentMap> {
-    let results = documentMap();
+  ): PersistencePromise<DocumentKeyMap<OverlayedDocument>> {
     let recalculateDocuments = mutableDocumentMap();
+    const mutatedFields = newDocumentKeyMap<FieldMask | null>();
+    const results = newDocumentKeyMap<OverlayedDocument>();
     docs.forEach((_, doc) => {
       const overlay = overlays.get(doc.key);
       // Recalculate an overlay if the document's existence state changed due to
@@ -196,25 +231,40 @@ export class LocalDocumentsView {
       ) {
         recalculateDocuments = recalculateDocuments.insert(doc.key, doc);
       } else if (overlay !== undefined) {
-        mutationApplyToLocalView(overlay.mutation, doc, null, Timestamp.now());
+        mutatedFields.set(doc.key, overlay.mutation.getFieldMask());
+        mutationApplyToLocalView(
+          overlay.mutation,
+          doc,
+          overlay.mutation.getFieldMask(),
+          Timestamp.now()
+        );
       }
     });
 
     return this.recalculateAndSaveOverlays(
       transaction,
       recalculateDocuments
-    ).next(() => {
-      docs.forEach((key, value) => {
-        results = results.insert(key, value);
-      });
+    ).next(recalculatedFields => {
+      recalculatedFields.forEach((documentKey, mask) =>
+        mutatedFields.set(documentKey, mask)
+      );
+      docs.forEach((documentKey, document) =>
+        results.set(
+          documentKey,
+          new OverlayedDocument(
+            document,
+            mutatedFields.get(documentKey) ?? null
+          )
+        )
+      );
       return results;
     });
   }
 
   private recalculateAndSaveOverlays(
     transaction: PersistenceTransaction,
     docs: MutableDocumentMap
-  ): PersistencePromise<void> {
+  ): PersistencePromise<DocumentKeyMap<FieldMask | null>> {
     const masks = newDocumentKeyMap<FieldMask | null>();
     // A reverse lookup map from batch id to the documents within that batch.
     let documentsByBatchId = new SortedMap<number, DocumentKeySet>(
@@ -274,7 +324,8 @@ export class LocalDocumentsView {
           );
         }
         return PersistencePromise.waitFor(promises);
-      });
+      })
+      .next(() => masks);
   }
 
   /**
@@ -284,7 +335,7 @@ export class LocalDocumentsView {
   recalculateAndSaveOverlaysForDocumentKeys(
     transaction: PersistenceTransaction,
     documentKeys: DocumentKeySet
-  ): PersistencePromise<void> {
+  ): PersistencePromise<DocumentKeyMap<FieldMask | null>> {
     return this.remoteDocumentCache
       .getEntries(transaction, documentKeys)
       .next(docs => this.recalculateAndSaveOverlays(transaction, docs));
@@ -407,7 +458,7 @@ export class LocalDocumentsView {
             mutationApplyToLocalView(
               overlay.mutation,
               document,
-              null,
+              FieldMask.empty(),
               Timestamp.now()
             );
           }

packages/firestore/src/local/local_store_impl.ts

@@ -28,8 +28,10 @@ import { canonifyTarget, Target, targetEquals } from '../core/target';
 import { BatchId, TargetId } from '../core/types';
 import { Timestamp } from '../lite-api/timestamp';
 import {
+  DocumentKeyMap,
   documentKeySet,
   DocumentKeySet,
+  documentMap,
   DocumentMap,
   mutableDocumentMap,
   MutableDocumentMap
@@ -75,6 +77,7 @@ import { LocalStore } from './local_store';
 import { LocalViewChanges } from './local_view_changes';
 import { LruGarbageCollector, LruResults } from './lru_garbage_collector';
 import { MutationQueue } from './mutation_queue';
+import { OverlayedDocument } from './overlayed_document';
 import { Persistence } from './persistence';
 import { PersistencePromise } from './persistence_promise';
 import { PersistenceTransaction } from './persistence_transaction';
@@ -315,7 +318,7 @@ export function localStoreWriteLocally(
   const localWriteTime = Timestamp.now();
   const keys = mutations.reduce((keys, m) => keys.add(m.key), documentKeySet());
 
-  let existingDocs: DocumentMap;
+  let overlayedDocuments: DocumentKeyMap<OverlayedDocument>;
   let mutationBatch: MutationBatch;
 
   return localStoreImpl.persistence
@@ -342,13 +345,13 @@ export function localStoreWriteLocally(
           // 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.
-          return localStoreImpl.localDocuments.getLocalViewOfDocuments(
+          return localStoreImpl.localDocuments.getOverlayedDocuments(
             txn,
             remoteDocs
           );
         })
         .next(docs => {
-          existingDocs = docs;
+          overlayedDocuments = docs;
 
           // For non-idempotent mutations (such as `FieldValue.increment()`),
           // we record the base state in a separate patch mutation. This is
@@ -360,7 +363,7 @@ export function localStoreWriteLocally(
           for (const mutation of mutations) {
             const baseValue = mutationExtractBaseValue(
               mutation,
-              existingDocs.get(mutation.key)!
+              overlayedDocuments.get(mutation.key)!.overlayedDocument
             );
             if (baseValue != null) {
               // NOTE: The base state should only be applied if there's some
@@ -387,7 +390,7 @@ export function localStoreWriteLocally(
         .next(batch => {
           mutationBatch = batch;
           const overlays = batch.applyToLocalDocumentSet(
-            existingDocs,
+            overlayedDocuments,
             docsWithoutRemoteVersion
           );
           return localStoreImpl.documentOverlayCache.saveOverlays(
@@ -398,7 +401,11 @@ export function localStoreWriteLocally(
         });
     })
     .then(() => {
-      return { batchId: mutationBatch.batchId, changes: existingDocs };
+      let documents = documentMap();
+      overlayedDocuments.forEach(
+        (key, val) => (documents = documents.insert(key, val.overlayedDocument))
+      );
+      return { batchId: mutationBatch.batchId, changes: documents };
     });
 }
 

packages/firestore/src/local/overlayed_document.ts

@@ -0,0 +1,35 @@
+/**
+ * @license
+ * 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.
+ */
+
+import { Document } from '../model/document';
+import { FieldMask } from '../model/field_mask';
+
+/**
+ * Represents a local view (overlay) of a document, and the fields that are
+ * locally mutated.
+ */
+export class OverlayedDocument {
+  constructor(
+    readonly overlayedDocument: Document,
+
+    /**
+     * The fields that are locally mutated by patch mutations. If the overlayed
+     * document is from set or delete mutations, this returns null.
+     */
+    readonly mutatedFields: FieldMask | null
+  ) {}
+}

packages/firestore/src/model/mutation.ts

@@ -213,6 +213,12 @@ export abstract class Mutation {
   abstract readonly key: DocumentKey;
   abstract readonly precondition: Precondition;
   abstract readonly fieldTransforms: FieldTransform[];
+  /**
+   * Returns a `FieldMask` representing the fields that will be changed by
+   * applying this mutation. Returns `null` if the mutation will overwrite the
+   * entire document.
+   */
+  abstract getFieldMask(): FieldMask | null;
 }
 
 /**
@@ -444,6 +450,10 @@ export class SetMutation extends Mutation {
   }
 
   readonly type: MutationType = MutationType.Set;
+
+  getFieldMask(): FieldMask | null {
+    return null;
+  }
 }
 
 function setMutationApplyToRemoteDocument(
@@ -516,6 +526,10 @@ export class PatchMutation extends Mutation {
   }
 
   readonly type: MutationType = MutationType.Patch;
+
+  getFieldMask(): FieldMask | null {
+    return this.fieldMask;
+  }
 }
 
 function patchMutationApplyToRemoteDocument(
@@ -670,6 +684,10 @@ export class DeleteMutation extends Mutation {
 
   readonly type: MutationType = MutationType.Delete;
   readonly fieldTransforms: FieldTransform[] = [];
+
+  getFieldMask(): FieldMask | null {
+    return null;
+  }
 }
 
 function deleteMutationApplyToRemoteDocument(
@@ -720,4 +738,8 @@ export class VerifyMutation extends Mutation {
 
   readonly type: MutationType = MutationType.Verify;
   readonly fieldTransforms: FieldTransform[] = [];
+
+  getFieldMask(): FieldMask | null {
+    return null;
+  }
 }