Port FSTDocumentKeySet to C++ DocumentKeySet (#1229)

* Define a Comparator for DocumentKey * Automated migration from FSTDocumentKeySet to DocumentKeySet * Manual fixups for DocumentKeySet * Delete FSTDocumentKeySet
author: Gil <mcg@google.com> 2018-05-05 08:10:51 -0700
committer: GitHub <noreply@github.com> 2018-05-05 08:10:51 -0700
commit: a080e481b5e6fcbc2b645920051cf20fc8cad7a7 (patch)
tree: 5d5ce3d39be316ac82a2fd34c92b05fc9e67f83a /Firestore
parent: d158e420c6fa04997ee3d2a6c44fd53a52883d81 (diff)
36 files changed, 446 insertions, 429 deletions
diff --git a/Firestore/Example/Tests/Core/FSTQueryListenerTests.mm b/Firestore/Example/Tests/Core/FSTQueryListenerTests.mm
index 0454152..7ae9704 100644
--- a/Firestore/Example/Tests/Core/FSTQueryListenerTests.mm
+++ b/Firestore/Example/Tests/Core/FSTQueryListenerTests.mm
@@ -27,6 +27,8 @@
 
 #import "Firestore/Example/Tests/Util/FSTHelpers.h"
 
+using firebase::firestore::model::DocumentKeySet;
+
 NS_ASSUME_NONNULL_BEGIN
 
 @interface FSTQueryListenerTests : XCTestCase
@@ -53,7 +55,7 @@ NS_ASSUME_NONNULL_BEGIN
   FSTQueryListener *listener = [self listenToQuery:query accumulatingSnapshots:accum];
   FSTQueryListener *otherListener = [self listenToQuery:query accumulatingSnapshots:otherAccum];
 
-  FSTView *view = [[FSTView alloc] initWithQuery:query remoteDocuments:[FSTDocumentKeySet keySet]];
+  FSTView *view = [[FSTView alloc] initWithQuery:query remoteDocuments:DocumentKeySet{}];
   FSTViewSnapshot *snap1 = FSTTestApplyChanges(view, @[ doc1, doc2 ], nil);
   FSTViewSnapshot *snap2 = FSTTestApplyChanges(view, @[ doc2prime ], nil);
 
@@ -107,7 +109,7 @@ NS_ASSUME_NONNULL_BEGIN
 
   FSTQueryListener *listener = [self listenToQuery:query accumulatingSnapshots:accum];
 
-  FSTView *view = [[FSTView alloc] initWithQuery:query remoteDocuments:[FSTDocumentKeySet keySet]];
+  FSTView *view = [[FSTView alloc] initWithQuery:query remoteDocuments:DocumentKeySet{}];
   FSTViewSnapshot *snap1 = FSTTestApplyChanges(view, @[], nil);
 
   FSTTargetChange *ackTarget =
@@ -136,7 +138,7 @@ NS_ASSUME_NONNULL_BEGIN
               [listener mute];
             }];
 
-  FSTView *view = [[FSTView alloc] initWithQuery:query remoteDocuments:[FSTDocumentKeySet keySet]];
+  FSTView *view = [[FSTView alloc] initWithQuery:query remoteDocuments:DocumentKeySet{}];
   FSTViewSnapshot *viewSnapshot1 = FSTTestApplyChanges(view, @[ doc1 ], nil);
   FSTViewSnapshot *viewSnapshot2 = FSTTestApplyChanges(view, @[ doc2 ], nil);
 
@@ -178,7 +180,7 @@ NS_ASSUME_NONNULL_BEGIN
   FSTQueryListener *fullListener =
       [self listenToQuery:query options:options accumulatingSnapshots:fullAccum];
 
-  FSTView *view = [[FSTView alloc] initWithQuery:query remoteDocuments:[FSTDocumentKeySet keySet]];
+  FSTView *view = [[FSTView alloc] initWithQuery:query remoteDocuments:DocumentKeySet{}];
   FSTViewSnapshot *snap1 = FSTTestApplyChanges(view, @[ doc1 ], nil);
 
   FSTTargetChange *ackTarget =
@@ -218,7 +220,7 @@ NS_ASSUME_NONNULL_BEGIN
   FSTQueryListener *fullListener =
       [self listenToQuery:query options:options accumulatingSnapshots:fullAccum];
 
-  FSTView *view = [[FSTView alloc] initWithQuery:query remoteDocuments:[FSTDocumentKeySet keySet]];
+  FSTView *view = [[FSTView alloc] initWithQuery:query remoteDocuments:DocumentKeySet{}];
   FSTViewSnapshot *snap1 = FSTTestApplyChanges(view, @[ doc1, doc2 ], nil);
   FSTViewSnapshot *snap2 = FSTTestApplyChanges(view, @[ doc1Prime ], nil);
   FSTViewSnapshot *snap3 = FSTTestApplyChanges(view, @[ doc3 ], nil);
@@ -265,7 +267,7 @@ NS_ASSUME_NONNULL_BEGIN
   FSTQueryListener *fullListener =
       [self listenToQuery:query options:options accumulatingSnapshots:fullAccum];
 
-  FSTView *view = [[FSTView alloc] initWithQuery:query remoteDocuments:[FSTDocumentKeySet keySet]];
+  FSTView *view = [[FSTView alloc] initWithQuery:query remoteDocuments:DocumentKeySet{}];
   FSTViewSnapshot *snap1 = FSTTestApplyChanges(view, @[ doc1, doc2 ], nil);
   FSTViewSnapshot *snap2 = FSTTestApplyChanges(view, @[ doc1Prime ], nil);
   FSTViewSnapshot *snap3 = FSTTestApplyChanges(view, @[ doc3 ], nil);
@@ -298,7 +300,7 @@ NS_ASSUME_NONNULL_BEGIN
   FSTQueryListener *filteredListener =
       [self listenToQuery:query accumulatingSnapshots:filteredAccum];
 
-  FSTView *view = [[FSTView alloc] initWithQuery:query remoteDocuments:[FSTDocumentKeySet keySet]];
+  FSTView *view = [[FSTView alloc] initWithQuery:query remoteDocuments:DocumentKeySet{}];
   FSTViewSnapshot *snap1 = FSTTestApplyChanges(view, @[ doc1, doc2 ], nil);
   FSTViewSnapshot *snap2 = FSTTestApplyChanges(view, @[ doc1Prime, doc3 ], nil);
 
@@ -331,7 +333,7 @@ NS_ASSUME_NONNULL_BEGIN
                                                                         waitForSyncWhenOnline:YES]
           accumulatingSnapshots:events];
 
