Fix overlay recalculation bug with multiple batches (#3495)

Author: wu-hui

firebase-firestore/CHANGELOG.md

@@ -12,6 +12,8 @@ by opting into a release at
   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.
+- [fixed] Fixed a NPE issue where mutations with multiple documents are not
+  handled correctly during previous mutation acknowledgement(#3490).
 
 # 24.0.1
 - [changed] Improved performance for databases that contain many document

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

@@ -165,8 +165,14 @@ private void recalculateAndSaveOverlays(Map<DocumentKey, MutableDocument> docs)
     // along the way.
     for (MutationBatch batch : batches) {
       for (DocumentKey key : batch.getKeys()) {
+        MutableDocument baseDoc = docs.get(key);
+        if (baseDoc == null) {
+          // If this batch has documents not included in passed in `docs`, skip them.
+          continue;
+        }
+
         FieldMask mask = masks.containsKey(key) ? masks.get(key) : FieldMask.EMPTY;
-        mask = batch.applyToLocalView(docs.get(key), mask);
+        mask = batch.applyToLocalView(baseDoc, mask);
         masks.put(key, mask);
         int batchId = batch.getBatchId();
         if (!documentsByBatchId.containsKey(batchId)) {

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

@@ -1599,4 +1599,21 @@ public void testOnlyPersistsUpdatesForDocumentsWhenVersionChanges() {
     assertContains(doc("foo/bar", 1, map("val", "old")));
     assertContains(doc("foo/baz", 2, map("val", "new")));
   }
+
+  @Test
+  public void testCanHandleBatchAckWhenPendingBatchesHaveOtherDocs() {
+    // 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.
+    // See: https://github.com/firebase/firebase-android-sdk/issues/3490
+    writeMutation(patchMutation("foo/bar", map("foo", "bar")));
+    writeMutations(
+        asList(
+            setMutation("foo/bar", map("foo", "bar-set")),
+            setMutation("foo/another", map("foo", "another"))));
+
+    rejectMutation();
+    assertContains(doc("foo/bar", 0, map("foo", "bar-set")).setHasLocalMutations());
+    assertContains(doc("foo/another", 0, map("foo", "another")).setHasLocalMutations());
+  }
 }