Port overlay recalculation bug (Android SDK PR #3495).

Author: ehsannas

packages/firestore/src/local/local_documents_view.ts

@@ -226,22 +226,25 @@ export class LocalDocumentsView {
       .next(batches => {
         for (const batch of batches) {
           batch.keys().forEach(key => {
-            let mask: FieldMask | null = masks.has(key)
-              ? masks.get(key)!
-              : FieldMask.empty();
-            mask = batch.applyToLocalView(docs.get(key)!, mask);
-            masks.set(key, mask);
-            if (documentsByBatchId.get(batch.batchId) === null) {
+            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,
-                documentKeySet()
+                newSet
               );
             }
-            const newSet = documentsByBatchId.get(batch.batchId)!.add(key);
-            documentsByBatchId = documentsByBatchId.insert(
-              batch.batchId,
-              newSet
-            );
           });
         }
       })

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

@@ -1801,6 +1801,25 @@ function genericLocalStoreTests(
     );
   });
 
+  it('can handle batch Ack when pending batches have other docs', () => {
+    // Prepare two batches, the first one will get rejected by the backend.
+    // When the first batch is rejected, overlay is recalculated with only the
+    // second batch, even though it has more documents than what is being
+    // rejected.
+    return expectLocalStore()
+      .afterMutations([patchMutation('foo/bar', { 'foo': 'bar' })])
+      .afterMutations([
+        setMutation('foo/bar', { 'foo': 'bar-set' }),
+        setMutation('foo/another', { 'foo': 'another' })
+      ])
+      .afterRejectingMutation()
+      .toContain(doc('foo/bar', 0, { 'foo': 'bar-set' }).setHasLocalMutations())
+      .toContain(
+        doc('foo/another', 0, { 'foo': 'another' }).setHasLocalMutations()
+      )
+      .finish();
+  });
+
   it('uses target mapping to execute queries', () => {
     if (gcIsEager) {
       return;