-  FSTView *view = [[FSTView alloc] initWithQuery:query remoteDocuments:[FSTDocumentKeySet keySet]];
+  FSTView *view = [[FSTView alloc] initWithQuery:query remoteDocuments:DocumentKeySet{}];
   FSTViewSnapshot *snap1 = FSTTestApplyChanges(view, @[ doc1 ], nil);
   FSTViewSnapshot *snap2 = FSTTestApplyChanges(view, @[ doc2 ], nil);
   FSTViewSnapshot *snap3 =
@@ -374,7 +376,7 @@ NS_ASSUME_NONNULL_BEGIN
                                                                         waitForSyncWhenOnline:YES]
           accumulatingSnapshots:events];
 
-  FSTView *view = [[FSTView alloc] initWithQuery:query remoteDocuments:[FSTDocumentKeySet keySet]];
+  FSTView *view = [[FSTView alloc] initWithQuery:query remoteDocuments:DocumentKeySet{}];
   FSTViewSnapshot *snap1 = FSTTestApplyChanges(view, @[ doc1 ], nil);
   FSTViewSnapshot *snap2 = FSTTestApplyChanges(view, @[ doc2 ], nil);
 
@@ -415,7 +417,7 @@ NS_ASSUME_NONNULL_BEGIN
                                            options:[FSTListenOptions defaultOptions]
                              accumulatingSnapshots:events];
 
-  FSTView *view = [[FSTView alloc] initWithQuery:query remoteDocuments:[FSTDocumentKeySet keySet]];
+  FSTView *view = [[FSTView alloc] initWithQuery:query remoteDocuments:DocumentKeySet{}];
   FSTViewSnapshot *snap1 = FSTTestApplyChanges(view, @[], nil);
 
   [listener applyChangedOnlineState:FSTOnlineStateOnline];   // no event
@@ -441,7 +443,7 @@ NS_ASSUME_NONNULL_BEGIN
                                            options:[FSTListenOptions defaultOptions]
                              accumulatingSnapshots:events];
 
-  FSTView *view = [[FSTView alloc] initWithQuery:query remoteDocuments:[FSTDocumentKeySet keySet]];
+  FSTView *view = [[FSTView alloc] initWithQuery:query remoteDocuments:DocumentKeySet{}];
   FSTViewSnapshot *snap1 = FSTTestApplyChanges(view, @[], nil);
 
   [listener applyChangedOnlineState:FSTOnlineStateOffline];  // no event
diff --git a/Firestore/Example/Tests/Core/FSTViewTests.mm b/Firestore/Example/Tests/Core/FSTViewTests.mm
index 63ce711..ec62d82 100644
--- a/Firestore/Example/Tests/Core/FSTViewTests.mm
+++ b/Firestore/Example/Tests/Core/FSTViewTests.mm
@@ -34,6 +34,7 @@
 
 namespace testutil = firebase::firestore::testutil;
 using firebase::firestore::model::ResourcePath;
+using firebase::firestore::model::DocumentKeySet;
 
 NS_ASSUME_NONNULL_BEGIN
 
