Port the feedback from the Java port (#1199)

* Port the feedback from the port

* Fix comment grammar

4 files changed, 76 insertions, 49 deletions

diff --git a/Firestore/Example/Tests/Remote/FSTRemoteEventTests.mm b/Firestore/Example/Tests/Remote/FSTRemoteEventTests.mm
index 35903e7..a88fdab 100644
--- a/Firestore/Example/Tests/Remote/FSTRemoteEventTests.mm
+++ b/Firestore/Example/Tests/Remote/FSTRemoteEventTests.mm
@@ -558,20 +558,10 @@ NS_ASSUME_NONNULL_BEGIN
 - (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 ]];
+  FSTWatchChangeAggregator *aggregator = [self aggregatorWithTargets:@[ @1 ]
+                                                         outstanding:_noPendingResponses
+                                                             changes:@[ shouldSynthesize ]];
 
   FSTRemoteEvent *event = [aggregator remoteEvent];
   DocumentKey synthesized = DocumentKey::FromPathString("docs/2");
@@ -583,12 +573,36 @@ NS_ASSUME_NONNULL_BEGIN
       [FSTDeletedDocument documentWithKey:synthesized version:event.snapshotVersion];
   XCTAssertEqualObjects(expected, event.documentUpdates.at(synthesized));
   XCTAssertTrue([event.limboDocumentChanges containsObject:synthesized]);
+}
+
+- (void)testDoesntSynthesizeDeletesForWrongState {
+  FSTWatchChange *wrongState =
+      [FSTWatchTargetChange changeWithState:FSTWatchTargetChangeStateNoChange targetIDs:@[ @2 ]];
+
+  FSTWatchChangeAggregator *aggregator =
+      [self aggregatorWithTargets:@[ @2 ] outstanding:_noPendingResponses changes:@[ wrongState ]];
+
+  FSTRemoteEvent *event = [aggregator remoteEvent];
 
   DocumentKey notSynthesized = DocumentKey::FromPathString("docs/no1");
   [event synthesizeDeleteForLimboTargetChange:event.targetChanges[@2] key:notSynthesized];
   XCTAssertEqual(event.documentUpdates.find(notSynthesized), event.documentUpdates.end());
   XCTAssertFalse([event.limboDocumentChanges containsObject:notSynthesized]);
+}
 
+- (void)testDoesntSynthesizeDeletesForExistingDoc {
+  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:@[ @3 ]
+                                                         outstanding:_noPendingResponses
+                                                             changes:@[ hasDocument, docChange ]];
+
+  FSTRemoteEvent *event = [aggregator remoteEvent];
   [event synthesizeDeleteForLimboTargetChange:event.targetChanges[@3] key:doc.key];
   FSTMaybeDocument *docData = event.documentUpdates.at(doc.key);
   XCTAssertFalse([docData isKindOfClass:[FSTDeletedDocument class]]);
@@ -603,21 +617,16 @@ NS_ASSUME_NONNULL_BEGIN
                                                                               documentKey:newDoc.key
                                                                                  document:newDoc];
 
-  FSTWatchTargetChange *resetTargetChange =
-      [FSTWatchTargetChange changeWithState:FSTWatchTargetChangeStateReset
-                                  targetIDs:@[ @2 ]
-                                resumeToken:_resumeToken1];
-
   FSTWatchChange *existingDocChange =
-      [[FSTDocumentWatchChange alloc] initWithUpdatedTargetIDs:@[ @1, @2 ]
+      [[FSTDocumentWatchChange alloc] initWithUpdatedTargetIDs:@[ @1 ]
                                               removedTargetIDs:@[]
                                                    documentKey:existingDoc.key
                                                       document:existingDoc];
 
   FSTWatchChangeAggregator *aggregator =
-      [self aggregatorWithTargets:@[ @1, @2 ]
+      [self aggregatorWithTargets:@[ @1 ]
                       outstanding:_noPendingResponses
-                          changes:@[ newDocChange, resetTargetChange, existingDocChange ]];
+                          changes:@[ newDocChange, existingDocChange ]];
   FSTRemoteEvent *event = [aggregator remoteEvent];
   FSTDocumentKeySet *existingKeys = [[FSTDocumentKeySet keySet] setByAddingObject:existingDoc.key];
 
@@ -628,17 +637,37 @@ NS_ASSUME_NONNULL_BEGIN
   FSTDocumentKey *newDocKey = newDoc.key;
   XCTAssertTrue([update.addedDocuments containsObject:existingDocKey]);
 
-  [event filterUpdatesFromTargetChange:updateChange existingDocuments:existingKeys];
+  [update filterUpdatesUsingExistingKeys:existingKeys];
   // Now it's been filtered, since it already existed.
   XCTAssertFalse([update.addedDocuments containsObject:existingDocKey]);
   XCTAssertTrue([update.addedDocuments containsObject:newDocKey]);
+}
+
+- (void)testDoesntFilterResets {
+  FSTDocument *existingDoc = FSTTestDoc("docs/existing", 1, @{@"some" : @"data"}, NO);
+  FSTDocumentKey *existingDocKey = existingDoc.key;
+  FSTWatchTargetChange *resetTargetChange =
+      [FSTWatchTargetChange changeWithState:FSTWatchTargetChangeStateReset
+                                  targetIDs:@[ @2 ]
+                                resumeToken:_resumeToken1];
+  FSTWatchChange *existingDocChange =
+      [[FSTDocumentWatchChange alloc] initWithUpdatedTargetIDs:@[ @2 ]
+                                              removedTargetIDs:@[]
+                                                   documentKey:existingDocKey
+                                                      document:existingDoc];
+  FSTWatchChangeAggregator *aggregator =
+      [self aggregatorWithTargets:@[ @2 ]
+                      outstanding:_noPendingResponses
+                          changes:@[ resetTargetChange, existingDocChange ]];
+  FSTRemoteEvent *event = [aggregator remoteEvent];
+  FSTDocumentKeySet *existingKeys = [[FSTDocumentKeySet keySet] setByAddingObject:existingDocKey];
 
   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];
