Port FSTDocumentKeySet to C++ DocumentKeySet (#1229)

* Define a Comparator for DocumentKey

* Automated migration from FSTDocumentKeySet to DocumentKeySet

* Manual fixups for DocumentKeySet

* Delete FSTDocumentKeySet

24 files changed, 315 insertions, 328 deletions

diff --git a/Firestore/Source/Core/FSTFirestoreClient.mm b/Firestore/Source/Core/FSTFirestoreClient.mm
index f1b4ffd..658cf57 100644
--- a/Firestore/Source/Core/FSTFirestoreClient.mm
+++ b/Firestore/Source/Core/FSTFirestoreClient.mm
@@ -56,6 +56,7 @@ using firebase::firestore::auth::CredentialsProvider;
 using firebase::firestore::auth::User;
 using firebase::firestore::core::DatabaseInfo;
 using firebase::firestore::model::DatabaseId;
+using firebase::firestore::model::DocumentKeySet;
 
 NS_ASSUME_NONNULL_BEGIN
 
@@ -312,9 +313,8 @@ NS_ASSUME_NONNULL_BEGIN
                                              NSError *_Nullable error))completion {
   [self.workerDispatchQueue dispatchAsync:^{
     FSTDocumentDictionary *docs = [self.localStore executeQuery:query.query];
-    FSTDocumentKeySet *remoteKeys = [FSTDocumentKeySet keySet];
 
-    FSTView *view = [[FSTView alloc] initWithQuery:query.query remoteDocuments:remoteKeys];
+    FSTView *view = [[FSTView alloc] initWithQuery:query.query remoteDocuments:DocumentKeySet{}];
     FSTViewDocumentChanges *viewDocChanges = [view computeChangesWithDocuments:docs];
     FSTViewChange *viewChange = [view applyChangesToDocuments:viewDocChanges];
     FSTAssert(viewChange.limboChanges.count == 0,
diff --git a/Firestore/Source/Core/FSTSyncEngine.mm b/Firestore/Source/Core/FSTSyncEngine.mm
index 8ec98e8..ed97d6c 100644
--- a/Firestore/Source/Core/FSTSyncEngine.mm
+++ b/Firestore/Source/Core/FSTSyncEngine.mm
@@ -21,6 +21,7 @@
 #include <map>
 #include <set>
 #include <unordered_map>
+#include <utility>
 
 #import "FIRFirestoreErrors.h"
 #import "Firestore/Source/Core/FSTQuery.h"
@@ -52,6 +53,7 @@ using firebase::firestore::core::TargetIdGenerator;
 using firebase::firestore::model::DocumentKey;
 using firebase::firestore::model::SnapshotVersion;
 using firebase::firestore::model::TargetId;
+using firebase::firestore::model::DocumentKeySet;
 
 NS_ASSUME_NONNULL_BEGIN
 
@@ -185,9 +187,9 @@ static const FSTListenSequenceNumber kIrrelevantSequenceNumber = -1;
 
   FSTQueryData *queryData = [self.localStore allocateQuery:query];
   FSTDocumentDictionary *docs = [self.localStore executeQuery:query];
-  FSTDocumentKeySet *remoteKeys = [self.localStore remoteDocumentKeysForTarget:queryData.targetID];
+  DocumentKeySet remoteKeys = [self.localStore remoteDocumentKeysForTarget:queryData.targetID];
 
-  FSTView *view = [[FSTView alloc] initWithQuery:query remoteDocuments:remoteKeys];
+  FSTView *view = [[FSTView alloc] initWithQuery:query remoteDocuments:std::move(remoteKeys)];
   FSTViewDocumentChanges *viewDocChanges = [view computeChangesWithDocuments:docs];
   FSTViewChange *viewChange = [view applyChangesToDocuments:viewDocChanges];
   FSTAssert(viewChange.limboChanges.count == 0,
@@ -347,13 +349,14 @@ static const FSTListenSequenceNumber kIrrelevantSequenceNumber = -1;
         [NSMutableDictionary dictionary];
     FSTDeletedDocument *doc =
         [FSTDeletedDocument documentWithKey:limboKey version:SnapshotVersion::None()];
-    FSTDocumentKeySet *limboDocuments = [[FSTDocumentKeySet keySet] setByAddingObject:doc.key];
-    FSTRemoteEvent *event = [[FSTRemoteEvent alloc] initWithSnapshotVersion:SnapshotVersion::None()
-                                                              targetChanges:targetChanges
-                                                            documentUpdates:{
-                                                              { limboKey, doc }
-                                                            }
-                                                             limboDocuments:limboDocuments];
+    DocumentKeySet limboDocuments = DocumentKeySet{doc.key};
+    FSTRemoteEvent *event =
+        [[FSTRemoteEvent alloc] initWithSnapshotVersion:SnapshotVersion::None()
+                                          targetChanges:targetChanges
+                                        documentUpdates:{
+                                          { limboKey, doc }
+                                        }
+                                         limboDocuments:std::move(limboDocuments)];
     [self applyRemoteEvent:event];
   } else {
     FSTQueryView *queryView = self.queryViewsByTarget[@(targetID)];
diff --git a/Firestore/Source/Core/FSTTransaction.mm b/Firestore/Source/Core/FSTTransaction.mm
index 57dcc48..5c36b20 100644
--- a/Firestore/Source/Core/FSTTransaction.mm
+++ b/Firestore/Source/Core/FSTTransaction.mm
@@ -24,19 +24,20 @@
 #import "FIRFirestoreErrors.h"
 #import "Firestore/Source/API/FSTUserDataConverter.h"
 #import "Firestore/Source/Model/FSTDocument.h"
-#import "Firestore/Source/Model/FSTDocumentKeySet.h"
 #import "Firestore/Source/Model/FSTMutation.h"
 #import "Firestore/Source/Remote/FSTDatastore.h"
 #import "Firestore/Source/Util/FSTAssert.h"
 #import "Firestore/Source/Util/FSTUsageValidation.h"
 
 #include "Firestore/core/src/firebase/firestore/model/document_key.h"
+#include "Firestore/core/src/firebase/firestore/model/document_key_set.h"
 #include "Firestore/core/src/firebase/firestore/model/precondition.h"
 #include "Firestore/core/src/firebase/firestore/model/snapshot_version.h"
 
 using firebase::firestore::model::DocumentKey;
 using firebase::firestore::model::Precondition;
 using firebase::firestore::model::SnapshotVersion;
+using firebase::firestore::model::DocumentKeySet;
 
 NS_ASSUME_NONNULL_BEGIN
 
@@ -216,15 +217,15 @@ NS_ASSUME_NONNULL_BEGIN
   }
 
   // Make a list of read documents that haven't been written.
-  FSTDocumentKeySet *unwritten = [FSTDocumentKeySet keySet];
+  DocumentKeySet unwritten;
   for (const auto &kv : _readVersions) {
-    unwritten = [unwritten setByAddingObject:kv.first];
+    unwritten = unwritten.insert(kv.first);
   };
   // For each mutation, note that the doc was written.
   for (FSTMutation *mutation in self.mutations) {
-    unwritten = [unwritten setByRemovingObject:mutation.key];
+    unwritten = unwritten.erase(mutation.key);
   }
-  if (unwritten.count) {
+  if (!unwritten.empty()) {
     // TODO(klimt): This is a temporary restriction, until "verify" is supported on the backend.
     completion([NSError
         errorWithDomain:FIRFirestoreErrorDomain
diff --git a/Firestore/Source/Core/FSTView.h b/Firestore/Source/Core/FSTView.h
index 431b863..fc6cead 100644
--- a/Firestore/Source/Core/FSTView.h
+++ b/Firestore/Source/Core/FSTView.h
@@ -18,9 +18,9 @@
 
 #import "Firestore/Source/Core/FSTTypes.h"
 #import "Firestore/Source/Model/FSTDocumentDictionary.h"
-#import "Firestore/Source/Model/FSTDocumentKeySet.h"
 
 #include "Firestore/core/src/firebase/firestore/model/document_key.h"
+#include "Firestore/core/src/firebase/firestore/model/document_key_set.h"
 
 @class FSTDocumentSet;
 @class FSTDocumentViewChangeSet;
@@ -38,6 +38,8 @@ NS_ASSUME_NONNULL_BEGIN
 
 - (instancetype)init NS_UNAVAILABLE;
 
+- (const firebase::firestore::model::DocumentKeySet &)mutatedKeys;
+
 /** The new set of docs that should be in the view. */
 @property(nonatomic, strong, readonly) FSTDocumentSet *documentSet;
 
@@ -50,8 +52,6 @@ NS_ASSUME_NONNULL_BEGIN
  */
 @property(nonatomic, assign, readonly) BOOL needsRefill;
 
-@property(nonatomic, strong, readonly) FSTDocumentKeySet *mutatedKeys;
-
 @end
 
 #pragma mark - FSTLimboDocumentChange
@@ -97,7 +97,8 @@ typedef NS_ENUM(NSInteger, FSTLimboDocumentChangeType) {
 - (instancetype)init NS_UNAVAILABLE;
 
 - (instancetype)initWithQuery:(FSTQuery *)query
-              remoteDocuments:(FSTDocumentKeySet *)remoteDocuments NS_DESIGNATED_INITIALIZER;
+              remoteDocuments:(firebase::firestore::model::DocumentKeySet)remoteDocuments
+    NS_DESIGNATED_INITIALIZER;
 
 /**
  * Iterates over a set of doc changes, applies the query limit, and computes what the new results
@@ -152,7 +153,7 @@ typedef NS_ENUM(NSInteger, FSTLimboDocumentChangeType) {
  * 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;
+- (const firebase::firestore::model::DocumentKeySet &)syncedDocuments;
 
 @end
 
diff --git a/Firestore/Source/Core/FSTView.mm b/Firestore/Source/Core/FSTView.mm
index d87951a..d254a82 100644
--- a/Firestore/Source/Core/FSTView.mm
+++ b/Firestore/Source/Core/FSTView.mm
@@ -29,6 +29,7 @@
 #include "Firestore/core/src/firebase/firestore/model/document_key.h"
 
 using firebase::firestore::model::DocumentKey;
+using firebase::firestore::model::DocumentKeySet;
 
 NS_ASSUME_NONNULL_BEGIN
 
@@ -40,26 +41,32 @@ NS_ASSUME_NONNULL_BEGIN
 - (instancetype)initWithDocumentSet:(FSTDocumentSet *)documentSet
                           changeSet:(FSTDocumentViewChangeSet *)changeSet
                         needsRefill:(BOOL)needsRefill
-                        mutatedKeys:(FSTDocumentKeySet *)mutatedKeys NS_DESIGNATED_INITIALIZER;
+                        mutatedKeys:(DocumentKeySet)mutatedKeys NS_DESIGNATED_INITIALIZER;
 
 @end
 
-@implementation FSTViewDocumentChanges
+@implementation FSTViewDocumentChanges {
+  DocumentKeySet _mutatedKeys;
+}
 
 - (instancetype)initWithDocumentSet:(FSTDocumentSet *)documentSet
                           changeSet:(FSTDocumentViewChangeSet *)changeSet
                         needsRefill:(BOOL)needsRefill
-                        mutatedKeys:(FSTDocumentKeySet *)mutatedKeys {
+                        mutatedKeys:(DocumentKeySet)mutatedKeys {
   self = [super init];
   if (self) {
     _documentSet = documentSet;
     _changeSet = changeSet;
     _needsRefill = needsRefill;
-    _mutatedKeys = mutatedKeys;
+    _mutatedKeys = std::move(mutatedKeys);
   }
   return self;
 }
 
+- (const DocumentKeySet &)mutatedKeys {
+  return _mutatedKeys;
+}
+
 @end
 
 #pragma mark - FSTLimboDocumentChange
@@ -165,32 +172,33 @@ static NSComparisonResult FSTCompareDocumentViewChangeTypes(FSTDocumentViewChang
 
 @property(nonatomic, strong) FSTDocumentSet *documentSet;
 
-/** Documents included in the remote target. */
-@property(nonatomic, strong) FSTDocumentKeySet *syncedDocuments;
-
-/** Documents in the view but not in the remote target */
-@property(nonatomic, strong) FSTDocumentKeySet *limboDocuments;
+@end
 
-/** Document Keys that have local changes. */
-@property(nonatomic, strong) FSTDocumentKeySet *mutatedKeys;
+@implementation FSTView {
+  /** Documents included in the remote target. */
+  DocumentKeySet _syncedDocuments;
 
-@end
+  /** Documents in the view but not in the remote target */
+  DocumentKeySet _limboDocuments;
 
-@implementation FSTView
+  /** Document Keys that have local changes. */
+  DocumentKeySet _mutatedKeys;
+}
 
-- (instancetype)initWithQuery:(FSTQuery *)query
-              remoteDocuments:(nonnull FSTDocumentKeySet *)remoteDocuments {
+- (instancetype)initWithQuery:(FSTQuery *)query remoteDocuments:(DocumentKeySet)remoteDocuments {
   self = [super init];
   if (self) {
     _query = query;
     _documentSet = [FSTDocumentSet documentSetWithComparator:query.comparator];
-    _syncedDocuments = remoteDocuments;
-    _limboDocuments = [FSTDocumentKeySet keySet];
-    _mutatedKeys = [FSTDocumentKeySet keySet];
+    _syncedDocuments = std::move(remoteDocuments);
   }
   return self;
 }
 
+- (const DocumentKeySet &)syncedDocuments {
+  return _syncedDocuments;
+}
+
 - (FSTViewDocumentChanges *)computeChangesWithDocuments:(FSTMaybeDocumentDictionary *)docChanges {
   return [self computeChangesWithDocuments:docChanges previousChanges:nil];
 }
@@ -202,8 +210,8 @@ static NSComparisonResult FSTCompareDocumentViewChangeTypes(FSTDocumentViewChang
       previousChanges ? previousChanges.changeSet : [FSTDocumentViewChangeSet changeSet];
   FSTDocumentSet *oldDocumentSet = previousChanges ? previousChanges.documentSet : self.documentSet;
 
-  __block FSTDocumentKeySet *newMutatedKeys =
-      previousChanges ? previousChanges.mutatedKeys : self.mutatedKeys;
+  __block DocumentKeySet newMutatedKeys =
+      previousChanges ? previousChanges.mutatedKeys : _mutatedKeys;
   __block FSTDocumentSet *newDocumentSet = oldDocumentSet;
   __block BOOL needsRefill = NO;
 
@@ -236,13 +244,13 @@ static NSComparisonResult FSTCompareDocumentViewChangeTypes(FSTDocumentViewChang
     if (newDoc) {
       newDocumentSet = [newDocumentSet documentSetByAddingDocument:newDoc];
       if (newDoc.hasLocalMutations) {
-        newMutatedKeys = [newMutatedKeys setByAddingObject:key];
+        newMutatedKeys = newMutatedKeys.insert(key);
       } else {
-        newMutatedKeys = [newMutatedKeys setByRemovingObject:key];
+        newMutatedKeys = newMutatedKeys.erase(key);
       }
     } else {
       newDocumentSet = [newDocumentSet documentSetByRemovingKey:key];
-      newMutatedKeys = [newMutatedKeys setByRemovingObject:key];
+      newMutatedKeys = newMutatedKeys.erase(key);
     }
 
     // Calculate change
@@ -311,7 +319,7 @@ static NSComparisonResult FSTCompareDocumentViewChangeTypes(FSTDocumentViewChang
 
   FSTDocumentSet *oldDocuments = self.documentSet;
   self.documentSet = docChanges.documentSet;
-  self.mutatedKeys = docChanges.mutatedKeys;
+  _mutatedKeys = docChanges.mutatedKeys;
 
   // Sort changes based on type and query comparator.
   NSArray<FSTDocumentViewChange *> *changes = [docChanges.changeSet changes];
@@ -325,7 +333,7 @@ static NSComparisonResult FSTCompareDocumentViewChangeTypes(FSTDocumentViewChang
   }];
   [self applyTargetChange:targetChange];
   NSArray<FSTLimboDocumentChange *> *limboChanges = [self updateLimboDocuments];
-  BOOL synced = self.limboDocuments.count == 0 && self.isCurrent;
+  BOOL synced = _limboDocuments.empty() && self.isCurrent;
   FSTSyncState newSyncState = synced ? FSTSyncStateSynced : FSTSyncStateLocal;
   BOOL syncStateChanged = newSyncState != self.syncState;
   self.syncState = newSyncState;
@@ -340,7 +348,7 @@ static NSComparisonResult FSTCompareDocumentViewChangeTypes(FSTDocumentViewChang
                                   oldDocuments:oldDocuments
                                documentChanges:changes
                                      fromCache:newSyncState == FSTSyncStateLocal
-                              hasPendingWrites:!docChanges.mutatedKeys.isEmpty
+                              hasPendingWrites:!docChanges.mutatedKeys.empty()
                               syncStateChanged:syncStateChanged];
 
     return [FSTViewChange changeWithSnapshot:snapshot limboChanges:limboChanges];
@@ -358,7 +366,7 @@ static NSComparisonResult FSTCompareDocumentViewChangeTypes(FSTDocumentViewChang
                                           initWithDocumentSet:self.documentSet
                                                     changeSet:[FSTDocumentViewChangeSet changeSet]
                                                   needsRefill:NO
-                                                  mutatedKeys:self.mutatedKeys]];
+                                                  mutatedKeys:_mutatedKeys]];
   } else {
     // No effect, just return a no-op FSTViewChange.
     return [[FSTViewChange alloc] initWithSnapshot:nil limboChanges:@[]];
@@ -370,7 +378,7 @@ static NSComparisonResult FSTCompareDocumentViewChangeTypes(FSTDocumentViewChang
 /** Returns whether the doc for the given key should be in limbo. */
 - (BOOL)shouldBeLimboDocumentKey:(const DocumentKey &)key {
   // If the remote end says it's part of this query, it's not in limbo.
-  if ([self.syncedDocuments containsObject:key]) {
+  if (_syncedDocuments.contains(key)) {
     return NO;
   }
   // The local store doesn't think it's a result, so it shouldn't be in limbo.
@@ -395,16 +403,14 @@ static NSComparisonResult FSTCompareDocumentViewChangeTypes(FSTDocumentViewChang
   if (targetChange) {
     FSTTargetMapping *targetMapping = targetChange.mapping;
     if ([targetMapping isKindOfClass:[FSTResetMapping class]]) {
-      self.syncedDocuments = ((FSTResetMapping *)targetMapping).documents;
+      _syncedDocuments = ((FSTResetMapping *)targetMapping).documents;
     } else if ([targetMapping isKindOfClass:[FSTUpdateMapping class]]) {
-      [((FSTUpdateMapping *)targetMapping).addedDocuments
-          enumerateObjectsUsingBlock:^(FSTDocumentKey *key, BOOL *stop) {
-            self.syncedDocuments = [self.syncedDocuments setByAddingObject:key];
-          }];
-      [((FSTUpdateMapping *)targetMapping).removedDocuments
-          enumerateObjectsUsingBlock:^(FSTDocumentKey *key, BOOL *stop) {
-            self.syncedDocuments = [self.syncedDocuments setByRemovingObject:key];
-          }];
+      for (const DocumentKey &key : ((FSTUpdateMapping *)targetMapping).addedDocuments) {
+        _syncedDocuments = _syncedDocuments.insert(key);
+      }
+      for (const DocumentKey &key : ((FSTUpdateMapping *)targetMapping).removedDocuments) {
+        _syncedDocuments = _syncedDocuments.erase(key);
+      }
     }
 
     switch (targetChange.currentStatusUpdate) {
@@ -428,29 +434,29 @@ static NSComparisonResult FSTCompareDocumentViewChangeTypes(FSTDocumentViewChang
   }
 
   // TODO(klimt): Do this incrementally so that it's not quadratic when updating many documents.
-  FSTDocumentKeySet *oldLimboDocuments = self.limboDocuments;
-  self.limboDocuments = [FSTDocumentKeySet keySet];
+  DocumentKeySet oldLimboDocuments = std::move(_limboDocuments);
+  _limboDocuments = DocumentKeySet{};
   for (FSTDocument *doc in self.documentSet.documentEnumerator) {
     if ([self shouldBeLimboDocumentKey:doc.key]) {
-      self.limboDocuments = [self.limboDocuments setByAddingObject:doc.key];
+      _limboDocuments = _limboDocuments.insert(doc.key);
     }
   }
 
   // Diff the new limbo docs with the old limbo docs.
   NSMutableArray<FSTLimboDocumentChange *> *changes =
-      [NSMutableArray arrayWithCapacity:(oldLimboDocuments.count + self.limboDocuments.count)];
-  [oldLimboDocuments enumerateObjectsUsingBlock:^(FSTDocumentKey *key, BOOL *stop) {
-    if (![self.limboDocuments containsObject:key]) {
+      [NSMutableArray arrayWithCapacity:(oldLimboDocuments.size() + _limboDocuments.size())];
+  for (const DocumentKey &key : oldLimboDocuments) {
+    if (!_limboDocuments.contains(key)) {
       [changes addObject:[FSTLimboDocumentChange changeWithType:FSTLimboDocumentChangeTypeRemoved
                                                             key:key]];
     }
-  }];
-  [self.limboDocuments enumerateObjectsUsingBlock:^(FSTDocumentKey *key, BOOL *stop) {
-    if (![oldLimboDocuments containsObject:key]) {
+  }
+  for (const DocumentKey &key : _limboDocuments) {
+    if (!oldLimboDocuments.contains(key)) {
       [changes addObject:[FSTLimboDocumentChange changeWithType:FSTLimboDocumentChangeTypeAdded
                                                             key:key]];
     }
-  }];
+  }
   return changes;
 }
 
diff --git a/Firestore/Source/Local/FSTLevelDBQueryCache.mm b/Firestore/Source/Local/FSTLevelDBQueryCache.mm
index 85b659a..97b218d 100644
--- a/Firestore/Source/Local/FSTLevelDBQueryCache.mm
+++ b/Firestore/Source/Local/FSTLevelDBQueryCache.mm
@@ -41,6 +41,7 @@ using firebase::firestore::model::SnapshotVersion;
 using leveldb::DB;
 using leveldb::Slice;
 using leveldb::Status;
+using firebase::firestore::model::DocumentKeySet;
 
 @interface FSTLevelDBQueryCache ()
 
@@ -282,30 +283,28 @@ using leveldb::Status;
 
 #pragma mark Matching Key tracking
 
-- (void)addMatchingKeys:(FSTDocumentKeySet *)keys forTargetID:(FSTTargetID)targetID {
+- (void)addMatchingKeys:(const DocumentKeySet &)keys forTargetID:(FSTTargetID)targetID {
   // Store an empty value in the index which is equivalent to serializing a GPBEmpty message. In the
   // future if we wanted to store some other kind of value here, we can parse these empty values as
   // with some other protocol buffer (and the parser will see all default values).
   std::string emptyBuffer;
 
-  [keys enumerateObjectsUsingBlock:^(FSTDocumentKey *documentKey, BOOL *stop) {
+  for (const DocumentKey &key : keys) {
     self->_db.currentTransaction->Put(
-        [FSTLevelDBTargetDocumentKey keyWithTargetID:targetID documentKey:documentKey],
-        emptyBuffer);
+        [FSTLevelDBTargetDocumentKey keyWithTargetID:targetID documentKey:key], emptyBuffer);
     self->_db.currentTransaction->Put(
-        [FSTLevelDBDocumentTargetKey keyWithDocumentKey:documentKey targetID:targetID],
-        emptyBuffer);
-  }];
+        [FSTLevelDBDocumentTargetKey keyWithDocumentKey:key targetID:targetID], emptyBuffer);
+  }
 }
 
-- (void)removeMatchingKeys:(FSTDocumentKeySet *)keys forTargetID:(FSTTargetID)targetID {
-  [keys enumerateObjectsUsingBlock:^(FSTDocumentKey *key, BOOL *stop) {
+- (void)removeMatchingKeys:(const DocumentKeySet &)keys forTargetID:(FSTTargetID)targetID {
+  for (const DocumentKey &key : keys) {
     self->_db.currentTransaction->Delete(
         [FSTLevelDBTargetDocumentKey keyWithTargetID:targetID documentKey:key]);
     self->_db.currentTransaction->Delete(
         [FSTLevelDBDocumentTargetKey keyWithDocumentKey:key targetID:targetID]);
     [self.garbageCollector addPotentialGarbageKey:key];
-  }];
+  }
 }
 
 - (void)removeMatchingKeysForTargetID:(FSTTargetID)targetID {
@@ -331,12 +330,12 @@ using leveldb::Status;
   }
 }
 
-- (FSTDocumentKeySet *)matchingKeysForTargetID:(FSTTargetID)targetID {
+- (DocumentKeySet)matchingKeysForTargetID:(FSTTargetID)targetID {
   std::string indexPrefix = [FSTLevelDBTargetDocumentKey keyPrefixWithTargetID:targetID];
   auto indexIterator = _db.currentTransaction->NewIterator();
   indexIterator->Seek(indexPrefix);
 
-  FSTDocumentKeySet *result = [FSTDocumentKeySet keySet];
+  DocumentKeySet result;
   FSTLevelDBTargetDocumentKey *rowKey = [[FSTLevelDBTargetDocumentKey alloc] init];
   for (; indexIterator->Valid(); indexIterator->Next()) {
     absl::string_view indexKey = indexIterator->key();
@@ -346,7 +345,7 @@ using leveldb::Status;
       break;
     }
 
-    result = [result setByAddingObject:rowKey.documentKey];
+    result = result.insert(rowKey.documentKey);
   }
 
   return result;
diff --git a/Firestore/Source/Local/FSTLocalDocumentsView.h b/Firestore/Source/Local/FSTLocalDocumentsView.h
index e75e0f3..bb5bb22 100644
--- a/Firestore/Source/Local/FSTLocalDocumentsView.h
+++ b/Firestore/Source/Local/FSTLocalDocumentsView.h
@@ -17,9 +17,9 @@
 #import <Foundation/Foundation.h>
 
 #import "Firestore/Source/Model/FSTDocumentDictionary.h"
-#import "Firestore/Source/Model/FSTDocumentKeySet.h"
 
 #include "Firestore/core/src/firebase/firestore/model/document_key.h"
+#include "Firestore/core/src/firebase/firestore/model/document_key_set.h"
 
 @class FSTMaybeDocument;
 @class FSTQuery;
@@ -53,7 +53,8 @@ NS_ASSUME_NONNULL_BEGIN
  * If we don't have cached state for a document in `keys`, a FSTDeletedDocument will be stored
  * for that key in the resulting set.
  */
-- (FSTMaybeDocumentDictionary *)documentsForKeys:(FSTDocumentKeySet *)keys;
+- (FSTMaybeDocumentDictionary *)documentsForKeys:
+    (const firebase::firestore::model::DocumentKeySet &)keys;
 
 /** Performs a query against the local view of all documents. */
 - (FSTDocumentDictionary *)documentsMatchingQuery:(FSTQuery *)query;
diff --git a/Firestore/Source/Local/FSTLocalDocumentsView.mm b/Firestore/Source/Local/FSTLocalDocumentsView.mm
index 703d74d..471840a 100644
--- a/Firestore/Source/Local/FSTLocalDocumentsView.mm
+++ b/Firestore/Source/Local/FSTLocalDocumentsView.mm
@@ -32,6 +32,7 @@
 using firebase::firestore::model::DocumentKey;
 using firebase::firestore::model::ResourcePath;
 using firebase::firestore::model::SnapshotVersion;
+using firebase::firestore::model::DocumentKeySet;
 
 NS_ASSUME_NONNULL_BEGIN
 
@@ -65,9 +66,9 @@ NS_ASSUME_NONNULL_BEGIN
   return [self localDocument:remoteDoc key:key];
 }
 
-- (FSTMaybeDocumentDictionary *)documentsForKeys:(FSTDocumentKeySet *)keys {
+- (FSTMaybeDocumentDictionary *)documentsForKeys:(const DocumentKeySet &)keys {
   FSTMaybeDocumentDictionary *results = [FSTMaybeDocumentDictionary maybeDocumentDictionary];
-  for (FSTDocumentKey *key in keys.objectEnumerator) {
+  for (const DocumentKey &key : keys) {
     // TODO(mikelehen): PERF: Consider fetching all remote documents at once rather than one-by-one.
     FSTMaybeDocument *maybeDoc = [self documentForKey:key];
     // TODO(http://b/32275378): Don't conflate missing / deleted.
@@ -106,7 +107,7 @@ NS_ASSUME_NONNULL_BEGIN
 
   // Now use the mutation queue to discover any other documents that may match the query after
   // applying mutations.
-  FSTDocumentKeySet *matchingKeys = [FSTDocumentKeySet keySet];
+  DocumentKeySet matchingKeys;
   NSArray<FSTMutationBatch *> *matchingMutationBatches =
       [self.mutationQueue allMutationBatchesAffectingQuery:query];
   for (FSTMutationBatch *batch in matchingMutationBatches) {
@@ -115,13 +116,13 @@ NS_ASSUME_NONNULL_BEGIN
 
       // If the key is already in the results, we can skip it.
       if (![results containsKey:mutation.key]) {
-        matchingKeys = [matchingKeys setByAddingObject:mutation.key];
+        matchingKeys = matchingKeys.insert(mutation.key);
       }
     }
   }
 
   // Now add in results for the matchingKeys.
-  for (FSTDocumentKey *key in matchingKeys.objectEnumerator) {
+  for (const DocumentKey &key : matchingKeys) {
     FSTMaybeDocument *doc = [self documentForKey:key];
     if ([doc isKindOfClass:[FSTDocument class]]) {
       results = [results dictionaryBySettingObject:(FSTDocument *)doc forKey:key];
diff --git a/Firestore/Source/Local/FSTLocalStore.h b/Firestore/Source/Local/FSTLocalStore.h
index 84aeb58..1312c28 100644
--- a/Firestore/Source/Local/FSTLocalStore.h
+++ b/Firestore/Source/Local/FSTLocalStore.h
@@ -18,11 +18,11 @@
 
 #import "Firestore/Source/Core/FSTTypes.h"
 #import "Firestore/Source/Model/FSTDocumentDictionary.h"
-#import "Firestore/Source/Model/FSTDocumentKeySet.h"
 #import "Firestore/Source/Model/FSTDocumentVersionDictionary.h"
 
 #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/document_key_set.h"
 #include "Firestore/core/src/firebase/firestore/model/snapshot_version.h"
 
 @class FSTLocalViewChanges;
@@ -157,7 +157,7 @@ NS_ASSUME_NONNULL_BEGIN
  * Returns the keys of the documents that are associated with the given targetID in the remote
  * table.
  */
-- (FSTDocumentKeySet *)remoteDocumentKeysForTarget:(FSTTargetID)targetID;
+- (firebase::firestore::model::DocumentKeySet)remoteDocumentKeysForTarget:(FSTTargetID)targetID;
 
 /**
  * Collects garbage if necessary.
diff --git a/Firestore/Source/Local/FSTLocalStore.mm b/Firestore/Source/Local/FSTLocalStore.mm
index 92462f0..e721579 100644
--- a/Firestore/Source/Local/FSTLocalStore.mm
+++ b/Firestore/Source/Local/FSTLocalStore.mm
@@ -48,6 +48,7 @@ using firebase::firestore::auth::User;
 using firebase::firestore::core::TargetIdGenerator;
 using firebase::firestore::model::DocumentKey;
 using firebase::firestore::model::SnapshotVersion;
+using firebase::firestore::model::DocumentKeySet;
 
 NS_ASSUME_NONNULL_BEGIN
 
@@ -186,11 +187,11 @@ NS_ASSUME_NONNULL_BEGIN
                                              mutationQueue:self.mutationQueue];
 
     // Union the old/new changed keys.
-    FSTDocumentKeySet *changedKeys = [FSTDocumentKeySet keySet];
+    DocumentKeySet changedKeys;
     for (NSArray<FSTMutationBatch *> *batches in @[ oldBatches, newBatches ]) {
       for (FSTMutationBatch *batch in batches) {
         for (FSTMutation *mutation in batch.mutations) {
-          changedKeys = [changedKeys setByAddingObject:mutation.key];
+          changedKeys = changedKeys.insert(mutation.key);
         }
       }
     }
@@ -205,7 +206,7 @@ NS_ASSUME_NONNULL_BEGIN
     FIRTimestamp *localWriteTime = [FIRTimestamp timestamp];
     FSTMutationBatch *batch =
         [self.mutationQueue addMutationBatchWithWriteTime:localWriteTime mutations:mutations];
-    FSTDocumentKeySet *keys = [batch keys];
+    DocumentKeySet keys = [batch keys];
     FSTMaybeDocumentDictionary *changedDocuments = [self.localDocuments documentsForKeys:keys];
     return [FSTLocalWriteResult resultForBatchID:batch.batchID changes:changedDocuments];
   });
@@ -217,10 +218,9 @@ NS_ASSUME_NONNULL_BEGIN
 
     [mutationQueue acknowledgeBatch:batchResult.batch streamToken:batchResult.streamToken];
 
-    FSTDocumentKeySet *affected;
+    DocumentKeySet affected;
     if ([self shouldHoldBatchResultWithVersion:batchResult.commitVersion]) {
       [self.heldBatchResults addObject:batchResult];
-      affected = [FSTDocumentKeySet keySet];
     } else {
       affected = [self releaseBatchResults:@[ batchResult ]];
     }
@@ -239,7 +239,7 @@ NS_ASSUME_NONNULL_BEGIN
     FSTBatchID lastAcked = [self.mutationQueue highestAcknowledgedBatchID];
     FSTAssert(batchID > lastAcked, @"Acknowledged batches can't be rejected.");
 
-    FSTDocumentKeySet *affected = [self removeMutationBatch:toReject];
+    DocumentKeySet affected = [self removeMutationBatch:toReject];
 
     [self.mutationQueue performConsistencyCheck];
 
@@ -307,11 +307,11 @@ NS_ASSUME_NONNULL_BEGIN
     }];
 
     // TODO(klimt): This could probably be an NSMutableDictionary.
-    FSTDocumentKeySet *changedDocKeys = [FSTDocumentKeySet keySet];
+    DocumentKeySet changedDocKeys;
     for (const auto &kv : remoteEvent.documentUpdates) {
       const DocumentKey &key = kv.first;
       FSTMaybeDocument *doc = kv.second;
-      changedDocKeys = [changedDocKeys setByAddingObject:key];
+      changedDocKeys = changedDocKeys.insert(key);
       FSTMaybeDocument *existingDoc = [self.remoteDocumentCache entryForKey:key];
       // Make sure we don't apply an old document version to the remote cache, though we
       // make an exception for SnapshotVersion::None() which can happen for manufactured
@@ -345,13 +345,13 @@ NS_ASSUME_NONNULL_BEGIN
       [self.queryCache setLastRemoteSnapshotVersion:remoteVersion];
     }
 
-    FSTDocumentKeySet *releasedWriteKeys = [self releaseHeldBatchResults];
+    DocumentKeySet releasedWriteKeys = [self releaseHeldBatchResults];
 
     // Union the two key sets.
-    __block FSTDocumentKeySet *keysToRecalc = changedDocKeys;
-    [releasedWriteKeys enumerateObjectsUsingBlock:^(FSTDocumentKey *key, BOOL *stop) {
-      keysToRecalc = [keysToRecalc setByAddingObject:key];
-    }];
+    DocumentKeySet keysToRecalc = changedDocKeys;
+    for (const DocumentKey &key : releasedWriteKeys) {
+      keysToRecalc = keysToRecalc.insert(key);
+    }
 
     return [self.localDocuments documentsForKeys:keysToRecalc];
   });
@@ -430,8 +430,8 @@ NS_ASSUME_NONNULL_BEGIN
   });
 }
 
-- (FSTDocumentKeySet *)remoteDocumentKeysForTarget:(FSTTargetID)targetID {
-  return self.persistence.run("RemoteDocumentKeysForTarget", [&]() -> FSTDocumentKeySet * {
+- (DocumentKeySet)remoteDocumentKeysForTarget:(FSTTargetID)targetID {
+  return self.persistence.run("RemoteDocumentKeysForTarget", [&]() -> DocumentKeySet {
     return [self.queryCache matchingKeysForTargetID:targetID];
   });
 }
@@ -455,7 +455,7 @@ NS_ASSUME_NONNULL_BEGIN
  *
  * @return the set of keys of docs that were modified by those writes.
  */
-- (FSTDocumentKeySet *)releaseHeldBatchResults {
+- (DocumentKeySet)releaseHeldBatchResults {
   NSMutableArray<FSTMutationBatchResult *> *toRelease = [NSMutableArray array];
   for (FSTMutationBatchResult *batchResult in self.heldBatchResults) {
     if (![self isRemoteUpToVersion:batchResult.commitVersion]) {
@@ -465,7 +465,7 @@ NS_ASSUME_NONNULL_BEGIN
   }
 
   if (toRelease.count == 0) {
-    return [FSTDocumentKeySet keySet];
+    return DocumentKeySet{};
   } else {
     [self.heldBatchResults removeObjectsInRange:NSMakeRange(0, toRelease.count)];
     return [self releaseBatchResults:toRelease];
@@ -482,7 +482,7 @@ NS_ASSUME_NONNULL_BEGIN
   return ![self isRemoteUpToVersion:version] || self.heldBatchResults.count > 0;
 }
 
-- (FSTDocumentKeySet *)releaseBatchResults:(NSArray<FSTMutationBatchResult *> *)batchResults {
+- (DocumentKeySet)releaseBatchResults:(NSArray<FSTMutationBatchResult *> *)batchResults {
   NSMutableArray<FSTMutationBatch *> *batches = [NSMutableArray array];
   for (FSTMutationBatchResult *batchResult in batchResults) {
     [self applyBatchResult:batchResult];
@@ -492,38 +492,35 @@ NS_ASSUME_NONNULL_BEGIN
   return [self removeMutationBatches:batches];
 }
 
-- (FSTDocumentKeySet *)removeMutationBatch:(FSTMutationBatch *)batch {
+- (DocumentKeySet)removeMutationBatch:(FSTMutationBatch *)batch {
   return [self removeMutationBatches:@[ batch ]];
 }
 
 /** Removes all the mutation batches named in the given array. */
-- (FSTDocumentKeySet *)removeMutationBatches:(NSArray<FSTMutationBatch *> *)batches {
-  // TODO(klimt): Could this be an NSMutableDictionary?
-  __block FSTDocumentKeySet *affectedDocs = [FSTDocumentKeySet keySet];
-
+- (DocumentKeySet)removeMutationBatches:(NSArray<FSTMutationBatch *> *)batches {
+  DocumentKeySet affectedDocs;
   for (FSTMutationBatch *batch in batches) {
     for (FSTMutation *mutation in batch.mutations) {
       const DocumentKey &key = mutation.key;
-      affectedDocs = [affectedDocs setByAddingObject:key];
+      affectedDocs = affectedDocs.insert(key);
     }
   }
 
   [self.mutationQueue removeMutationBatches:batches];
-
   return affectedDocs;
 }
 
 - (void)applyBatchResult:(FSTMutationBatchResult *)batchResult {
   FSTMutationBatch *batch = batchResult.batch;
-  FSTDocumentKeySet *docKeys = batch.keys;
-  [docKeys enumerateObjectsUsingBlock:^(FSTDocumentKey *docKey, BOOL *stop) {
+  DocumentKeySet docKeys = batch.keys;
+  for (const DocumentKey &docKey : docKeys) {
     FSTMaybeDocument *_Nullable remoteDoc = [self.remoteDocumentCache entryForKey:docKey];
     FSTMaybeDocument *_Nullable doc = remoteDoc;
     // TODO(zxu): Once ported to use C++ version of FSTMutationBatchResult, we should be able to
     // check ackVersion instead, which will be an absl::optional type.
-    FSTAssert(batchResult.docVersions[docKey],
+    FSTAssert(batchResult.docVersions[static_cast<FSTDocumentKey *>(docKey)],
               @"docVersions should contain every doc in the write.");
-    SnapshotVersion ackVersion = batchResult.docVersions[docKey];
+    SnapshotVersion ackVersion = batchResult.docVersions[static_cast<FSTDocumentKey *>(docKey)];
     if (!doc || SnapshotVersion{doc.version} < ackVersion) {
       doc = [batch applyTo:doc documentKey:docKey mutationBatchResult:batchResult];
       if (!doc) {
@@ -533,7 +530,7 @@ NS_ASSUME_NONNULL_BEGIN
         [self.remoteDocumentCache addEntry:doc];
       }
     }
-  }];
+  }
 }
 
 @end
diff --git a/Firestore/Source/Local/FSTLocalViewChanges.h b/Firestore/Source/Local/FSTLocalViewChanges.h
index eb84642..143d010 100644
--- a/Firestore/Source/Local/FSTLocalViewChanges.h
+++ b/Firestore/Source/Local/FSTLocalViewChanges.h
@@ -16,7 +16,7 @@
 
 #import <Foundation/Foundation.h>
 
-#import "Firestore/Source/Model/FSTDocumentKeySet.h"
+#include "Firestore/core/src/firebase/firestore/model/document_key_set.h"
 
 @class FSTDocumentSet;
 @class FSTMutation;
@@ -34,16 +34,17 @@ NS_ASSUME_NONNULL_BEGIN
 @interface FSTLocalViewChanges : NSObject
 
 + (instancetype)changesForQuery:(FSTQuery *)query
-                      addedKeys:(FSTDocumentKeySet *)addedKeys
-                    removedKeys:(FSTDocumentKeySet *)removedKeys;
+                      addedKeys:(firebase::firestore::model::DocumentKeySet)addedKeys
+                    removedKeys:(firebase::firestore::model::DocumentKeySet)removedKeys;
 
 + (instancetype)changesForViewSnapshot:(FSTViewSnapshot *)viewSnapshot;
 
 - (id)init NS_UNAVAILABLE;
 
 @property(nonatomic, strong, readonly) FSTQuery *query;
-@property(nonatomic, strong) FSTDocumentKeySet *addedKeys;
-@property(nonatomic, strong) FSTDocumentKeySet *removedKeys;
+
+- (const firebase::firestore::model::DocumentKeySet &)addedKeys;
+- (const firebase::firestore::model::DocumentKeySet &)removedKeys;
 
 @end
 
diff --git a/Firestore/Source/Local/FSTLocalViewChanges.mm b/Firestore/Source/Local/FSTLocalViewChanges.mm
index 9a7f445..eb6b259 100644
--- a/Firestore/Source/Local/FSTLocalViewChanges.mm
+++ b/Firestore/Source/Local/FSTLocalViewChanges.mm
@@ -16,31 +16,38 @@
 
 #import "Firestore/Source/Local/FSTLocalViewChanges.h"
 
+#include <utility>
+
 #import "Firestore/Source/Core/FSTViewSnapshot.h"
 #import "Firestore/Source/Model/FSTDocument.h"
 
+using firebase::firestore::model::DocumentKeySet;
+
 NS_ASSUME_NONNULL_BEGIN
 
 @interface FSTLocalViewChanges ()
 - (instancetype)initWithQuery:(FSTQuery *)query
-                    addedKeys:(FSTDocumentKeySet *)addedKeys
-                  removedKeys:(FSTDocumentKeySet *)removedKeys NS_DESIGNATED_INITIALIZER;
+                    addedKeys:(DocumentKeySet)addedKeys
+                  removedKeys:(DocumentKeySet)removedKeys NS_DESIGNATED_INITIALIZER;
 @end
 
-@implementation FSTLocalViewChanges
+@implementation FSTLocalViewChanges {
+  DocumentKeySet _addedKeys;
+  DocumentKeySet _removedKeys;
+}
 
 + (instancetype)changesForViewSnapshot:(FSTViewSnapshot *)viewSnapshot {
-  FSTDocumentKeySet *addedKeys = [FSTDocumentKeySet keySet];
-  FSTDocumentKeySet *removedKeys = [FSTDocumentKeySet keySet];
+  DocumentKeySet addedKeys;
+  DocumentKeySet removedKeys;
 
   for (FSTDocumentViewChange *docChange in viewSnapshot.documentChanges) {
     switch (docChange.type) {
       case FSTDocumentViewChangeTypeAdded:
-        addedKeys = [addedKeys setByAddingObject:docChange.document.key];
+        addedKeys = addedKeys.insert(docChange.document.key);
         break;
 
       case FSTDocumentViewChangeTypeRemoved:
-        removedKeys = [removedKeys setByAddingObject:docChange.document.key];
+        removedKeys = removedKeys.insert(docChange.document.key);
         break;
 
       default:
@@ -49,28 +56,39 @@ NS_ASSUME_NONNULL_BEGIN
     }
   }
 
-  return [self changesForQuery:viewSnapshot.query addedKeys:addedKeys removedKeys:removedKeys];
+  return [self changesForQuery:viewSnapshot.query
+                     addedKeys:std::move(addedKeys)
+                   removedKeys:std::move(removedKeys)];
 }
 
 + (instancetype)changesForQuery:(FSTQuery *)query
-                      addedKeys:(FSTDocumentKeySet *)addedKeys
-                    removedKeys:(FSTDocumentKeySet *)removedKeys {
-  return
-      [[FSTLocalViewChanges alloc] initWithQuery:query addedKeys:addedKeys removedKeys:removedKeys];
+                      addedKeys:(DocumentKeySet)addedKeys
+                    removedKeys:(DocumentKeySet)removedKeys {
+  return [[FSTLocalViewChanges alloc] initWithQuery:query
+                                          addedKeys:std::move(addedKeys)
+                                        removedKeys:std::move(removedKeys)];
 }
 
 - (instancetype)initWithQuery:(FSTQuery *)query
-                    addedKeys:(FSTDocumentKeySet *)addedKeys
-                  removedKeys:(FSTDocumentKeySet *)removedKeys {
+                    addedKeys:(DocumentKeySet)addedKeys
+                  removedKeys:(DocumentKeySet)removedKeys {
   self = [super init];
   if (self) {
     _query = query;
-    _addedKeys = addedKeys;
-    _removedKeys = removedKeys;
+    _addedKeys = std::move(addedKeys);
+    _removedKeys = std::move(removedKeys);
   }
   return self;
 }
 
+- (const DocumentKeySet &)addedKeys {
+  return _addedKeys;
+}
+
+- (const DocumentKeySet &)removedKeys {
+  return _removedKeys;
+}
+
 @end
 
 NS_ASSUME_NONNULL_END
diff --git a/Firestore/Source/Local/FSTMemoryMutationQueue.mm b/Firestore/Source/Local/FSTMemoryMutationQueue.mm
index 8028bb3..b0c8760 100644
--- a/Firestore/Source/Local/FSTMemoryMutationQueue.mm
+++ b/Firestore/Source/Local/FSTMemoryMutationQueue.mm
@@ -21,6 +21,7 @@
 #import "Firestore/Source/Model/FSTMutation.h"
 #import "Firestore/Source/Model/FSTMutationBatch.h"
 #import "Firestore/Source/Util/FSTAssert.h"
+#import "Firestore/third_party/Immutable/FSTImmutableSortedSet.h"
 
 #include "Firestore/core/src/firebase/firestore/model/document_key.h"
 #include "Firestore/core/src/firebase/firestore/model/resource_path.h"
diff --git a/Firestore/Source/Local/FSTMemoryQueryCache.mm b/Firestore/Source/Local/FSTMemoryQueryCache.mm
index 57c2703..991faa9 100644
--- a/Firestore/Source/Local/FSTMemoryQueryCache.mm
+++ b/Firestore/Source/Local/FSTMemoryQueryCache.mm
@@ -26,6 +26,7 @@
 #include "Firestore/core/src/firebase/firestore/model/snapshot_version.h"
 
 using firebase::firestore::model::SnapshotVersion;
+using firebase::firestore::model::DocumentKeySet;
 
 NS_ASSUME_NONNULL_BEGIN
 
@@ -116,11 +117,11 @@ NS_ASSUME_NONNULL_BEGIN
 
 #pragma mark Reference tracking
 
-- (void)addMatchingKeys:(FSTDocumentKeySet *)keys forTargetID:(FSTTargetID)targetID {
+- (void)addMatchingKeys:(const DocumentKeySet &)keys forTargetID:(FSTTargetID)targetID {
   [self.references addReferencesToKeys:keys forID:targetID];
 }
 
-- (void)removeMatchingKeys:(FSTDocumentKeySet *)keys forTargetID:(FSTTargetID)targetID {
+- (void)removeMatchingKeys:(const DocumentKeySet &)keys forTargetID:(FSTTargetID)targetID {
   [self.references removeReferencesToKeys:keys forID:targetID];
 }
 
@@ -128,7 +129,7 @@ NS_ASSUME_NONNULL_BEGIN
   [self.references removeReferencesForID:targetID];
 }
 
-- (FSTDocumentKeySet *)matchingKeysForTargetID:(FSTTargetID)targetID {
+- (DocumentKeySet)matchingKeysForTargetID:(FSTTargetID)targetID {
   return [self.references referencedKeysForID:targetID];
 }
 
diff --git a/Firestore/Source/Local/FSTQueryCache.h b/Firestore/Source/Local/FSTQueryCache.h
index de41dd7..1ad46aa 100644
--- a/Firestore/Source/Local/FSTQueryCache.h
+++ b/Firestore/Source/Local/FSTQueryCache.h
@@ -18,8 +18,8 @@
 
 #import "Firestore/Source/Core/FSTTypes.h"
 #import "Firestore/Source/Local/FSTGarbageCollector.h"
-#import "Firestore/Source/Model/FSTDocumentKeySet.h"
 
+#include "Firestore/core/src/firebase/firestore/model/document_key_set.h"
 #include "Firestore/core/src/firebase/firestore/model/snapshot_version.h"
 
 @class FSTDocumentSet;
@@ -105,15 +105,17 @@ NS_ASSUME_NONNULL_BEGIN
 - (nullable FSTQueryData *)queryDataForQuery:(FSTQuery *)query;
 
 /** Adds the given document keys to cached query results of the given target ID. */
-- (void)addMatchingKeys:(FSTDocumentKeySet *)keys forTargetID:(FSTTargetID)targetID;
+- (void)addMatchingKeys:(const firebase::firestore::model::DocumentKeySet &)keys
+            forTargetID:(FSTTargetID)targetID;
 
 /** Removes the given document keys from the cached query results of the given target ID. */
-- (void)removeMatchingKeys:(FSTDocumentKeySet *)keys forTargetID:(FSTTargetID)targetID;
+- (void)removeMatchingKeys:(const firebase::firestore::model::DocumentKeySet &)keys
+               forTargetID:(FSTTargetID)targetID;
 
 /** Removes all the keys in the query results of the given target ID. */
 - (void)removeMatchingKeysForTargetID:(FSTTargetID)targetID;
 
-- (FSTDocumentKeySet *)matchingKeysForTargetID:(FSTTargetID)targetID;
+- (firebase::firestore::model::DocumentKeySet)matchingKeysForTargetID:(FSTTargetID)targetID;
 
 @end
 
diff --git a/Firestore/Source/Local/FSTReferenceSet.h b/Firestore/Source/Local/FSTReferenceSet.h
index 9d842cb..9a90a40 100644
--- a/Firestore/Source/Local/FSTReferenceSet.h
+++ b/Firestore/Source/Local/FSTReferenceSet.h
@@ -18,7 +18,8 @@
 
 #import "Firestore/Source/Core/FSTTypes.h"
 #import "Firestore/Source/Local/FSTGarbageCollector.h"
-#import "Firestore/Source/Model/FSTDocumentKeySet.h"
+
+#include "Firestore/core/src/firebase/firestore/model/document_key_set.h"
 
 NS_ASSUME_NONNULL_BEGIN
 
@@ -47,13 +48,14 @@ NS_ASSUME_NONNULL_BEGIN
 - (void)addReferenceToKey:(const firebase::firestore::model::DocumentKey &)key forID:(int)ID;
 
 /** Add references to the given document keys for the given ID. */
-- (void)addReferencesToKeys:(FSTDocumentKeySet *)keys forID:(int)ID;
+- (void)addReferencesToKeys:(const firebase::firestore::model::DocumentKeySet &)keys forID:(int)ID;
 
 /** Removes a reference to the given document key for the given ID. */
 - (void)removeReferenceToKey:(const firebase::firestore::model::DocumentKey &)key forID:(int)ID;
 
 /** Removes references to the given document keys for the given ID. */
-- (void)removeReferencesToKeys:(FSTDocumentKeySet *)keys forID:(int)ID;
+- (void)removeReferencesToKeys:(const firebase::firestore::model::DocumentKeySet &)keys
+                         forID:(int)ID;
 
 /** Clears all references with a given ID. Calls -removeReferenceToKey: for each key removed. */
 - (void)removeReferencesForID:(int)ID;
@@ -62,7 +64,7 @@ NS_ASSUME_NONNULL_BEGIN
 - (void)removeAllReferences;
 
 /** Returns all of the document keys that have had references added for the given ID. */
-- (FSTDocumentKeySet *)referencedKeysForID:(int)ID;
+- (firebase::firestore::model::DocumentKeySet)referencedKeysForID:(int)ID;
 
 @end
 
diff --git a/Firestore/Source/Local/FSTReferenceSet.mm b/Firestore/Source/Local/FSTReferenceSet.mm
index 14f5d47..6b34725 100644
--- a/Firestore/Source/Local/FSTReferenceSet.mm
+++ b/Firestore/Source/Local/FSTReferenceSet.mm
@@ -17,10 +17,12 @@
 #import "Firestore/Source/Local/FSTReferenceSet.h"
 
 #import "Firestore/Source/Local/FSTDocumentReference.h"
+#import "Firestore/third_party/Immutable/FSTImmutableSortedSet.h"
 
 #include "Firestore/core/src/firebase/firestore/model/document_key.h"
 
 using firebase::firestore::model::DocumentKey;
+using firebase::firestore::model::DocumentKeySet;
 
 NS_ASSUME_NONNULL_BEGIN
 
@@ -68,20 +70,20 @@ NS_ASSUME_NONNULL_BEGIN
   self.referencesByID = [self.referencesByID setByAddingObject:reference];
 }
 
-- (void)addReferencesToKeys:(FSTDocumentKeySet *)keys forID:(int)ID {
-  [keys enumerateObjectsUsingBlock:^(FSTDocumentKey *key, BOOL *stop) {
+- (void)addReferencesToKeys:(const DocumentKeySet &)keys forID:(int)ID {
+  for (const DocumentKey &key : keys) {
     [self addReferenceToKey:key forID:ID];
-  }];
+  }
 }
 
 - (void)removeReferenceToKey:(const DocumentKey &)key forID:(int)ID {
   [self removeReference:[[FSTDocumentReference alloc] initWithKey:key ID:ID]];
 }
 
-- (void)removeReferencesToKeys:(FSTDocumentKeySet *)keys forID:(int)ID {
-  [keys enumerateObjectsUsingBlock:^(FSTDocumentKey *key, BOOL *stop) {
+- (void)removeReferencesToKeys:(const DocumentKeySet &)keys forID:(int)ID {
+  for (const DocumentKey &key : keys) {
     [self removeReferenceToKey:key forID:ID];
-  }];
+  }
 }
 
 - (void)removeReferencesForID:(int)ID {
@@ -109,17 +111,17 @@ NS_ASSUME_NONNULL_BEGIN
   [self.garbageCollector addPotentialGarbageKey:reference.key];
 }
 
-- (FSTDocumentKeySet *)referencedKeysForID:(int)ID {
+- (DocumentKeySet)referencedKeysForID:(int)ID {
   FSTDocumentReference *start =
       [[FSTDocumentReference alloc] initWithKey:DocumentKey::Empty() ID:ID];
   FSTDocumentReference *end =
       [[FSTDocumentReference alloc] initWithKey:DocumentKey::Empty() ID:(ID + 1)];
 
-  __block FSTDocumentKeySet *keys = [FSTDocumentKeySet keySet];
+  __block DocumentKeySet keys;
   [self.referencesByID enumerateObjectsFrom:start
                                          to:end
                                  usingBlock:^(FSTDocumentReference *reference, BOOL *stop) {
-                                   keys = [keys setByAddingObject:reference.key];
+                                   keys = keys.insert(reference.key);
                                  }];
   return keys;
 }
diff --git a/Firestore/Source/Model/FSTDocumentKeySet.h b/Firestore/Source/Model/FSTDocumentKeySet.h
deleted file mode 100644
index 80f6624..0000000
--- a/Firestore/Source/Model/FSTDocumentKeySet.h
+++ /dev/null
@@ -1,35 +0,0 @@
-/*
- * Copyright 2017 Google
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-#import <Foundation/Foundation.h>
-
-#import "Firestore/third_party/Immutable/FSTImmutableSortedSet.h"
-
-@class FSTDocumentKey;
-
-NS_ASSUME_NONNULL_BEGIN
-
-/** Convenience type for a set of keys, since they are so common. */
-typedef FSTImmutableSortedSet<FSTDocumentKey *> FSTDocumentKeySet;
-
-@interface FSTImmutableSortedSet (FSTDocumentKey)
-
-/** Returns a new set using the DocumentKeyComparator. */
-+ (FSTDocumentKeySet *)keySet;
-
-@end
-
-NS_ASSUME_NONNULL_END
diff --git a/Firestore/Source/Model/FSTDocumentKeySet.mm b/Firestore/Source/Model/FSTDocumentKeySet.mm
deleted file mode 100644
index f07b785..0000000
--- a/Firestore/Source/Model/FSTDocumentKeySet.mm
+++ /dev/null
@@ -1,31 +0,0 @@
-/*
- * Copyright 2017 Google
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-#import "Firestore/Source/Model/FSTDocumentKeySet.h"
-
-#import "Firestore/Source/Model/FSTDocumentKey.h"
-
-NS_ASSUME_NONNULL_BEGIN
-
-@implementation FSTImmutableSortedSet (FSTDocumentKey)
-
-+ (instancetype)keySet {
-  return [FSTDocumentKeySet setWithComparator:FSTDocumentKeyComparator];
-}
-
-@end
-
-NS_ASSUME_NONNULL_END
diff --git a/Firestore/Source/Model/FSTMutationBatch.h b/Firestore/Source/Model/FSTMutationBatch.h
index d3d68a9..f95b145 100644
--- a/Firestore/Source/Model/FSTMutationBatch.h
+++ b/Firestore/Source/Model/FSTMutationBatch.h
@@ -17,10 +17,10 @@
 #import <Foundation/Foundation.h>
 
 #import "Firestore/Source/Core/FSTTypes.h"
-#import "Firestore/Source/Model/FSTDocumentKeySet.h"
 #import "Firestore/Source/Model/FSTDocumentVersionDictionary.h"
 
 #include "Firestore/core/src/firebase/firestore/model/document_key.h"
+#include "Firestore/core/src/firebase/firestore/model/document_key_set.h"
 #include "Firestore/core/src/firebase/firestore/model/snapshot_version.h"
 
 @class FSTMutation;
@@ -85,7 +85,7 @@ extern const FSTBatchID kFSTBatchIDUnknown;
 - (FSTMutationBatch *)toTombstone;
 
 /** Returns the set of unique keys referenced by all mutations in the batch. */
-- (FSTDocumentKeySet *)keys;
+- (firebase::firestore::model::DocumentKeySet)keys;
 
 @property(nonatomic, assign, readonly) FSTBatchID batchID;
 @property(nonatomic, strong, readonly) FIRTimestamp *localWriteTime;
diff --git a/Firestore/Source/Model/FSTMutationBatch.mm b/Firestore/Source/Model/FSTMutationBatch.mm
index 3c03bdc..9926bd4 100644
--- a/Firestore/Source/Model/FSTMutationBatch.mm
+++ b/Firestore/Source/Model/FSTMutationBatch.mm
@@ -28,6 +28,7 @@
 
 using firebase::firestore::model::DocumentKey;
 using firebase::firestore::model::SnapshotVersion;
+using firebase::firestore::model::DocumentKeySet;
 
 NS_ASSUME_NONNULL_BEGIN
 
@@ -115,10 +116,10 @@ const FSTBatchID kFSTBatchIDUnknown = -1;
 }
 
 // TODO(klimt): This could use NSMutableDictionary instead.
-- (FSTDocumentKeySet *)keys {
-  FSTDocumentKeySet *set = [FSTDocumentKeySet keySet];
+- (DocumentKeySet)keys {
+  DocumentKeySet set;
   for (FSTMutation *mutation in self.mutations) {
-    set = [set setByAddingObject:mutation.key];
+    set = set.insert(mutation.key);
   }
   return set;
 }
diff --git a/Firestore/Source/Remote/FSTRemoteEvent.h b/Firestore/Source/Remote/FSTRemoteEvent.h
index 06099c7..c84e34d 100644
--- a/Firestore/Source/Remote/FSTRemoteEvent.h
+++ b/Firestore/Source/Remote/FSTRemoteEvent.h
@@ -20,9 +20,9 @@
 
 #import "Firestore/Source/Core/FSTTypes.h"
 #import "Firestore/Source/Model/FSTDocumentDictionary.h"
-#import "Firestore/Source/Model/FSTDocumentKeySet.h"
 
 #include "Firestore/core/src/firebase/firestore/model/document_key.h"
+#include "Firestore/core/src/firebase/firestore/model/document_key_set.h"
 #include "Firestore/core/src/firebase/firestore/model/snapshot_version.h"
 
 @class FSTDocument;
@@ -49,7 +49,8 @@ NS_ASSUME_NONNULL_BEGIN
  * 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;
+- (void)filterUpdatesUsingExistingKeys:
+    (const firebase::firestore::model::DocumentKeySet &)existingKeys;
 
 @end
 
@@ -65,7 +66,7 @@ NS_ASSUME_NONNULL_BEGIN
 + (FSTResetMapping *)mappingWithDocuments:(NSArray<FSTDocument *> *)documents;
 
 /** The new set of documents for the target. */
-@property(nonatomic, strong, readonly) FSTDocumentKeySet *documents;
+- (const firebase::firestore::model::DocumentKeySet &)documents;
 @end
 
 #pragma mark - FSTUpdateMapping
@@ -82,12 +83,13 @@ NS_ASSUME_NONNULL_BEGIN
 + (FSTUpdateMapping *)mappingWithAddedDocuments:(NSArray<FSTDocument *> *)added
                                removedDocuments:(NSArray<FSTDocument *> *)removed;
 
-- (FSTDocumentKeySet *)applyTo:(FSTDocumentKeySet *)keys;
+- (firebase::firestore::model::DocumentKeySet)applyTo:
+    (const firebase::firestore::model::DocumentKeySet &)keys;
 
 /** The documents added to the target. */
-@property(nonatomic, strong, readonly) FSTDocumentKeySet *addedDocuments;
+- (const firebase::firestore::model::DocumentKeySet &)addedDocuments;
 /** The documents removed from the target. */
-@property(nonatomic, strong, readonly) FSTDocumentKeySet *removedDocuments;
+- (const firebase::firestore::model::DocumentKeySet &)removedDocuments;
 @end
 
 #pragma mark - FSTTargetChange
@@ -166,7 +168,7 @@ initWithSnapshotVersion:(firebase::firestore::model::SnapshotVersion)snapshotVer
           targetChanges:(NSMutableDictionary<NSNumber *, FSTTargetChange *> *)targetChanges
         documentUpdates:
             (std::map<firebase::firestore::model::DocumentKey, FSTMaybeDocument *>)documentUpdates
-         limboDocuments:(FSTDocumentKeySet *)limboDocuments;
+         limboDocuments:(firebase::firestore::model::DocumentKeySet)limboDocuments;
 
 /** The snapshot version this event brings us up to. */
 - (const firebase::firestore::model::SnapshotVersion &)snapshotVersion;
@@ -175,14 +177,14 @@ initWithSnapshotVersion:(firebase::firestore::model::SnapshotVersion)snapshotVer
 @property(nonatomic, strong, readonly)
     NSDictionary<FSTBoxedTargetID *, FSTTargetChange *> *targetChanges;
 
-@property(nonatomic, strong, readonly) FSTDocumentKeySet *limboDocumentChanges;
-
 /**
  * A set of which documents have changed or been deleted, along with the doc's new values
  * (if not deleted).
  */
 - (const std::map<firebase::firestore::model::DocumentKey, FSTMaybeDocument *> &)documentUpdates;
 
+- (const firebase::firestore::model::DocumentKeySet &)limboDocumentChanges;
+
 /** Adds a document update to this remote event */
 - (void)addDocumentUpdate:(FSTMaybeDocument *)document;
 
diff --git a/Firestore/Source/Remote/FSTRemoteEvent.mm b/Firestore/Source/Remote/FSTRemoteEvent.mm
index 718fd11..438072e 100644
--- a/Firestore/Source/Remote/FSTRemoteEvent.mm
+++ b/Firestore/Source/Remote/FSTRemoteEvent.mm
@@ -26,9 +26,12 @@
 #import "Firestore/Source/Util/FSTClasses.h"
 #import "Firestore/Source/Util/FSTLogger.h"
 #include "Firestore/core/src/firebase/firestore/model/document_key.h"
+#include "Firestore/core/src/firebase/firestore/util/hashing.h"
 
 using firebase::firestore::model::DocumentKey;
 using firebase::firestore::model::SnapshotVersion;
+using firebase::firestore::util::Hash;
+using firebase::firestore::model::DocumentKeySet;
 
 NS_ASSUME_NONNULL_BEGIN
 
@@ -54,7 +57,7 @@ NS_ASSUME_NONNULL_BEGIN
   @throw FSTAbstractMethodException();  // NOLINT
 }
 
-- (void)filterUpdatesUsingExistingKeys:(FSTDocumentKeySet *)existingKeys {
+- (void)filterUpdatesUsingExistingKeys:(const DocumentKeySet &)existingKeys {
   @throw FSTAbstractMethodException();  // NOLINT
 }
 
@@ -62,28 +65,30 @@ NS_ASSUME_NONNULL_BEGIN
 
 #pragma mark - FSTResetMapping
 
-@interface FSTResetMapping ()
-@property(nonatomic, strong) FSTDocumentKeySet *documents;
-@end
-
-@implementation FSTResetMapping
+@implementation FSTResetMapping {
+  DocumentKeySet _documents;
+}
 
 + (instancetype)mappingWithDocuments:(NSArray<FSTDocument *> *)documents {
-  FSTResetMapping *mapping = [[FSTResetMapping alloc] init];
+  DocumentKeySet keys;
   for (FSTDocument *doc in documents) {
-    mapping.documents = [mapping.documents setByAddingObject:doc.key];
+    keys = keys.insert(doc.key);
   }
-  return mapping;
+  return [[FSTResetMapping alloc] initWithDocuments:std::move(keys)];
 }
 
-- (instancetype)init {
+- (instancetype)initWithDocuments:(DocumentKeySet)documents {
   self = [super init];
   if (self) {
-    _documents = [FSTDocumentKeySet keySet];
+    _documents = std::move(documents);
   }
   return self;
 }
 
+- (const DocumentKeySet &)documents {
+  return _documents;
+}
+
 - (BOOL)isEqual:(id)other {
   if (other == self) {
     return YES;
@@ -93,22 +98,22 @@ NS_ASSUME_NONNULL_BEGIN
   }
 
   FSTResetMapping *otherMapping = (FSTResetMapping *)other;
-  return [self.documents isEqual:otherMapping.documents];
+  return _documents == otherMapping.documents;
 }
 
 - (NSUInteger)hash {
-  return self.documents.hash;
+  return Hash(_documents);
 }
 
 - (void)addDocumentKey:(const DocumentKey &)documentKey {
-  self.documents = [self.documents setByAddingObject:documentKey];
+  _documents = _documents.insert(documentKey);
 }
 
 - (void)removeDocumentKey:(const DocumentKey &)documentKey {
-  self.documents = [self.documents setByRemovingObject:documentKey];
+  _documents = _documents.erase(documentKey);
 }
 
-- (void)filterUpdatesUsingExistingKeys:(FSTDocumentKeySet *)existingKeys {
+- (void)filterUpdatesUsingExistingKeys:(const DocumentKeySet &)existingKeys {
   // No-op. Resets are not filtered.
 }
 
@@ -116,34 +121,43 @@ NS_ASSUME_NONNULL_BEGIN
 
 #pragma mark - FSTUpdateMapping
 
-@interface FSTUpdateMapping ()
-@property(nonatomic, strong) FSTDocumentKeySet *addedDocuments;
-@property(nonatomic, strong) FSTDocumentKeySet *removedDocuments;
-@end
-
-@implementation FSTUpdateMapping
+@implementation FSTUpdateMapping {
+  DocumentKeySet _addedDocuments;
+  DocumentKeySet _removedDocuments;
+}
 
 + (FSTUpdateMapping *)mappingWithAddedDocuments:(NSArray<FSTDocument *> *)added
                                removedDocuments:(NSArray<FSTDocument *> *)removed {
-  FSTUpdateMapping *mapping = [[FSTUpdateMapping alloc] init];
+  DocumentKeySet addedDocuments;
+  DocumentKeySet removedDocuments;
   for (FSTDocument *doc in added) {
-    mapping.addedDocuments = [mapping.addedDocuments setByAddingObject:doc.key];
+    addedDocuments = addedDocuments.insert(doc.key);
   }
   for (FSTDocument *doc in removed) {
-    mapping.removedDocuments = [mapping.removedDocuments setByAddingObject:doc.key];
+    removedDocuments = removedDocuments.insert(doc.key);
   }
-  return mapping;
+  return [[FSTUpdateMapping alloc] initWithAddedDocuments:std::move(addedDocuments)
+                                         removedDocuments:std::move(removedDocuments)];
 }
 
-- (instancetype)init {
+- (instancetype)initWithAddedDocuments:(DocumentKeySet)addedDocuments
+                      removedDocuments:(DocumentKeySet)removedDocuments {
   self = [super init];
   if (self) {
-    _addedDocuments = [FSTDocumentKeySet keySet];
-    _removedDocuments = [FSTDocumentKeySet keySet];
+    _addedDocuments = std::move(addedDocuments);
+    _removedDocuments = std::move(removedDocuments);
   }
   return self;
 }
 
+- (const DocumentKeySet &)addedDocuments {
+  return _addedDocuments;
+}
+
+- (const DocumentKeySet &)removedDocuments {
+  return _removedDocuments;
+}
+
 - (BOOL)isEqual:(id)other {
   if (other == self) {
     return YES;
@@ -153,42 +167,42 @@ NS_ASSUME_NONNULL_BEGIN
   }
 
   FSTUpdateMapping *otherMapping = (FSTUpdateMapping *)other;
-  return [self.addedDocuments isEqual:otherMapping.addedDocuments] &&
-         [self.removedDocuments isEqual:otherMapping.removedDocuments];
+  return _addedDocuments == otherMapping.addedDocuments &&
+         _removedDocuments == otherMapping.removedDocuments;
 }
 
 - (NSUInteger)hash {
-  return self.addedDocuments.hash * 31 + self.removedDocuments.hash;
+  return Hash(_addedDocuments, _removedDocuments);
 }
 
-- (FSTDocumentKeySet *)applyTo:(FSTDocumentKeySet *)keys {
-  __block FSTDocumentKeySet *result = keys;
-  [self.addedDocuments enumerateObjectsUsingBlock:^(FSTDocumentKey *key, BOOL *stop) {
-    result = [result setByAddingObject:key];
-  }];
-  [self.removedDocuments enumerateObjectsUsingBlock:^(FSTDocumentKey *key, BOOL *stop) {
-    result = [result setByRemovingObject:key];
-  }];
+- (DocumentKeySet)applyTo:(const DocumentKeySet &)keys {
+  DocumentKeySet result = keys;
+  for (const DocumentKey &key : _addedDocuments) {
+    result = result.insert(key);
+  }
+  for (const DocumentKey &key : _removedDocuments) {
+    result = result.erase(key);
+  }
   return result;
 }
 
 - (void)addDocumentKey:(const DocumentKey &)documentKey {
-  self.addedDocuments = [self.addedDocuments setByAddingObject:documentKey];
-  self.removedDocuments = [self.removedDocuments setByRemovingObject:documentKey];
+  _addedDocuments = _addedDocuments.insert(documentKey);
+  _removedDocuments = _removedDocuments.erase(documentKey);
 }
 
 - (void)removeDocumentKey:(const DocumentKey &)documentKey {
-  self.addedDocuments = [self.addedDocuments setByRemovingObject:documentKey];
-  self.removedDocuments = [self.removedDocuments setByAddingObject:documentKey];
+  _addedDocuments = _addedDocuments.erase(documentKey);
+  _removedDocuments = _removedDocuments.insert(documentKey);
 }
 
-- (void)filterUpdatesUsingExistingKeys:(FSTDocumentKeySet *)existingKeys {
-  __block FSTDocumentKeySet *result = _addedDocuments;
-  [_addedDocuments enumerateObjectsUsingBlock:^(FSTDocumentKey *docKey, BOOL *stop) {
-    if ([existingKeys containsObject:docKey]) {
-      result = [result setByRemovingObject:docKey];
+- (void)filterUpdatesUsingExistingKeys:(const DocumentKeySet &)existingKeys {
+  DocumentKeySet result = _addedDocuments;
+  for (const DocumentKey &key : _addedDocuments) {
+    if (existingKeys.contains(key)) {
+      result = result.erase(key);
     }
-  }];
+  }
   _addedDocuments = result;
 }
 
@@ -227,14 +241,19 @@ NS_ASSUME_NONNULL_BEGIN
 
 + (instancetype)changeWithDocuments:(NSArray<FSTMaybeDocument *> *)docs
                 currentStatusUpdate:(FSTCurrentStatusUpdate)currentStatusUpdate {
-  FSTUpdateMapping *mapping = [[FSTUpdateMapping alloc] init];
+  DocumentKeySet addedDocuments;
+  DocumentKeySet removedDocuments;
   for (FSTMaybeDocument *doc in docs) {
     if ([doc isKindOfClass:[FSTDeletedDocument class]]) {
-      mapping.removedDocuments = [mapping.removedDocuments setByAddingObject:doc.key];
+      removedDocuments = removedDocuments.insert(doc.key);
     } else {
-      mapping.addedDocuments = [mapping.addedDocuments setByAddingObject:doc.key];
+      addedDocuments = addedDocuments.insert(doc.key);
     }
   }
+  FSTUpdateMapping *mapping =
+      [[FSTUpdateMapping alloc] initWithAddedDocuments:std::move(addedDocuments)
+                                      removedDocuments:std::move(removedDocuments)];
+
   FSTTargetChange *change = [[FSTTargetChange alloc] init];
   change.mapping = mapping;
   change.currentStatusUpdate = currentStatusUpdate;
@@ -273,38 +292,44 @@ NS_ASSUME_NONNULL_BEGIN
 
 #pragma mark - FSTRemoteEvent
 
-@interface FSTRemoteEvent () {
-  NSMutableDictionary<FSTBoxedTargetID *, FSTTargetChange *> *_targetChanges;
-}
-
-- (instancetype)
-initWithSnapshotVersion:(SnapshotVersion)snapshotVersion
-          targetChanges:(NSMutableDictionary<FSTBoxedTargetID *, FSTTargetChange *> *)targetChanges
-        documentUpdates:(std::map<DocumentKey, FSTMaybeDocument *>)documentUpdates
-         limboDocuments:(FSTDocumentKeySet *)limboDocuments;
-
-@end
-
 @implementation FSTRemoteEvent {
-  std::map<DocumentKey, FSTMaybeDocument *> _documentUpdates;
   SnapshotVersion _snapshotVersion;
+  NSMutableDictionary<FSTBoxedTargetID *, FSTTargetChange *> *_targetChanges;
+  std::map<DocumentKey, FSTMaybeDocument *> _documentUpdates;
+  DocumentKeySet _limboDocumentChanges;
 }
 
 - (instancetype)initWithSnapshotVersion:(SnapshotVersion)snapshotVersion
                           targetChanges:
                               (NSMutableDictionary<NSNumber *, FSTTargetChange *> *)targetChanges
                         documentUpdates:(std::map<DocumentKey, FSTMaybeDocument *>)documentUpdates
-                         limboDocuments:(FSTDocumentKeySet *)limboDocuments {
+                         limboDocuments:(DocumentKeySet)limboDocuments {
   self = [super init];
   if (self) {
     _snapshotVersion = std::move(snapshotVersion);
     _targetChanges = targetChanges;
     _documentUpdates = std::move(documentUpdates);
-    _limboDocumentChanges = limboDocuments;
+    _limboDocumentChanges = std::move(limboDocuments);
   }
   return self;
 }
 
+- (NSDictionary<FSTBoxedTargetID *, FSTTargetChange *> *)targetChanges {
+  return _targetChanges;
+}
+
+- (const DocumentKeySet &)limboDocumentChanges {
+  return _limboDocumentChanges;
+}
+
+- (const std::map<DocumentKey, FSTMaybeDocument *> &)documentUpdates {
+  return _documentUpdates;
+}
+
+- (const SnapshotVersion &)snapshotVersion {
+  return _snapshotVersion;
+}
+
 - (void)synthesizeDeleteForLimboTargetChange:(FSTTargetChange *)targetChange
                                          key:(const DocumentKey &)key {
   if (targetChange.currentStatusUpdate == FSTCurrentStatusUpdateMarkCurrent &&
@@ -328,22 +353,10 @@ initWithSnapshotVersion:(SnapshotVersion)snapshotVersion
     // 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];
-    _limboDocumentChanges = [_limboDocumentChanges setByAddingObject:key];
+    _limboDocumentChanges = _limboDocumentChanges.insert(key);
   }
 }
 
-- (NSDictionary<FSTBoxedTargetID *, FSTTargetChange *> *)targetChanges {
-  return _targetChanges;
-}
-
-- (const std::map<DocumentKey, FSTMaybeDocument *> &)documentUpdates {
-  return _documentUpdates;
-}
-
-- (const SnapshotVersion &)snapshotVersion {
-  return _snapshotVersion;
-}
-
 /** Adds a document update to this remote event */
 - (void)addDocumentUpdate:(FSTMaybeDocument *)document {
   _documentUpdates[document.key] = document;
@@ -393,7 +406,7 @@ initWithSnapshotVersion:(SnapshotVersion)snapshotVersion
   /** Keeps track of document to update */
   std::map<DocumentKey, FSTMaybeDocument *> _documentUpdates;
 
-  FSTDocumentKeySet *_limboDocuments;
+  DocumentKeySet _limboDocuments;
   /** The snapshot version for every target change this creates. */
   SnapshotVersion _snapshotVersion;
 }
@@ -410,7 +423,7 @@ initWithSnapshotVersion:(SnapshotVersion)snapshotVersion
     _targetChanges = [NSMutableDictionary dictionary];
     _listenTargets = listenTargets;
     _pendingTargetResponses = [NSMutableDictionary dictionaryWithDictionary:pendingTargetResponses];
-    _limboDocuments = [FSTDocumentKeySet keySet];
+    _limboDocuments = DocumentKeySet{};
     _existenceFilters = [NSMutableDictionary dictionary];
   }
   return self;
@@ -467,7 +480,7 @@ initWithSnapshotVersion:(SnapshotVersion)snapshotVersion
     // We haven't seen a document update for this key yet.
     if (queryData.purpose == FSTQueryPurposeLimboResolution) {
       // We haven't seen this document before, and this target is a limbo target.
-      _limboDocuments = [_limboDocuments setByAddingObject:documentKey];
+      _limboDocuments = _limboDocuments.insert(documentKey);
       return YES;
     } else {
       // We haven't seen the document before, but this is a non-limbo target.
@@ -481,7 +494,7 @@ initWithSnapshotVersion:(SnapshotVersion)snapshotVersion
   } else {
     // We haven't marked this as non-limbo yet, but this target is not a limbo target.
     // Mark the key as non-limbo and make sure it isn't in our set.
-    _limboDocuments = [_limboDocuments setByRemovingObject:documentKey];
+    _limboDocuments = _limboDocuments.erase(documentKey);
     return NO;
   }
 }
diff --git a/Firestore/Source/Remote/FSTRemoteStore.mm b/Firestore/Source/Remote/FSTRemoteStore.mm
index d15d2e5..0ea4887 100644
--- a/Firestore/Source/Remote/FSTRemoteStore.mm
+++ b/Firestore/Source/Remote/FSTRemoteStore.mm
@@ -43,6 +43,7 @@ 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
 
@@ -393,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]]) {
@@ -406,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