@@ -49,7 +50,7 @@ NS_ASSUME_NONNULL_BEGIN
 
 - (void)testAddsDocumentsBasedOnQuery {
   FSTQuery *query = [self queryForMessages];
-  FSTView *view = [[FSTView alloc] initWithQuery:query remoteDocuments:[FSTDocumentKeySet keySet]];
+  FSTView *view = [[FSTView alloc] initWithQuery:query remoteDocuments:DocumentKeySet{}];
 
   FSTDocument *doc1 = FSTTestDoc("rooms/eros/messages/1", 0, @{@"text" : @"msg1"}, NO);
   FSTDocument *doc2 = FSTTestDoc("rooms/eros/messages/2", 0, @{@"text" : @"msg2"}, NO);
@@ -77,7 +78,7 @@ NS_ASSUME_NONNULL_BEGIN
 
 - (void)testRemovesDocuments {
   FSTQuery *query = [self queryForMessages];
-  FSTView *view = [[FSTView alloc] initWithQuery:query remoteDocuments:[FSTDocumentKeySet keySet]];
+  FSTView *view = [[FSTView alloc] initWithQuery:query remoteDocuments:DocumentKeySet{}];
 
   FSTDocument *doc1 = FSTTestDoc("rooms/eros/messages/1", 0, @{@"text" : @"msg1"}, NO);
   FSTDocument *doc2 = FSTTestDoc("rooms/eros/messages/2", 0, @{@"text" : @"msg2"}, NO);
@@ -108,7 +109,7 @@ NS_ASSUME_NONNULL_BEGIN
 
 - (void)testReturnsNilIfThereAreNoChanges {
   FSTQuery *query = [self queryForMessages];
-  FSTView *view = [[FSTView alloc] initWithQuery:query remoteDocuments:[FSTDocumentKeySet keySet]];
+  FSTView *view = [[FSTView alloc] initWithQuery:query remoteDocuments:DocumentKeySet{}];
 
   FSTDocument *doc1 = FSTTestDoc("rooms/eros/messages/1", 0, @{@"text" : @"msg1"}, NO);
   FSTDocument *doc2 = FSTTestDoc("rooms/eros/messages/2", 0, @{@"text" : @"msg2"}, NO);
@@ -123,7 +124,7 @@ NS_ASSUME_NONNULL_BEGIN
 
 - (void)testDoesNotReturnNilForFirstChanges {
   FSTQuery *query = [self queryForMessages];
-  FSTView *view = [[FSTView alloc] initWithQuery:query remoteDocuments:[FSTDocumentKeySet keySet]];
+  FSTView *view = [[FSTView alloc] initWithQuery:query remoteDocuments:DocumentKeySet{}];
 
   FSTViewSnapshot *snapshot = FSTTestApplyChanges(view, @[], nil);
   XCTAssertNotNil(snapshot);
@@ -137,7 +138,7 @@ NS_ASSUME_NONNULL_BEGIN
                                    value:[FSTDoubleValue doubleValue:2]];
   query = [query queryByAddingFilter:filter];
 
-  FSTView *view = [[FSTView alloc] initWithQuery:query remoteDocuments:[FSTDocumentKeySet keySet]];
+  FSTView *view = [[FSTView alloc] initWithQuery:query remoteDocuments:DocumentKeySet{}];
   FSTDocument *doc1 = FSTTestDoc("rooms/eros/messages/1", 0, @{ @"sort" : @1 }, NO);
   FSTDocument *doc2 = FSTTestDoc("rooms/eros/messages/2", 0, @{ @"sort" : @2 }, NO);
   FSTDocument *doc3 = FSTTestDoc("rooms/eros/messages/3", 0, @{ @"sort" : @3 }, NO);
@@ -169,7 +170,7 @@ NS_ASSUME_NONNULL_BEGIN
                                    value:[FSTDoubleValue doubleValue:2]];
   query = [query queryByAddingFilter:filter];
 
-  FSTView *view = [[FSTView alloc] initWithQuery:query remoteDocuments:[FSTDocumentKeySet keySet]];
+  FSTView *view = [[FSTView alloc] initWithQuery:query remoteDocuments:DocumentKeySet{}];
   FSTDocument *doc1 = FSTTestDoc("rooms/eros/messages/1", 0, @{ @"sort" : @1 }, NO);
   FSTDocument *doc2 = FSTTestDoc("rooms/eros/messages/2", 0, @{ @"sort" : @3 }, NO);
   FSTDocument *doc3 = FSTTestDoc("rooms/eros/messages/3", 0, @{ @"sort" : @2 }, NO);
@@ -205,7 +206,7 @@ NS_ASSUME_NONNULL_BEGIN
 - (void)testRemovesDocumentsForQueryWithLimit {
   FSTQuery *query = [self queryForMessages];
   query = [query queryBySettingLimit:2];
-  FSTView *view = [[FSTView alloc] initWithQuery:query remoteDocuments:[FSTDocumentKeySet keySet]];
+  FSTView *view = [[FSTView alloc] initWithQuery:query remoteDocuments:DocumentKeySet{}];
 
   FSTDocument *doc1 = FSTTestDoc("rooms/eros/messages/1", 0, @{@"text" : @"msg1"}, NO);
   FSTDocument *doc2 = FSTTestDoc("rooms/eros/messages/2", 0, @{@"text" : @"msg2"}, NO);
@@ -239,7 +240,7 @@ NS_ASSUME_NONNULL_BEGIN
   query = [query queryByAddingSortOrder:[FSTSortOrder sortOrderWithFieldPath:testutil::Field("num")
                                                                    ascending:YES]];
   query = [query queryBySettingLimit:2];
-  FSTView *view = [[FSTView alloc] initWithQuery:query remoteDocuments:[FSTDocumentKeySet keySet]];
+  FSTView *view = [[FSTView alloc] initWithQuery:query remoteDocuments:DocumentKeySet{}];
 
   FSTDocument *doc1 = FSTTestDoc("rooms/eros/messages/1", 0, @{ @"num" : @1 }, NO);
   FSTDocument *doc2 = FSTTestDoc("rooms/eros/messages/2", 0, @{ @"num" : @2 }, NO);
@@ -283,7 +284,7 @@ NS_ASSUME_NONNULL_BEGIN
 
 - (void)testKeepsTrackOfLimboDocuments {
   FSTQuery *query = [self queryForMessages];
-  FSTView *view = [[FSTView alloc] initWithQuery:query remoteDocuments:[FSTDocumentKeySet keySet]];
+  FSTView *view = [[FSTView alloc] initWithQuery:query remoteDocuments:DocumentKeySet{}];
 
   FSTDocument *doc1 = FSTTestDoc("rooms/eros/messages/0", 0, @{}, NO);
   FSTDocument *doc2 = FSTTestDoc("rooms/eros/messages/1", 0, @{}, NO);
@@ -339,8 +340,8 @@ NS_ASSUME_NONNULL_BEGIN
 
   // Unlike other cases, here the view is initialized with a set of previously synced documents
   // which happens when listening to a previously listened-to query.
-  FSTView *view = [[FSTView alloc] initWithQuery:query
-                                 remoteDocuments:FSTTestDocKeySet(@[ doc1.key, doc2.key ])];
+  FSTView *view =
+      [[FSTView alloc] initWithQuery:query remoteDocuments:DocumentKeySet{doc1.key, doc2.key}];
 
   FSTTargetChange *markCurrent =
       [FSTTargetChange changeWithDocuments:@[]
@@ -361,7 +362,7 @@ NS_ASSUME_NONNULL_BEGIN
   FSTQuery *query = [[self queryForMessages] queryBySettingLimit:2];
   FSTDocument *doc1 = FSTTestDoc("rooms/eros/messages/0", 0, @{}, NO);
   FSTDocument *doc2 = FSTTestDoc("rooms/eros/messages/1", 0, @{}, NO);
-  FSTView *view = [[FSTView alloc] initWithQuery:query remoteDocuments:[FSTDocumentKeySet keySet]];
+  FSTView *view = [[FSTView alloc] initWithQuery:query remoteDocuments:DocumentKeySet{}];
 
   // Start with a full view.
   FSTViewDocumentChanges *changes =
@@ -394,7 +395,7 @@ NS_ASSUME_NONNULL_BEGIN
   FSTDocument *doc1 = FSTTestDoc("rooms/eros/messages/0", 0, @{ @"order" : @1 }, NO);
   FSTDocument *doc2 = FSTTestDoc("rooms/eros/messages/1", 0, @{ @"order" : @2 }, NO);
   FSTDocument *doc3 = FSTTestDoc("rooms/eros/messages/2", 0, @{ @"order" : @3 }, NO);
-  FSTView *view = [[FSTView alloc] initWithQuery:query remoteDocuments:[FSTDocumentKeySet keySet]];
+  FSTView *view = [[FSTView alloc] initWithQuery:query remoteDocuments:DocumentKeySet{}];
 
   // Start with a full view.
   FSTViewDocumentChanges *changes =
@@ -430,7 +431,7 @@ NS_ASSUME_NONNULL_BEGIN
   FSTDocument *doc3 = FSTTestDoc("rooms/eros/messages/2", 0, @{ @"order" : @3 }, NO);
   FSTDocument *doc4 = FSTTestDoc("rooms/eros/messages/3", 0, @{ @"order" : @4 }, NO);
   FSTDocument *doc5 = FSTTestDoc("rooms/eros/messages/4", 0, @{ @"order" : @5 }, NO);
-  FSTView *view = [[FSTView alloc] initWithQuery:query remoteDocuments:[FSTDocumentKeySet keySet]];
+  FSTView *view = [[FSTView alloc] initWithQuery:query remoteDocuments:DocumentKeySet{}];
 
   // Start with a full view.
   FSTViewDocumentChanges *changes =
@@ -460,7 +461,7 @@ NS_ASSUME_NONNULL_BEGIN
   FSTDocument *doc3 = FSTTestDoc("rooms/eros/messages/2", 0, @{ @"order" : @3 }, NO);
   FSTDocument *doc4 = FSTTestDoc("rooms/eros/messages/3", 0, @{ @"order" : @4 }, NO);
   FSTDocument *doc5 = FSTTestDoc("rooms/eros/messages/4", 0, @{ @"order" : @5 }, NO);
-  FSTView *view = [[FSTView alloc] initWithQuery:query remoteDocuments:[FSTDocumentKeySet keySet]];
+  FSTView *view = [[FSTView alloc] initWithQuery:query remoteDocuments:DocumentKeySet{}];
 
   // Start with a full view.
   FSTViewDocumentChanges *changes =
@@ -483,7 +484,7 @@ NS_ASSUME_NONNULL_BEGIN
   FSTQuery *query = [[self queryForMessages] queryBySettingLimit:2];
   FSTDocument *doc1 = FSTTestDoc("rooms/eros/messages/0", 0, @{}, NO);
   FSTDocument *doc2 = FSTTestDoc("rooms/eros/messages/1", 0, @{}, NO);
-  FSTView *view = [[FSTView alloc] initWithQuery:query remoteDocuments:[FSTDocumentKeySet keySet]];
+  FSTView *view = [[FSTView alloc] initWithQuery:query remoteDocuments:DocumentKeySet{}];
 
   // Start with a full view.
   FSTViewDocumentChanges *changes =
@@ -506,7 +507,7 @@ NS_ASSUME_NONNULL_BEGIN
   FSTQuery *query = [[self queryForMessages] queryBySettingLimit:20];
   FSTDocument *doc1 = FSTTestDoc("rooms/eros/messages/0", 0, @{}, NO);
   FSTDocument *doc2 = FSTTestDoc("rooms/eros/messages/1", 0, @{}, NO);
-  FSTView *view = [[FSTView alloc] initWithQuery:query remoteDocuments:[FSTDocumentKeySet keySet]];
+  FSTView *view = [[FSTView alloc] initWithQuery:query remoteDocuments:DocumentKeySet{}];
 
   FSTViewDocumentChanges *changes =
       [view computeChangesWithDocuments:FSTTestDocUpdates(@[ doc1, doc2 ])];
@@ -528,7 +529,7 @@ NS_ASSUME_NONNULL_BEGIN
   FSTQuery *query = [[self queryForMessages] queryBySettingLimit:2];
   FSTDocument *doc1 = FSTTestDoc("rooms/eros/messages/0", 0, @{}, NO);
   FSTDocument *doc2 = FSTTestDoc("rooms/eros/messages/1", 0, @{}, NO);
-  FSTView *view = [[FSTView alloc] initWithQuery:query remoteDocuments:[FSTDocumentKeySet keySet]];
+  FSTView *view = [[FSTView alloc] initWithQuery:query remoteDocuments:DocumentKeySet{}];
 
   // Start with a full view.
   FSTViewDocumentChanges *changes =
@@ -551,69 +552,69 @@ NS_ASSUME_NONNULL_BEGIN
   FSTQuery *query = [self queryForMessages];
   FSTDocument *doc1 = FSTTestDoc("rooms/eros/messages/0", 0, @{}, NO);
   FSTDocument *doc2 = FSTTestDoc("rooms/eros/messages/1", 0, @{}, NO);
-  FSTView *view = [[FSTView alloc] initWithQuery:query remoteDocuments:[FSTDocumentKeySet keySet]];
+  FSTView *view = [[FSTView alloc] initWithQuery:query remoteDocuments:DocumentKeySet{}];
 
   // Start with a full view.
   FSTViewDocumentChanges *changes =
       [view computeChangesWithDocuments:FSTTestDocUpdates(@[ doc1, doc2 ])];
   [view applyChangesToDocuments:changes];
-  XCTAssertEqualObjects(changes.mutatedKeys, FSTTestDocKeySet(@[]));
+  XCTAssertEqual(changes.mutatedKeys, DocumentKeySet{});
 
   FSTDocument *doc3 = FSTTestDoc("rooms/eros/messages/2", 0, @{}, YES);
   changes = [view computeChangesWithDocuments:FSTTestDocUpdates(@[ doc3 ])];
-  XCTAssertEqualObjects(changes.mutatedKeys, FSTTestDocKeySet(@[ doc3.key ]));
+  XCTAssertEqual(changes.mutatedKeys, DocumentKeySet{doc3.key});
 }
 
 - (void)testRemovesKeysFromMutatedKeysWhenNewDocHasNoLocalChanges {
   FSTQuery *query = [self queryForMessages];
   FSTDocument *doc1 = FSTTestDoc("rooms/eros/messages/0", 0, @{}, NO);
   FSTDocument *doc2 = FSTTestDoc("rooms/eros/messages/1", 0, @{}, YES);
-  FSTView *view = [[FSTView alloc] initWithQuery:query remoteDocuments:[FSTDocumentKeySet keySet]];
+  FSTView *view = [[FSTView alloc] initWithQuery:query remoteDocuments:DocumentKeySet{}];
 
   // Start with a full view.
   FSTViewDocumentChanges *changes =
       [view computeChangesWithDocuments:FSTTestDocUpdates(@[ doc1, doc2 ])];
   [view applyChangesToDocuments:changes];
-  XCTAssertEqualObjects(changes.mutatedKeys, FSTTestDocKeySet(@[ doc2.key ]));
+  XCTAssertEqual(changes.mutatedKeys, (DocumentKeySet{doc2.key}));
 
   FSTDocument *doc2Prime = FSTTestDoc("rooms/eros/messages/1", 0, @{}, NO);
   changes = [view computeChangesWithDocuments:FSTTestDocUpdates(@[ doc2Prime ])];
   [view applyChangesToDocuments:changes];
-  XCTAssertEqualObjects(changes.mutatedKeys, FSTTestDocKeySet(@[]));
+  XCTAssertEqual(changes.mutatedKeys, DocumentKeySet{});
 }
 
 - (void)testRemembersLocalMutationsFromPreviousSnapshot {
   FSTQuery *query = [self queryForMessages];
   FSTDocument *doc1 = FSTTestDoc("rooms/eros/messages/0", 0, @{}, NO);
   FSTDocument *doc2 = FSTTestDoc("rooms/eros/messages/1", 0, @{}, YES);
-  FSTView *view = [[FSTView alloc] initWithQuery:query remoteDocuments:[FSTDocumentKeySet keySet]];
+  FSTView *view = [[FSTView alloc] initWithQuery:query remoteDocuments:DocumentKeySet{}];
 
   // Start with a full view.
   FSTViewDocumentChanges *changes =
       [view computeChangesWithDocuments:FSTTestDocUpdates(@[ doc1, doc2 ])];
   [view applyChangesToDocuments:changes];
-  XCTAssertEqualObjects(changes.mutatedKeys, FSTTestDocKeySet(@[ doc2.key ]));
+  XCTAssertEqual(changes.mutatedKeys, (DocumentKeySet{doc2.key}));
 
   FSTDocument *doc3 = FSTTestDoc("rooms/eros/messages/2", 0, @{}, NO);
   changes = [view computeChangesWithDocuments:FSTTestDocUpdates(@[ doc3 ])];
   [view applyChangesToDocuments:changes];
-  XCTAssertEqualObjects(changes.mutatedKeys, FSTTestDocKeySet(@[ doc2.key ]));
+  XCTAssertEqual(changes.mutatedKeys, (DocumentKeySet{doc2.key}));
 }
 
 - (void)testRemembersLocalMutationsFromPreviousCallToComputeChangesWithDocuments {
   FSTQuery *query = [self queryForMessages];
   FSTDocument *doc1 = FSTTestDoc("rooms/eros/messages/0", 0, @{}, NO);
   FSTDocument *doc2 = FSTTestDoc("rooms/eros/messages/1", 0, @{}, YES);
-  FSTView *view = [[FSTView alloc] initWithQuery:query remoteDocuments:[FSTDocumentKeySet keySet]];
+  FSTView *view = [[FSTView alloc] initWithQuery:query remoteDocuments:DocumentKeySet{}];
 
   // Start with a full view.
   FSTViewDocumentChanges *changes =
       [view computeChangesWithDocuments:FSTTestDocUpdates(@[ doc1, doc2 ])];
-  XCTAssertEqualObjects(changes.mutatedKeys, FSTTestDocKeySet(@[ doc2.key ]));
+  XCTAssertEqual(changes.mutatedKeys, (DocumentKeySet{doc2.key}));
 
   FSTDocument *doc3 = FSTTestDoc("rooms/eros/messages/2", 0, @{}, NO);
   changes = [view computeChangesWithDocuments:FSTTestDocUpdates(@[ doc3 ]) previousChanges:changes];
-  XCTAssertEqualObjects(changes.mutatedKeys, FSTTestDocKeySet(@[ doc2.key ]));
+  XCTAssertEqual(changes.mutatedKeys, (DocumentKeySet{doc2.key}));
 }
 
 @end
diff --git a/Firestore/Example/Tests/Local/FSTLocalSerializerTests.mm b/Firestore/Example/Tests/Local/FSTLocalSerializerTests.mm
index 421f244..0e160c0 100644
--- a/Firestore/Example/Tests/Local/FSTLocalSerializerTests.mm
+++ b/Firestore/Example/Tests/Local/FSTLocalSerializerTests.mm
@@ -125,7 +125,7 @@ NS_ASSUME_NONNULL_BEGIN
   XCTAssertEqual(decoded.batchID, model.batchID);
   XCTAssertEqualObjects(decoded.localWriteTime, model.localWriteTime);
   XCTAssertEqualObjects(decoded.mutations, model.mutations);
-  XCTAssertEqualObjects([decoded keys], [model keys]);
+  XCTAssertEqual([decoded keys], [model keys]);
 }
 
 - (void)testEncodesDocumentAsMaybeDocument {
diff --git a/Firestore/Example/Tests/Local/FSTLocalStoreTests.mm b/Firestore/Example/Tests/Local/FSTLocalStoreTests.mm
index 4175118..c417646 100644
--- a/Firestore/Example/Tests/Local/FSTLocalStoreTests.mm
+++ b/Firestore/Example/Tests/Local/FSTLocalStoreTests.mm
@@ -46,6 +46,7 @@
 namespace testutil = firebase::firestore::testutil;
 using firebase::firestore::auth::User;
 using firebase::firestore::model::SnapshotVersion;
+using firebase::firestore::model::DocumentKeySet;
 
 NS_ASSUME_NONNULL_BEGIN
 
@@ -230,8 +231,8 @@ FSTDocumentVersionDictionary *FSTVersionDictionary(FSTMutation *mutation,
   FSTMutationBatch *batch = [[FSTMutationBatch alloc] initWithBatchID:1
                                                        localWriteTime:[FIRTimestamp timestamp]
                                                             mutations:@[ set1, set2 ]];
-  FSTDocumentKeySet *keys = [batch keys];
-  XCTAssertEqual(keys.count, 2);
+  DocumentKeySet keys = [batch keys];
+  XCTAssertEqual(keys.size(), 2u);
 }
 
 - (void)testHandlesSetMutation {
@@ -856,13 +857,14 @@ FSTDocumentVersionDictionary *FSTVersionDictionary(FSTMutation *mutation,
 
   [self.localStore locallyWriteMutations:@[ FSTTestSetMutation(@"foo/bonk", @{@"a" : @"b"}) ]];
 
-  FSTDocumentKeySet *keys = [self.localStore remoteDocumentKeysForTarget:2];
-  FSTAssertEqualSets(keys, (@[ FSTTestDocKey(@"foo/bar"), FSTTestDocKey(@"foo/baz") ]));
+  DocumentKeySet keys = [self.localStore remoteDocumentKeysForTarget:2];
+  DocumentKeySet expected{testutil::Key("foo/bar"), testutil::Key("foo/baz")};
+  XCTAssertEqual(keys, expected);
 
   [self restartWithNoopGarbageCollector];
 
   keys = [self.localStore remoteDocumentKeysForTarget:2];
-  FSTAssertEqualSets(keys, (@[ FSTTestDocKey(@"foo/bar"), FSTTestDocKey(@"foo/baz") ]));
+  XCTAssertEqual(keys, (DocumentKeySet{testutil::Key("foo/bar"), testutil::Key("foo/baz")}));
 }
 
 @end
diff --git a/Firestore/Example/Tests/Local/FSTQueryCacheTests.mm b/Firestore/Example/Tests/Local/FSTQueryCacheTests.mm
index 5d07516..44b49de 100644
--- a/Firestore/Example/Tests/Local/FSTQueryCacheTests.mm
+++ b/Firestore/Example/Tests/Local/FSTQueryCacheTests.mm
@@ -32,6 +32,7 @@
 namespace testutil = firebase::firestore::testutil;
 using firebase::firestore::model::DocumentKey;
 using firebase::firestore::model::SnapshotVersion;
+using firebase::firestore::model::DocumentKeySet;
 
 NS_ASSUME_NONNULL_BEGIN
 
@@ -278,12 +279,12 @@ NS_ASSUME_NONNULL_BEGIN
     [self addMatchingKey:key2 forTargetID:1];
     [self addMatchingKey:key3 forTargetID:2];
 
-    FSTAssertEqualSets([self.queryCache matchingKeysForTargetID:1], (@[ key1, key2 ]));
-    FSTAssertEqualSets([self.queryCache matchingKeysForTargetID:2], @[ key3 ]);
+    XCTAssertEqual([self.queryCache matchingKeysForTargetID:1], (DocumentKeySet{key1, key2}));
+    XCTAssertEqual([self.queryCache matchingKeysForTargetID:2], (DocumentKeySet{key3}));
 
     [self addMatchingKey:key1 forTargetID:2];
-    FSTAssertEqualSets([self.queryCache matchingKeysForTargetID:1], (@[ key1, key2 ]));
-    FSTAssertEqualSets([self.queryCache matchingKeysForTargetID:2], (@[ key1, key3 ]));
+    XCTAssertEqual([self.queryCache matchingKeysForTargetID:1], (DocumentKeySet{key1, key2}));
+    XCTAssertEqual([self.queryCache matchingKeysForTargetID:2], (DocumentKeySet{key1, key3}));
   });
 }
 
@@ -428,14 +429,12 @@ NS_ASSUME_NONNULL_BEGIN
 }
 
 - (void)addMatchingKey:(const DocumentKey &)key forTargetID:(FSTTargetID)targetID {
-  FSTDocumentKeySet *keys = [FSTDocumentKeySet keySet];
-  keys = [keys setByAddingObject:key];
+  DocumentKeySet keys{key};
   [self.queryCache addMatchingKeys:keys forTargetID:targetID];
 }
 
 - (void)removeMatchingKey:(const DocumentKey &)key forTargetID:(FSTTargetID)targetID {
-  FSTDocumentKeySet *keys = [FSTDocumentKeySet keySet];
-  keys = [keys setByAddingObject:key];
+  DocumentKeySet keys{key};
   [self.queryCache removeMatchingKeys:keys forTargetID:targetID];
 }
 
diff --git a/Firestore/Example/Tests/Remote/FSTRemoteEventTests.mm b/Firestore/Example/Tests/Remote/FSTRemoteEventTests.mm
index e2cb8cc..84d0fa1 100644
--- a/Firestore/Example/Tests/Remote/FSTRemoteEventTests.mm
+++ b/Firestore/Example/Tests/Remote/FSTRemoteEventTests.mm
@@ -33,6 +33,7 @@
 
 namespace testutil = firebase::firestore::testutil;
 using firebase::firestore::model::DocumentKey;
+using firebase::firestore::model::DocumentKeySet;
 
 NS_ASSUME_NONNULL_BEGIN
 
@@ -575,7 +576,7 @@ NS_ASSUME_NONNULL_BEGIN
   FSTDeletedDocument *expected =
       [FSTDeletedDocument documentWithKey:synthesized version:event.snapshotVersion];
   XCTAssertEqualObjects(expected, event.documentUpdates.at(synthesized));
-  XCTAssertTrue([event.limboDocumentChanges containsObject:synthesized]);
+  XCTAssertTrue(event.limboDocumentChanges.contains(synthesized));
 }
 
 - (void)testDoesntSynthesizeDeletesForWrongState {
@@ -590,7 +591,7 @@ NS_ASSUME_NONNULL_BEGIN
   DocumentKey notSynthesized = DocumentKey::FromPathString("docs/no1");
   [event synthesizeDeleteForLimboTargetChange:event.targetChanges[@2] key:notSynthesized];
   XCTAssertEqual(event.documentUpdates.find(notSynthesized), event.documentUpdates.end());
-  XCTAssertFalse([event.limboDocumentChanges containsObject:notSynthesized]);
+  XCTAssertFalse(event.limboDocumentChanges.contains(notSynthesized));
 }
 
 - (void)testDoesntSynthesizeDeletesForExistingDoc {
@@ -609,7 +610,7 @@ NS_ASSUME_NONNULL_BEGIN
   [event synthesizeDeleteForLimboTargetChange:event.targetChanges[@3] key:doc.key];
   FSTMaybeDocument *docData = event.documentUpdates.at(doc.key);
   XCTAssertFalse([docData isKindOfClass:[FSTDeletedDocument class]]);
-  XCTAssertFalse([event.limboDocumentChanges containsObject:doc.key]);
+  XCTAssertFalse(event.limboDocumentChanges.contains(doc.key));
 }
 
 - (void)testFilterUpdates {
@@ -631,24 +632,24 @@ NS_ASSUME_NONNULL_BEGIN
                       outstanding:_noPendingResponses
                           changes:@[ newDocChange, existingDocChange ]];
   FSTRemoteEvent *event = [aggregator remoteEvent];
