Port FSTDocumentKeySet to C++ DocumentKeySet (#1229)

* Define a Comparator for DocumentKey

* Automated migration from FSTDocumentKeySet to DocumentKeySet

* Manual fixups for DocumentKeySet

* Delete FSTDocumentKeySet

12 files changed, 118 insertions, 93 deletions

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;
 }