Filter out document updates from target association changes (#1122)

* Filter out document updates from target association changes

* Move remote-event-modifying methods onto remote event

* Style

5 files changed, 154 insertions, 27 deletions

diff --git a/Firestore/Example/Tests/Remote/FSTRemoteEventTests.mm b/Firestore/Example/Tests/Remote/FSTRemoteEventTests.mm
index c4800d4..6ac2a6b 100644
--- a/Firestore/Example/Tests/Remote/FSTRemoteEventTests.mm
+++ b/Firestore/Example/Tests/Remote/FSTRemoteEventTests.mm
@@ -23,10 +23,13 @@
 #import "Firestore/Source/Model/FSTDocumentKey.h"
 #import "Firestore/Source/Remote/FSTExistenceFilter.h"
 #import "Firestore/Source/Remote/FSTWatchChange.h"
+#include "Firestore/core/src/firebase/firestore/model/document_key.h"
 
 #import "Firestore/Example/Tests/Remote/FSTWatchChange+Testing.h"
 #import "Firestore/Example/Tests/Util/FSTHelpers.h"
 
+using firebase::firestore::model::DocumentKey;
+
 NS_ASSUME_NONNULL_BEGIN
 
 @interface FSTRemoteEventTests : XCTestCase
@@ -551,6 +554,91 @@ NS_ASSUME_NONNULL_BEGIN
   XCTAssertEqualObjects(event.targetChanges[@2].resumeToken, resumeToken3);
 }
 
+- (void)testSynthesizeDeletes {
+  FSTWatchChange *shouldSynthesize =
+      [FSTWatchTargetChange changeWithState:FSTWatchTargetChangeStateCurrent targetIDs:@[ @1 ]];
+  FSTWatchChange *wrongState =
+      [FSTWatchTargetChange changeWithState:FSTWatchTargetChangeStateNoChange targetIDs:@[ @2 ]];
+  FSTWatchChange *hasDocument =
+      [FSTWatchTargetChange changeWithState:FSTWatchTargetChangeStateCurrent targetIDs:@[ @3 ]];
+  FSTDocument *doc = FSTTestDoc("docs/1", 1, @{ @"value" : @1 }, NO);
+  FSTWatchChange *docChange = [[FSTDocumentWatchChange alloc] initWithUpdatedTargetIDs:@[ @3 ]
+                                                                      removedTargetIDs:@[]
+                                                                           documentKey:doc.key
+                                                                              document:doc];
+
+  FSTWatchChangeAggregator *aggregator =
+      [self aggregatorWithTargets:@[ @1, @2, @3 ]
+                      outstanding:_noPendingResponses
+                          changes:@[ shouldSynthesize, wrongState, hasDocument, docChange ]];
+
+  FSTRemoteEvent *event = [aggregator remoteEvent];
+  DocumentKey synthesized = DocumentKey::FromPathString("docs/2");
+  XCTAssertEqual(event.documentUpdates.find(synthesized), event.documentUpdates.end());
+
+  FSTTargetChange *limboTargetChange = event.targetChanges[@1];
+  [event synthesizeDeleteForLimboTargetChange:limboTargetChange key:synthesized];
+  FSTDeletedDocument *expected =
+      [FSTDeletedDocument documentWithKey:synthesized version:event.snapshotVersion];
+  XCTAssertEqualObjects(expected, event.documentUpdates.at(synthesized));
+
+  DocumentKey notSynthesized1 = DocumentKey::FromPathString("docs/no1");
+  [event synthesizeDeleteForLimboTargetChange:event.targetChanges[@2] key:notSynthesized1];
+  XCTAssertEqual(event.documentUpdates.find(notSynthesized1), event.documentUpdates.end());
+
+  [event synthesizeDeleteForLimboTargetChange:event.targetChanges[@3] key:doc.key];
+  FSTMaybeDocument *docData = event.documentUpdates.at(doc.key);
+  XCTAssertFalse([docData isKindOfClass:[FSTDeletedDocument class]]);
+}
+
+- (void)testFilterUpdates {
+  FSTDocument *newDoc = FSTTestDoc("docs/new", 1, @{@"key" : @"value"}, NO);
+  FSTDocument *existingDoc = FSTTestDoc("docs/existing", 1, @{@"some" : @"data"}, NO);
+  FSTWatchChange *newDocChange = [[FSTDocumentWatchChange alloc] initWithUpdatedTargetIDs:@[ @1 ]
+                                                                         removedTargetIDs:@[]
+                                                                              documentKey:newDoc.key
+                                                                                 document:newDoc];
+
+  FSTWatchTargetChange *resetTargetChange =
+      [FSTWatchTargetChange changeWithState:FSTWatchTargetChangeStateReset
+                                  targetIDs:@[ @2 ]
+                                resumeToken:_resumeToken1];
+
+  FSTWatchChange *existingDocChange =
+      [[FSTDocumentWatchChange alloc] initWithUpdatedTargetIDs:@[ @1, @2 ]
+                                              removedTargetIDs:@[]
+                                                   documentKey:existingDoc.key
+                                                      document:existingDoc];
+
+  FSTWatchChangeAggregator *aggregator =
+      [self aggregatorWithTargets:@[ @1, @2 ]
+                      outstanding:_noPendingResponses
+                          changes:@[ newDocChange, resetTargetChange, existingDocChange ]];
+  FSTRemoteEvent *event = [aggregator remoteEvent];
+  FSTDocumentKeySet *existingKeys = [[FSTDocumentKeySet keySet] setByAddingObject:existingDoc.key];
+
+  FSTTargetChange *updateChange = event.targetChanges[@1];
+  XCTAssertTrue([updateChange.mapping isKindOfClass:[FSTUpdateMapping class]]);
+  FSTUpdateMapping *update = (FSTUpdateMapping *)updateChange.mapping;
+  FSTDocumentKey *existingDocKey = existingDoc.key;
+  FSTDocumentKey *newDocKey = newDoc.key;
+  XCTAssertTrue([update.addedDocuments containsObject:existingDocKey]);
+
+  [event filterUpdatesFromTargetChange:updateChange existingDocuments:existingKeys];
+  // Now it's been filtered, since it already existed.
+  XCTAssertFalse([update.addedDocuments containsObject:existingDocKey]);
+  XCTAssertTrue([update.addedDocuments containsObject:newDocKey]);
+
+  FSTTargetChange *resetChange = event.targetChanges[@2];
+  XCTAssertTrue([resetChange.mapping isKindOfClass:[FSTResetMapping class]]);
+  FSTResetMapping *resetMapping = (FSTResetMapping *)resetChange.mapping;
+  XCTAssertTrue([resetMapping.documents containsObject:existingDocKey]);
+
+  [event filterUpdatesFromTargetChange:resetChange existingDocuments:existingKeys];
+  // Document is still there, even though it already exists. Reset mappings don't get filtered.
+  XCTAssertTrue([resetMapping.documents containsObject:existingDocKey]);
+}
+
 @end
 
 NS_ASSUME_NONNULL_END