-  FSTDocumentKeySet *existingKeys = [[FSTDocumentKeySet keySet] setByAddingObject:existingDoc.key];
+  DocumentKeySet existingKeys = DocumentKeySet{existingDoc.key};
 
   FSTTargetChange *updateChange = event.targetChanges[@1];
   XCTAssertTrue([updateChange.mapping isKindOfClass:[FSTUpdateMapping class]]);
   FSTUpdateMapping *update = (FSTUpdateMapping *)updateChange.mapping;
   FSTDocumentKey *existingDocKey = existingDoc.key;
   FSTDocumentKey *newDocKey = newDoc.key;
-  XCTAssertTrue([update.addedDocuments containsObject:existingDocKey]);
+  XCTAssertTrue(update.addedDocuments.contains(existingDocKey));
 
   [update filterUpdatesUsingExistingKeys:existingKeys];
   // Now it's been filtered, since it already existed.
-  XCTAssertFalse([update.addedDocuments containsObject:existingDocKey]);
-  XCTAssertTrue([update.addedDocuments containsObject:newDocKey]);
+  XCTAssertFalse(update.addedDocuments.contains(existingDocKey));
+  XCTAssertTrue(update.addedDocuments.contains(newDocKey));
 }
 
 - (void)testDoesntFilterResets {
   FSTDocument *existingDoc = FSTTestDoc("docs/existing", 1, @{@"some" : @"data"}, NO);
-  FSTDocumentKey *existingDocKey = existingDoc.key;
+  const DocumentKey &existingDocKey = existingDoc.key;
   FSTWatchTargetChange *resetTargetChange =
       [FSTWatchTargetChange changeWithState:FSTWatchTargetChangeStateReset
                                   targetIDs:@[ @2 ]
@@ -663,16 +664,16 @@ NS_ASSUME_NONNULL_BEGIN
                       outstanding:_noPendingResponses
                           changes:@[ resetTargetChange, existingDocChange ]];
   FSTRemoteEvent *event = [aggregator remoteEvent];
-  FSTDocumentKeySet *existingKeys = [[FSTDocumentKeySet keySet] setByAddingObject:existingDocKey];
+  DocumentKeySet existingKeys = DocumentKeySet{existingDocKey};
 
   FSTTargetChange *resetChange = event.targetChanges[@2];
   XCTAssertTrue([resetChange.mapping isKindOfClass:[FSTResetMapping class]]);
   FSTResetMapping *resetMapping = (FSTResetMapping *)resetChange.mapping;
-  XCTAssertTrue([resetMapping.documents containsObject:existingDocKey]);
+  XCTAssertTrue(resetMapping.documents.contains(existingDocKey));
 
   [resetMapping filterUpdatesUsingExistingKeys:existingKeys];
   // Document is still there, even though it already exists. Reset mappings don't get filtered.
-  XCTAssertTrue([resetMapping.documents containsObject:existingDocKey]);
+  XCTAssertTrue(resetMapping.documents.contains(existingDocKey));
 }
 
 - (void)testTracksLimboDocuments {
@@ -715,13 +716,13 @@ NS_ASSUME_NONNULL_BEGIN
   [aggregator addWatchChanges:@[ docChange1, docChange2, docChange3, targetsChange ]];
 
   FSTRemoteEvent *event = [aggregator remoteEvent];
-  FSTDocumentKeySet *limboDocChanges = event.limboDocumentChanges;
+  DocumentKeySet limboDocChanges = event.limboDocumentChanges;
   // Doc1 is in both limbo and non-limbo targets, therefore not tracked as limbo
-  XCTAssertFalse([limboDocChanges containsObject:doc1.key]);
+  XCTAssertFalse(limboDocChanges.contains(doc1.key));
   // Doc2 is only in the limbo target, so is tracked as a limbo document
-  XCTAssertTrue([limboDocChanges containsObject:doc2.key]);
+  XCTAssertTrue(limboDocChanges.contains(doc2.key));
   // Doc3 is only in the non-limbo target, therefore not tracked as limbo
-  XCTAssertFalse([limboDocChanges containsObject:doc3.key]);
+  XCTAssertFalse(limboDocChanges.contains(doc3.key));
 }
 
 @end
