1 files changed, 12 insertions, 11 deletions

diff --git a/Firestore/Source/Remote/FSTRemoteStore.mm b/Firestore/Source/Remote/FSTRemoteStore.mm
index 39d285a..0ea4887 100644
--- a/Firestore/Source/Remote/FSTRemoteStore.mm
+++ b/Firestore/Source/Remote/FSTRemoteStore.mm
@@ -19,7 +19,6 @@
 #include <cinttypes>
 
 #import "Firestore/Source/Core/FSTQuery.h"
-#import "Firestore/Source/Core/FSTSnapshotVersion.h"
 #import "Firestore/Source/Core/FSTTransaction.h"
 #import "Firestore/Source/Local/FSTLocalStore.h"
 #import "Firestore/Source/Local/FSTQueryData.h"
@@ -37,11 +36,14 @@
 
 #include "Firestore/core/src/firebase/firestore/auth/user.h"
 #include "Firestore/core/src/firebase/firestore/model/document_key.h"
+#include "Firestore/core/src/firebase/firestore/model/snapshot_version.h"
 #include "Firestore/core/src/firebase/firestore/util/string_apple.h"
 
 namespace util = firebase::firestore::util;
 using firebase::firestore::auth::User;
 using firebase::firestore::model::DocumentKey;
+using firebase::firestore::model::SnapshotVersion;
+using firebase::firestore::model::DocumentKeySet;
 
 NS_ASSUME_NONNULL_BEGIN
 
@@ -299,7 +301,7 @@ static const int kMaxPendingWrites = 10;
 }
 
 - (void)watchStreamDidChange:(FSTWatchChange *)change
-             snapshotVersion:(FSTSnapshotVersion *)snapshotVersion {
+             snapshotVersion:(const SnapshotVersion &)snapshotVersion {
   // Mark the connection as Online because we got a message from the server.
   [self.onlineStateTracker updateState:FSTOnlineStateOnline];
 
@@ -315,10 +317,8 @@ static const int kMaxPendingWrites = 10;
     // older than a previous snapshot we've processed (can happen after we resume a target
     // using a resume token).
     [self.accumulatedChanges addObject:change];
-    FSTAssert(snapshotVersion, @"snapshotVersion must not be nil.");
-    if ([snapshotVersion isEqual:[FSTSnapshotVersion noVersion]] ||
-        [snapshotVersion compare:[self.localStore lastRemoteSnapshotVersion]] ==
-            NSOrderedAscending) {
+    if (snapshotVersion == SnapshotVersion::None() ||
+        snapshotVersion < [self.localStore lastRemoteSnapshotVersion]) {
       return;
     }
 
@@ -354,7 +354,7 @@ static const int kMaxPendingWrites = 10;
  * on to the SyncEngine.
  */
 - (void)processBatchedWatchChanges:(NSArray<FSTWatchChange *> *)changes
-                   snapshotVersion:(FSTSnapshotVersion *)snapshotVersion {
+                   snapshotVersion:(const SnapshotVersion &)snapshotVersion {
   FSTWatchChangeAggregator *aggregator =
       [[FSTWatchChangeAggregator alloc] initWithSnapshotVersion:snapshotVersion
                                                   listenTargets:self.listenTargets
@@ -394,7 +394,7 @@ static const int kMaxPendingWrites = 10;
 
     } else {
       // Not a document query.
-      FSTDocumentKeySet *trackedRemote = [self.localStore remoteDocumentKeysForTarget:targetID];
+      DocumentKeySet trackedRemote = [self.localStore remoteDocumentKeysForTarget:targetID];
       FSTTargetMapping *mapping = remoteEvent.targetChanges[target].mapping;
       if (mapping) {
         if ([mapping isKindOfClass:[FSTUpdateMapping class]]) {
@@ -407,7 +407,7 @@ static const int kMaxPendingWrites = 10;
         }
       }
 
-      if (trackedRemote.count != (NSUInteger)filter.count) {
+      if (trackedRemote.size() != static_cast<size_t>(filter.count)) {
         FSTLog(@"Existence filter mismatch, resetting mapping");
 
         // Make sure the mismatch is exposed in the remote event
@@ -449,7 +449,8 @@ static const int kMaxPendingWrites = 10;
       if (queryData) {
         self->_listenTargets[target] =
             [queryData queryDataByReplacingSnapshotVersion:change.snapshotVersion
-                                               resumeToken:resumeToken];
+                                               resumeToken:resumeToken
+                                            sequenceNumber:queryData.sequenceNumber];
       }
     }
   }];
@@ -566,7 +567,7 @@ static const int kMaxPendingWrites = 10;
 }
 
 /** Handles a successful StreamingWriteResponse from the server that contains a mutation result. */
-- (void)writeStreamDidReceiveResponseWithVersion:(FSTSnapshotVersion *)commitVersion
+- (void)writeStreamDidReceiveResponseWithVersion:(const SnapshotVersion &)commitVersion
                                  mutationResults:(NSArray<FSTMutationResult *> *)results {
   // This is a response to a write containing mutations and should be correlated to the first
   // pending write.