diff --git a/Firestore/Source/Core/FSTSyncEngine.mm b/Firestore/Source/Core/FSTSyncEngine.mm
index 0a4fc94..138fb41 100644
--- a/Firestore/Source/Core/FSTSyncEngine.mm
+++ b/Firestore/Source/Core/FSTSyncEngine.mm
@@ -292,38 +292,19 @@ static const FSTListenSequenceNumber kIrrelevantSequenceNumber = -1;
 - (void)applyRemoteEvent:(FSTRemoteEvent *)remoteEvent {
   [self assertDelegateExistsForSelector:_cmd];
 
-  // Make sure limbo documents are deleted if there were no results
+  // Make sure limbo documents are deleted if there were no results.
+  // Filter out document additions to targets that they already belong to.
   [remoteEvent.targetChanges enumerateKeysAndObjectsUsingBlock:^(
                                  FSTBoxedTargetID *_Nonnull targetID,
                                  FSTTargetChange *_Nonnull targetChange, BOOL *_Nonnull stop) {
     const auto iter = self->_limboKeysByTarget.find([targetID intValue]);
     if (iter == self->_limboKeysByTarget.end()) {
-      return;
-    }
-    const DocumentKey &limboKey = iter->second;
-    if (targetChange.currentStatusUpdate == FSTCurrentStatusUpdateMarkCurrent &&
-        remoteEvent.documentUpdates.find(limboKey) == remoteEvent.documentUpdates.end()) {
-      // When listening to a query the server responds with a snapshot containing documents
-      // matching the query and a current marker telling us we're now in sync. It's possible for
-      // these to arrive as separate remote events or as a single remote event. For a document
-      // query, there will be no documents sent in the response if the document doesn't exist.
-      //
-      // If the snapshot arrives separately from the current marker, we handle it normally and
-      // updateTrackedLimboDocumentsWithChanges:targetID: will resolve the limbo status of the
-      // document, removing it from limboDocumentRefs. This works because clients only initiate
-      // limbo resolution when a target is current and because all current targets are always at a
-      // consistent snapshot.
-      //
-      // However, if the document doesn't exist and the current marker arrives, the document is
-      // not present in the snapshot and our normal view handling would consider the document to
-      // remain in limbo indefinitely because there are no updates to the document. To avoid this,
-      // we specially handle this just this case here: synthesizing a delete.
-      //
-      // TODO(dimond): Ideally we would have an explicit lookup query instead resulting in an
-      // explicit delete message and we could remove this special logic.
-      [remoteEvent
-          addDocumentUpdate:[FSTDeletedDocument documentWithKey:limboKey
-                                                        version:remoteEvent.snapshotVersion]];
+      FSTQueryView *qv = self.queryViewsByTarget[targetID];
+      FSTAssert(qv, @"Missing queryview for non-limbo query: %i", [targetID intValue]);
+      [remoteEvent filterUpdatesFromTargetChange:targetChange
+                               existingDocuments:qv.view.syncedDocuments];
+    } else {
+      [remoteEvent synthesizeDeleteForLimboTargetChange:targetChange key:iter->second];
     }
   }];
 
diff --git a/Firestore/Source/Core/FSTView.h b/Firestore/Source/Core/FSTView.h
index 38f57fc..431b863 100644
--- a/Firestore/Source/Core/FSTView.h
+++ b/Firestore/Source/Core/FSTView.h
@@ -148,6 +148,12 @@ typedef NS_ENUM(NSInteger, FSTLimboDocumentChangeType) {
  */
 - (FSTViewChange *)applyChangedOnlineState:(FSTOnlineState)onlineState;
 
+/**
+ * The set of remote documents that the server has told us belongs to the target associated with
+ * this view.
+ */
+@property(nonatomic, strong, readonly) FSTDocumentKeySet *syncedDocuments;
+
 @end
 
 NS_ASSUME_NONNULL_END
diff --git a/Firestore/Source/Remote/FSTRemoteEvent.h b/Firestore/Source/Remote/FSTRemoteEvent.h
index dd45117..0f6b6c7 100644
--- a/Firestore/Source/Remote/FSTRemoteEvent.h
+++ b/Firestore/Source/Remote/FSTRemoteEvent.h
@@ -172,6 +172,17 @@ eventWithSnapshotVersion:(FSTSnapshotVersion *)snapshotVersion
 /** Handles an existence filter mismatch */
 - (void)handleExistenceFilterMismatchForTargetID:(FSTBoxedTargetID *)targetID;
 
+- (void)synthesizeDeleteForLimboTargetChange:(FSTTargetChange *)targetChange
+                                         key:(const firebase::firestore::model::DocumentKey &)key;
+
+/**
+ * Strips out mapping changes that aren't actually changes. That is, if the document already
+ * existed in the target, and is being added in the target, and this is not a reset, we can
+ * skip doing the work to associate the document with the target because it has already been done.
+ */
+- (void)filterUpdatesFromTargetChange:(FSTTargetChange *)targetChange
+                    existingDocuments:(FSTDocumentKeySet *)existingDocuments;
+
 @end
 
 #pragma mark - FSTWatchChangeAggregator
diff --git a/Firestore/Source/Remote/FSTRemoteEvent.mm b/Firestore/Source/Remote/FSTRemoteEvent.mm
index a3c85df..30aa0d3 100644
--- a/Firestore/Source/Remote/FSTRemoteEvent.mm
+++ b/Firestore/Source/Remote/FSTRemoteEvent.mm
@@ -281,6 +281,47 @@ eventWithSnapshotVersion:(FSTSnapshotVersion *)snapshotVersion
   return self;
 }
 
+- (void)filterUpdatesFromTargetChange:(FSTTargetChange *)targetChange
+                    existingDocuments:(FSTDocumentKeySet *)existingDocuments {
+  if ([targetChange.mapping isKindOfClass:[FSTUpdateMapping class]]) {
+    FSTUpdateMapping *update = (FSTUpdateMapping *)targetChange.mapping;
+    FSTDocumentKeySet *added = update.addedDocuments;
+    __block FSTDocumentKeySet *result = added;
+    [added enumerateObjectsUsingBlock:^(FSTDocumentKey *docKey, BOOL *stop) {
+      if ([existingDocuments containsObject:docKey]) {
+        result = [result setByRemovingObject:docKey];
+      }
+    }];
+    update.addedDocuments = result;
+  }
+}
+
+- (void)synthesizeDeleteForLimboTargetChange:(FSTTargetChange *)targetChange
+                                         key:(const DocumentKey &)key {
+  if (targetChange.currentStatusUpdate == FSTCurrentStatusUpdateMarkCurrent &&
+      _documentUpdates.find(key) == _documentUpdates.end()) {
+    // When listening to a query the server responds with a snapshot containing documents
+    // matching the query and a current marker telling us we're now in sync. It's possible for
+    // these to arrive as separate remote events or as a single remote event. For a document
+    // query, there will be no documents sent in the response if the document doesn't exist.
+    //
+    // If the snapshot arrives separately from the current marker, we handle it normally and
+    // updateTrackedLimboDocumentsWithChanges:targetID: will resolve the limbo status of the
+    // document, removing it from limboDocumentRefs. This works because clients only initiate
+    // limbo resolution when a target is current and because all current targets are always at a
+    // consistent snapshot.
+    //
+    // However, if the document doesn't exist and the current marker arrives, the document is
+    // not present in the snapshot and our normal view handling would consider the document to
+    // remain in limbo indefinitely because there are no updates to the document. To avoid this,
+    // we specially handle this just this case here: synthesizing a delete.
+    //
+    // TODO(dimond): Ideally we would have an explicit lookup query instead resulting in an
+    // explicit delete message and we could remove this special logic.
+    _documentUpdates[key] = [FSTDeletedDocument documentWithKey:key version:_snapshotVersion];
+  }
+}
+
 - (NSDictionary<FSTBoxedTargetID *, FSTTargetChange *> *)targetChanges {
   return static_cast<NSDictionary<FSTBoxedTargetID *, FSTTargetChange *> *>(_targetChanges);
 }