diff --git a/Firestore/Example/Tests/SpecTests/FSTSpecTests.mm b/Firestore/Example/Tests/SpecTests/FSTSpecTests.mm
index 12be590..5a7cb72 100644
--- a/Firestore/Example/Tests/SpecTests/FSTSpecTests.mm
+++ b/Firestore/Example/Tests/SpecTests/FSTSpecTests.mm
@@ -603,10 +603,13 @@ static NSString *const kNoIOSTag = @"no-ios";
     // XCTAssertEqualObjects(actualTargets[targetID], queryData);
 
     FSTQueryData *actual = actualTargets[targetID];
-    XCTAssertEqualObjects(actual.query, queryData.query);
-    XCTAssertEqual(actual.targetID, queryData.targetID);
-    XCTAssertEqual(actual.snapshotVersion, queryData.snapshotVersion);
-    XCTAssertEqualObjects(actual.resumeToken, queryData.resumeToken);
+    XCTAssertNotNil(actual);
+    if (actual) {
+      XCTAssertEqualObjects(actual.query, queryData.query);
+      XCTAssertEqual(actual.targetID, queryData.targetID);
+      XCTAssertEqual(actual.snapshotVersion, queryData.snapshotVersion);
+      XCTAssertEqualObjects(actual.resumeToken, queryData.resumeToken);
+    }
 
     [actualTargets removeObjectForKey:targetID];
   }];