+  [resetMapping filterUpdatesUsingExistingKeys:existingKeys];
   // Document is still there, even though it already exists. Reset mappings don't get filtered.
   XCTAssertTrue([resetMapping.documents containsObject:existingDocKey]);
 }
diff --git a/Firestore/Source/Core/FSTSyncEngine.mm b/Firestore/Source/Core/FSTSyncEngine.mm
index ca15e64..775f865 100644
--- a/Firestore/Source/Core/FSTSyncEngine.mm
+++ b/Firestore/Source/Core/FSTSyncEngine.mm
@@ -303,8 +303,7 @@ static const FSTListenSequenceNumber kIrrelevantSequenceNumber = -1;
     if (iter == self->_limboKeysByTarget.end()) {
       FSTQueryView *qv = self.queryViewsByTarget[targetID];
       FSTAssert(qv, @"Missing queryview for non-limbo query: %i", [targetID intValue]);
-      [remoteEvent filterUpdatesFromTargetChange:targetChange
-                               existingDocuments:qv.view.syncedDocuments];
+      [targetChange.mapping filterUpdatesUsingExistingKeys:qv.view.syncedDocuments];
     } else {
       [remoteEvent synthesizeDeleteForLimboTargetChange:targetChange key:iter->second];
     }
diff --git a/Firestore/Source/Remote/FSTRemoteEvent.h b/Firestore/Source/Remote/FSTRemoteEvent.h
index 322cb63..06099c7 100644
--- a/Firestore/Source/Remote/FSTRemoteEvent.h
+++ b/Firestore/Source/Remote/FSTRemoteEvent.h
@@ -43,6 +43,14 @@ NS_ASSUME_NONNULL_BEGIN
  * base class.
  */
 @interface FSTTargetMapping : NSObject
+
+/**
+ * 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)filterUpdatesUsingExistingKeys:(FSTDocumentKeySet *)existingKeys;
+
 @end
 
 #pragma mark - FSTResetMapping
@@ -184,14 +192,6 @@ initWithSnapshotVersion:(firebase::firestore::model::SnapshotVersion)snapshotVer
 - (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 26c32a2..4f7e344 100644
--- a/Firestore/Source/Remote/FSTRemoteEvent.mm
+++ b/Firestore/Source/Remote/FSTRemoteEvent.mm
@@ -104,6 +104,10 @@ NS_ASSUME_NONNULL_BEGIN
   self.documents = [self.documents setByRemovingObject:documentKey];
 }
 
+- (void)filterUpdatesUsingExistingKeys:(FSTDocumentKeySet *)existingKeys {
+  // No-op. Resets are not filtered.
+}
+
 @end
 
 #pragma mark - FSTUpdateMapping
@@ -174,6 +178,16 @@ NS_ASSUME_NONNULL_BEGIN
   self.removedDocuments = [self.removedDocuments setByAddingObject:documentKey];
 }
 
+- (void)filterUpdatesUsingExistingKeys:(FSTDocumentKeySet *)existingKeys {
+  __block FSTDocumentKeySet *result = _addedDocuments;
+  [_addedDocuments enumerateObjectsUsingBlock:^(FSTDocumentKey *docKey, BOOL *stop) {
+    if ([existingKeys containsObject:docKey]) {
+      result = [result setByRemovingObject:docKey];
+    }
+  }];
+  _addedDocuments = result;
+}
+
 @end
 
 #pragma mark - FSTTargetChange
@@ -287,21 +301,6 @@ initWithSnapshotVersion:(SnapshotVersion)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 &&
@@ -320,7 +319,7 @@ initWithSnapshotVersion:(SnapshotVersion)snapshotVersion
     // 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.
+    // we specially handle 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.