Integrate Document Overlay with the SDK.

packages/firestore/src/local/local_documents_view.ts

@@ -183,8 +183,8 @@ export class LocalDocumentsView {
     let recalculateDocuments = mutableDocumentMap();
     docs.forEach((_, doc) => {
       const overlay = overlays.get(doc.key);
-      // 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
+      // Recalculate an overlay if the document's existence state 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.
       // NOTE: we recalculate when `overlay` is undefined as well, because there
@@ -227,24 +227,19 @@ export class LocalDocumentsView {
         for (const batch of batches) {
           batch.keys().forEach(key => {
             const baseDoc = docs.get(key);
-            if (baseDoc !== null) {
-              let mask: FieldMask | null = masks.has(key)
-                ? masks.get(key)!
-                : FieldMask.empty();
-              mask = batch.applyToLocalView(baseDoc, mask);
-              masks.set(key, mask);
-              if (documentsByBatchId.get(batch.batchId) === null) {
-                documentsByBatchId = documentsByBatchId.insert(
-                  batch.batchId,
-                  documentKeySet()
-                );
-              }
-              const newSet = documentsByBatchId.get(batch.batchId)!.add(key);
-              documentsByBatchId = documentsByBatchId.insert(
-                batch.batchId,
-                newSet
-              );
+            if (baseDoc === null) {
+              return;
             }
+            let mask: FieldMask | null = masks.get(key) ?? FieldMask.empty();
+            mask = batch.applyToLocalView(baseDoc, mask);
+            masks.set(key, mask);
+            const newSet = (
+              documentsByBatchId.get(batch.batchId) ?? documentKeySet()
+            ).add(key);
+            documentsByBatchId = documentsByBatchId.insert(
+              batch.batchId,
+              newSet
+            );
           });
         }
       })
@@ -292,9 +287,7 @@ export class LocalDocumentsView {
   ): PersistencePromise<void> {
     return this.remoteDocumentCache
       .getEntries(transaction, documentKeys)
-      .next(docs => {
-        return this.recalculateAndSaveOverlays(transaction, docs);
-      });
+      .next(docs => this.recalculateAndSaveOverlays(transaction, docs));
   }
 
   /**

packages/firestore/src/local/local_store_impl.ts

@@ -323,8 +323,9 @@ export function localStoreWriteLocally(
       // 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.
+      // TODO(Overlay): Is there a better way to determine this? Using the
+      //  document version does not work because local mutations set them back
+      //  to 0.
       let remoteDocs = mutableDocumentMap();
       let docsWithoutRemoteVersion = documentKeySet();
       return localStoreImpl.remoteDocuments
@@ -705,16 +706,16 @@ function populateDocumentChangeBuffer(
   documents: MutableDocumentMap
 ): PersistencePromise<DocumentChangeResult> {
   let updatedKeys = documentKeySet();
-  let conditionChanged = documentKeySet();
+  let existenceChangedKeys = documentKeySet();
   documents.forEach(k => (updatedKeys = updatedKeys.add(k)));
   return documentBuffer.getEntries(txn, updatedKeys).next(existingDocs => {
-    let changedDocs = mutableDocumentMap();
+    let changedDocuments = mutableDocumentMap();
     documents.forEach((key, doc) => {
       const existingDoc = existingDocs.get(key)!;
 
       // Check if see if there is a existence state change for this document.
       if (doc.isFoundDocument() !== existingDoc.isFoundDocument()) {
-        conditionChanged = conditionChanged.add(key);
+        existenceChangedKeys = existenceChangedKeys.add(key);
       }
 
       // Note: The order of the steps below is important, since we want
@@ -726,7 +727,7 @@ function populateDocumentChangeBuffer(
         // events. We remove these documents from cache since we lost
         // access.
         documentBuffer.removeEntry(key, doc.readTime);
-        changedDocs = changedDocs.insert(key, doc);
+        changedDocuments = changedDocuments.insert(key, doc);
       } else if (
         !existingDoc.isValidDocument() ||
         doc.version.compareTo(existingDoc.version) > 0 ||
@@ -738,7 +739,7 @@ function populateDocumentChangeBuffer(
           'Cannot add a document when the remote version is zero'
         );
         documentBuffer.addEntry(doc);
-        changedDocs = changedDocs.insert(key, doc);
+        changedDocuments = changedDocuments.insert(key, doc);
       } else {
         logDebug(
           LOG_TAG,
@@ -751,10 +752,7 @@ function populateDocumentChangeBuffer(
         );
       }
     });
-    return {
-      changedDocuments: changedDocs,
-      existenceChangedKeys: conditionChanged
-    };
+    return { changedDocuments, existenceChangedKeys };
   });
 }
 

packages/firestore/src/model/field_mask.ts

@@ -17,6 +17,7 @@
 
 import { debugAssert } from '../util/assert';
 import { arrayEquals } from '../util/misc';
+import { SortedSet } from '../util/sorted_set';
 
 import { FieldPath } from './path';
 
@@ -46,6 +47,21 @@ export class FieldMask {
     return new FieldMask([]);
   }
 
+  /**
+   * Returns a new FieldMask object that is the result of adding all the given
+   * fields paths to this field mask.
+   */
+  unionWith(extraFields: FieldPath[]): FieldMask {
+    let mergedMaskSet = new SortedSet<FieldPath>(FieldPath.comparator);
+    for (const fieldPath of this.fields) {
+      mergedMaskSet = mergedMaskSet.add(fieldPath);
+    }
+    for (const fieldPath of extraFields) {
+      mergedMaskSet = mergedMaskSet.add(fieldPath);
+    }
+    return new FieldMask(mergedMaskSet.toArray());
+  }
+
   /**
    * Verifies that `fieldPath` is included by at least one field in this field
    * mask.

packages/firestore/src/model/mutation.ts

@@ -571,17 +571,9 @@ function patchMutationApplyToLocalView(
     return null;
   }
 
-  let mergedMaskSet = new SortedSet<FieldPath>(FieldPath.comparator);
-  for (const fieldPath of previousMask.fields) {
-    mergedMaskSet = mergedMaskSet.add(fieldPath);
-  }
-  for (const fieldPath of mutation.fieldMask.fields) {
-    mergedMaskSet = mergedMaskSet.add(fieldPath);
-  }
-  mutation.fieldTransforms
-    .map(transform => transform.field)
-    .forEach(fieldPath => (mergedMaskSet = mergedMaskSet.add(fieldPath)));
-  return new FieldMask(mergedMaskSet.toArray());
+  return previousMask
+    .unionWith(mutation.fieldMask.fields)
+    .unionWith(mutation.fieldTransforms.map(transform => transform.field));
 }
 
 /**

packages/firestore/test/unit/local/query_engine.test.ts

@@ -454,28 +454,8 @@ describe('QueryEngine', () => {
     // Add an unacknowledged mutation
     await addMutation(deleteMutation('coll/b'));
     const docs = await expectFullCollectionQuery(() =>
-      persistence.runTransaction('runQuery', 'readonly', txn => {
-        return targetCache
-          .getMatchingKeysForTargetId(txn, TEST_TARGET_ID)
-          .next(remoteKeys => {
-            return queryEngine
-              .getDocumentsMatchingQuery(
-                txn,
-                query1,
-                LAST_LIMBO_FREE_SNAPSHOT,
-                remoteKeys
-              )
-              .next(documentMap => {
-                let result = new DocumentSet();
-                documentMap.forEach((_, v) => {
-                  result = result.add(v);
-                });
-                return result;
-              });
-          });
-      })
+      runQuery(query1, LAST_LIMBO_FREE_SNAPSHOT)
     );
-
     verifyResult(docs, [MATCHING_DOC_A]);
   });
 });

packages/firestore/test/unit/model/mutation.test.ts

@@ -132,7 +132,7 @@ describe('Mutation', () => {
       mutationApplyToLocalView(
         overlay,
         docForOverlay,
-        /* previousMask */ null,
+        /* previousMask= */ null,
         timestamp
       );
     }
@@ -148,7 +148,12 @@ describe('Mutation', () => {
     const document = doc('collection/key', 0, docData);
 
     const set = setMutation('collection/key', { bar: 'bar-value' });
-    mutationApplyToLocalView(set, document, /* previousMask */ null, timestamp);
+    mutationApplyToLocalView(
+      set,
+      document,
+      /* previousMask= */ null,
+      timestamp
+    );
     expect(document).to.deep.equal(
       doc('collection/key', 0, { bar: 'bar-value' }).setHasLocalMutations()
     );
@@ -165,7 +170,7 @@ describe('Mutation', () => {
     mutationApplyToLocalView(
       patch,
       document,
-      /* previousMask */ null,
+      /* previousMask= */ null,
       timestamp
     );
     expect(document).to.deep.equal(
@@ -189,7 +194,7 @@ describe('Mutation', () => {
     mutationApplyToLocalView(
       patch,
       document,
-      /* previousMask */ null,
+      /* previousMask= */ null,
       timestamp
     );
     expect(document).to.deep.equal(
@@ -212,7 +217,7 @@ describe('Mutation', () => {
     mutationApplyToLocalView(
       patch,
       document,
-      /* previousMask */ null,
+      /* previousMask= */ null,
       timestamp
     );
     expect(document).to.deep.equal(
@@ -233,7 +238,7 @@ describe('Mutation', () => {
     mutationApplyToLocalView(
       patch,
       document,
-      /* previousMask */ null,
+      /* previousMask= */ null,
       timestamp
     );
     expect(document).to.deep.equal(
@@ -255,7 +260,7 @@ describe('Mutation', () => {
     mutationApplyToLocalView(
       patch,
       document,
-      /* previousMask */ null,
+      /* previousMask= */ null,
       timestamp
     );
     expect(document).to.deep.equal(
@@ -272,7 +277,7 @@ describe('Mutation', () => {
     mutationApplyToLocalView(
       patch,
       document,
-      /* previousMask */ null,
+      /* previousMask= */ null,
       timestamp
     );
     expect(document).to.deep.equal(deletedDoc('collection/key', 0));
@@ -289,7 +294,7 @@ describe('Mutation', () => {
     mutationApplyToLocalView(
       transform,
       document,
-      /* previousMask */ null,
+      /* previousMask= */ null,
       timestamp
     );
 
@@ -459,7 +464,7 @@ describe('Mutation', () => {
       mutationApplyToLocalView(
         transform,
         document,
-        /* previousMask */ null,
+        /* previousMask= */ null,
         timestamp
       );
     }
@@ -603,7 +608,7 @@ describe('Mutation', () => {
     mutationApplyToLocalView(
       mutation,
       document,
-      /* previousMask */ null,
+      /* previousMask= */ null,
       Timestamp.now()
     );
     expect(document).to.deep.equal(
@@ -758,13 +763,13 @@ describe('Mutation', () => {
     mutationApplyToLocalView(
       transform,
       document,
-      /* previousMask */ null,
+      /* previousMask= */ null,
       Timestamp.now()
     );
     mutationApplyToLocalView(
       transform,
       document,
-      /* previousMask */ null,
+      /* previousMask= */ null,
       Timestamp.now()
     );