diff --git a/Firestore/Example/Tests/Util/FSTHelpers.h b/Firestore/Example/Tests/Util/FSTHelpers.h
index f80f163..ccc01ca 100644
--- a/Firestore/Example/Tests/Util/FSTHelpers.h
+++ b/Firestore/Example/Tests/Util/FSTHelpers.h
@@ -21,7 +21,6 @@
 
 #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/field_path.h"
 #include "Firestore/core/src/firebase/firestore/model/field_value.h"
@@ -181,9 +180,6 @@ FSTObjectValue *FSTTestObjectValue(NSDictionary<NSString *, id> *data);
 /** A convenience method for creating document keys for tests. */
 FSTDocumentKey *FSTTestDocKey(NSString *path);
 
-/** A convenience method for creating a document key set for tests. */
-FSTDocumentKeySet *FSTTestDocKeySet(NSArray<FSTDocumentKey *> *keys);
-
 /** Allow tests to just use an int literal for versions. */
 typedef int64_t FSTTestSnapshotVersion;
 
diff --git a/Firestore/Example/Tests/Util/FSTHelpers.mm b/Firestore/Example/Tests/Util/FSTHelpers.mm
index 92eb730..dbd19aa 100644
--- a/Firestore/Example/Tests/Util/FSTHelpers.mm
+++ b/Firestore/Example/Tests/Util/FSTHelpers.mm
@@ -44,6 +44,7 @@
 
 #include "Firestore/core/src/firebase/firestore/model/database_id.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/field_mask.h"
 #include "Firestore/core/src/firebase/firestore/model/field_transform.h"
 #include "Firestore/core/src/firebase/firestore/model/field_value.h"
@@ -66,6 +67,7 @@ using firebase::firestore::model::Precondition;
 using firebase::firestore::model::ResourcePath;
 using firebase::firestore::model::ServerTimestampTransform;
 using firebase::firestore::model::TransformOperation;
+using firebase::firestore::model::DocumentKeySet;
 
 NS_ASSUME_NONNULL_BEGIN
 
@@ -146,14 +148,6 @@ FSTDocumentKey *FSTTestDocKey(NSString *path) {
   return [FSTDocumentKey keyWithPathString:path];
 }
 
-FSTDocumentKeySet *FSTTestDocKeySet(NSArray<FSTDocumentKey *> *keys) {
-  FSTDocumentKeySet *result = [FSTDocumentKeySet keySet];
-  for (FSTDocumentKey *key in keys) {
-    result = [result setByAddingObject:key];
-  }
-  return result;
-}
-
 FSTDocument *FSTTestDoc(const absl::string_view path,
                         FSTTestSnapshotVersion version,
                         NSDictionary<NSString *, id> *data,
@@ -343,17 +337,17 @@ NSData *_Nullable FSTTestResumeTokenFromSnapshotVersion(FSTTestSnapshotVersion s
 FSTLocalViewChanges *FSTTestViewChanges(FSTQuery *query,
                                         NSArray<NSString *> *addedKeys,
                                         NSArray<NSString *> *removedKeys) {
-  FSTDocumentKeySet *added = [FSTDocumentKeySet keySet];
+  DocumentKeySet added;
   for (NSString *keyPath in addedKeys) {
-    FSTDocumentKey *key = FSTTestDocKey(keyPath);
-    added = [added setByAddingObject:key];
+    added = added.insert(testutil::Key(util::MakeStringView(keyPath)));
   }
-  FSTDocumentKeySet *removed = [FSTDocumentKeySet keySet];
+  DocumentKeySet removed;
   for (NSString *keyPath in removedKeys) {
-    FSTDocumentKey *key = FSTTestDocKey(keyPath);
-    removed = [removed setByAddingObject:key];
+    removed = removed.insert(testutil::Key(util::MakeStringView(keyPath)));
   }
-  return [FSTLocalViewChanges changesForQuery:query addedKeys:added removedKeys:removed];
+  return [FSTLocalViewChanges changesForQuery:query
+                                    addedKeys:std::move(added)
+                                  removedKeys:std::move(removed)];
 }
 
 NS_ASSUME_NONNULL_END
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
diff --git a/Firestore/core/src/firebase/firestore/immutable/sorted_set.h b/Firestore/core/src/firebase/firestore/immutable/sorted_set.h
index 6828106..d78fd61 100644
--- a/Firestore/core/src/firebase/firestore/immutable/sorted_set.h
+++ b/Firestore/core/src/firebase/firestore/immutable/sorted_set.h
@@ -61,6 +61,13 @@ class SortedSet {
   explicit SortedSet(M&& map) : map_{std::move(map)} {
   }
 
+  SortedSet(std::initializer_list<value_type> entries, const C& comparator = {})
+      : map_{comparator} {
+    for (auto&& value : entries) {
+      map_ = map_.insert(value, {});
+    }
+  }
+
   bool empty() const {
     return map_.empty();
   }
@@ -124,7 +131,14 @@ class SortedSet {
   }
 
   friend bool operator==(const SortedSet& lhs, const SortedSet& rhs) {
-    return std::equal(lhs.begin(), lhs.end(), rhs.begin(), rhs.end());
+    if (lhs.size() != rhs.size()) {
+      return false;
+    }
+    return std::equal(lhs.begin(), lhs.end(), rhs.begin());
+  }
+
+  friend bool operator!=(const SortedSet& lhs, const SortedSet& rhs) {
+    return !(lhs == rhs);
   }
 
  private:
diff --git a/Firestore/core/src/firebase/firestore/model/document_key.h b/Firestore/core/src/firebase/firestore/model/document_key.h
index 7f6478f..6532b2e 100644
--- a/Firestore/core/src/firebase/firestore/model/document_key.h
+++ b/Firestore/core/src/firebase/firestore/model/document_key.h
@@ -17,6 +17,7 @@
 #ifndef FIRESTORE_CORE_SRC_FIREBASE_FIRESTORE_MODEL_DOCUMENT_KEY_H_
 #define FIRESTORE_CORE_SRC_FIREBASE_FIRESTORE_MODEL_DOCUMENT_KEY_H_
 
+#include <functional>
 #include <initializer_list>
 #include <memory>
 #include <string>
@@ -26,6 +27,7 @@
 #endif  // defined(__OBJC__)
 
 #include "Firestore/core/src/firebase/firestore/model/resource_path.h"
+#include "Firestore/core/src/firebase/firestore/util/comparison.h"
 #include "Firestore/core/src/firebase/firestore/util/hashing.h"
 #include "absl/strings/string_view.h"
 
@@ -118,6 +120,13 @@ inline bool operator>=(const DocumentKey& lhs, const DocumentKey& rhs) {
 }
 
 }  // namespace model
+
+namespace util {
+
+template <>
+struct Comparator<model::DocumentKey> : public std::less<model::DocumentKey> {};
+
+}  // namespace util
 }  // namespace firestore
 }  // namespace firebase
 
diff --git a/Firestore/core/test/firebase/firestore/model/document_key_test.cc b/Firestore/core/test/firebase/firestore/model/document_key_test.cc
index 619ee7f..71b78d1 100644
--- a/Firestore/core/test/firebase/firestore/model/document_key_test.cc
+++ b/Firestore/core/test/firebase/firestore/model/document_key_test.cc
@@ -21,8 +21,11 @@
 
 #include "Firestore/core/src/firebase/firestore/model/document_key.h"
 #include "Firestore/core/src/firebase/firestore/model/resource_path.h"
+#include "Firestore/core/test/firebase/firestore/testutil/testutil.h"
 #include "gtest/gtest.h"
 
+using firebase::firestore::testutil::Key;
+
 namespace firebase {
 namespace firestore {
 namespace model {
@@ -112,16 +115,16 @@ TEST(DocumentKey, IsDocumentKey) {
 }
 
 TEST(DocumentKey, Comparison) {
-  const DocumentKey abcd({"a", "b", "c", "d"});
-  const DocumentKey abcd_too({"a", "b", "c", "d"});
-  const DocumentKey xyzw({"x", "y", "z", "w"});
+  DocumentKey abcd = Key("a/b/c/d");
+  DocumentKey abcd_too = Key("a/b/c/d");
+  DocumentKey xyzw = Key("x/y/z/w");
   EXPECT_EQ(abcd, abcd_too);
   EXPECT_NE(abcd, xyzw);
 
-  const DocumentKey empty;
-  const DocumentKey a({"a", "a"});
-  const DocumentKey b({"b", "b"});
-  const DocumentKey ab({"a", "a", "b", "b"});
+  DocumentKey empty;
+  DocumentKey a = Key("a/a");
+  DocumentKey b = Key("b/b");
+  DocumentKey ab = Key("a/a/b/b");
 
   EXPECT_FALSE(empty < empty);
   EXPECT_TRUE(empty <= empty);
@@ -148,6 +151,12 @@ TEST(DocumentKey, Comparison) {
   EXPECT_TRUE(ab >= a);
 }
 
+TEST(DocumentKey, Comparator) {
+  DocumentKey abcd = Key("a/b/c/d");
+  DocumentKey xyzw = Key("x/y/z/w");
+  EXPECT_TRUE(util::Comparator<DocumentKey>{}(abcd, xyzw));
+}
+
 }  // namespace model
 }  // namespace firestore
 }  // namespace firebase
author	Gil <mcg@google.com>	2018-05-05 08:10:51 -0700
committer	GitHub <noreply@github.com>	2018-05-05 08:10:51 -0700
commit	a080e481b5e6fcbc2b645920051cf20fc8cad7a7 (patch)
tree	5d5ce3d39be316ac82a2fd34c92b05fc9e67f83a /Firestore
parent	d158e420c6fa04997ee3d2a6c44fd53a52883d81 (diff)