port C++ `DocumentKey` to `Local/*` (#963)

* port C++ DocumentKey to Local's

* address changes

30 files changed, 242 insertions, 179 deletions

diff --git a/Firestore/Example/Tests/Local/FSTEagerGarbageCollectorTests.mm b/Firestore/Example/Tests/Local/FSTEagerGarbageCollectorTests.mm
index 53f0202..2cf530c 100644
--- a/Firestore/Example/Tests/Local/FSTEagerGarbageCollectorTests.mm
+++ b/Firestore/Example/Tests/Local/FSTEagerGarbageCollectorTests.mm
@@ -16,12 +16,18 @@
 
 #import "Firestore/Source/Local/FSTEagerGarbageCollector.h"
 
+#include <set>
+
 #import <XCTest/XCTest.h>
 
 #import "Firestore/Source/Local/FSTReferenceSet.h"
 #import "Firestore/Source/Model/FSTDocumentKey.h"
 
-#import "Firestore/Example/Tests/Util/FSTHelpers.h"
+#include "Firestore/core/src/firebase/firestore/model/document_key.h"
+#include "Firestore/core/test/firebase/firestore/testutil/testutil.h"
+
+namespace testutil = firebase::firestore::testutil;
+using firebase::firestore::model::DocumentKey;
 
 NS_ASSUME_NONNULL_BEGIN
 
@@ -35,13 +41,13 @@ NS_ASSUME_NONNULL_BEGIN
   FSTReferenceSet *referenceSet = [[FSTReferenceSet alloc] init];
   [gc addGarbageSource:referenceSet];
 
-  FSTDocumentKey *key = FSTTestDocKey(@"foo/bar");
+  DocumentKey key = testutil::Key("foo/bar");
   [referenceSet addReferenceToKey:key forID:1];
-  FSTAssertEqualSets([gc collectGarbage], @[]);
+  XCTAssertEqual([gc collectGarbage], std::set<DocumentKey>({}));
   XCTAssertFalse([referenceSet isEmpty]);
 
   [referenceSet removeReferenceToKey:key forID:1];
-  FSTAssertEqualSets([gc collectGarbage], @[ key ]);
+  XCTAssertEqual([gc collectGarbage], std::set<DocumentKey>({key}));
   XCTAssertTrue([referenceSet isEmpty]);
 }
 
@@ -50,20 +56,20 @@ NS_ASSUME_NONNULL_BEGIN
   FSTReferenceSet *referenceSet = [[FSTReferenceSet alloc] init];
   [gc addGarbageSource:referenceSet];
 
-  FSTDocumentKey *key1 = FSTTestDocKey(@"foo/bar");
-  FSTDocumentKey *key2 = FSTTestDocKey(@"foo/baz");
-  FSTDocumentKey *key3 = FSTTestDocKey(@"foo/blah");
+  DocumentKey key1 = testutil::Key("foo/bar");
+  DocumentKey key2 = testutil::Key("foo/baz");
+  DocumentKey key3 = testutil::Key("foo/blah");
   [referenceSet addReferenceToKey:key1 forID:1];
   [referenceSet addReferenceToKey:key2 forID:1];
   [referenceSet addReferenceToKey:key3 forID:2];
   XCTAssertFalse([referenceSet isEmpty]);
 
   [referenceSet removeReferencesForID:1];
-  FSTAssertEqualSets([gc collectGarbage], (@[ key1, key2 ]));
+  XCTAssertEqual([gc collectGarbage], std::set<DocumentKey>({key1, key2}));
   XCTAssertFalse([referenceSet isEmpty]);
 
   [referenceSet removeReferencesForID:2];
-  FSTAssertEqualSets([gc collectGarbage], @[ key3 ]);
+  XCTAssertEqual([gc collectGarbage], std::set<DocumentKey>({key3}));
   XCTAssertTrue([referenceSet isEmpty]);
 }
 
@@ -77,12 +83,12 @@ NS_ASSUME_NONNULL_BEGIN
   [gc addGarbageSource:localViews];
   [gc addGarbageSource:mutations];
 
-  FSTDocumentKey *key1 = FSTTestDocKey(@"foo/bar");
+  DocumentKey key1 = testutil::Key("foo/bar");
   [remoteTargets addReferenceToKey:key1 forID:1];
   [localViews addReferenceToKey:key1 forID:1];
   [mutations addReferenceToKey:key1 forID:10];
 
-  FSTDocumentKey *key2 = FSTTestDocKey(@"foo/baz");
+  DocumentKey key2 = testutil::Key("foo/baz");
   [mutations addReferenceToKey:key2 forID:10];
 
   XCTAssertFalse([remoteTargets isEmpty]);
@@ -90,16 +96,16 @@ NS_ASSUME_NONNULL_BEGIN
   XCTAssertFalse([mutations isEmpty]);
 
   [localViews removeReferencesForID:1];
-  FSTAssertEqualSets([gc collectGarbage], @[]);
+  XCTAssertEqual([gc collectGarbage], std::set<DocumentKey>({}));
 
   [remoteTargets removeReferencesForID:1];
-  FSTAssertEqualSets([gc collectGarbage], @[]);
+  XCTAssertEqual([gc collectGarbage], std::set<DocumentKey>({}));
 
   [mutations removeReferenceToKey:key1 forID:10];
-  FSTAssertEqualSets([gc collectGarbage], @[ key1 ]);
+  XCTAssertEqual([gc collectGarbage], std::set<DocumentKey>({key1}));
 
   [mutations removeReferenceToKey:key2 forID:10];
-  FSTAssertEqualSets([gc collectGarbage], @[ key2 ]);
+  XCTAssertEqual([gc collectGarbage], std::set<DocumentKey>({key2}));
 
   XCTAssertTrue([remoteTargets isEmpty]);
   XCTAssertTrue([localViews isEmpty]);
diff --git a/Firestore/Example/Tests/Local/FSTMutationQueueTests.mm b/Firestore/Example/Tests/Local/FSTMutationQueueTests.mm
index 57572ec..cad62cf 100644
--- a/Firestore/Example/Tests/Local/FSTMutationQueueTests.mm
+++ b/Firestore/Example/Tests/Local/FSTMutationQueueTests.mm
@@ -28,8 +28,12 @@
 #import "Firestore/Example/Tests/Util/FSTHelpers.h"
 
 #include "Firestore/core/src/firebase/firestore/auth/user.h"
+#include "Firestore/core/src/firebase/firestore/model/document_key.h"
+#include "Firestore/core/test/firebase/firestore/testutil/testutil.h"
 
+namespace testutil = firebase::firestore::testutil;
 using firebase::firestore::auth::User;
+using firebase::firestore::model::DocumentKey;
 
 NS_ASSUME_NONNULL_BEGIN
 
@@ -283,7 +287,7 @@ NS_ASSUME_NONNULL_BEGIN
 
   NSArray<FSTMutationBatch *> *expected = @[ batches[1], batches[2] ];
   NSArray<FSTMutationBatch *> *matches =
-      [self.mutationQueue allMutationBatchesAffectingDocumentKey:FSTTestDocKey(@"foo/bar")];
+      [self.mutationQueue allMutationBatchesAffectingDocumentKey:testutil::Key("foo/bar")];
 
   XCTAssertEqualObjects(matches, expected);
 }
@@ -395,28 +399,29 @@ NS_ASSUME_NONNULL_BEGIN
   ]];
 
   [self removeMutationBatches:@[ batches[0] ]];
-  NSSet<FSTDocumentKey *> *garbage = [garbageCollector collectGarbage];
-  FSTAssertEqualSets(garbage, @[]);
+  std::set<DocumentKey> garbage = [garbageCollector collectGarbage];
+  XCTAssertEqual(garbage, std::set<DocumentKey>({}));
 
   [self removeMutationBatches:@[ batches[1] ]];
   garbage = [garbageCollector collectGarbage];
-  FSTAssertEqualSets(garbage, @[ FSTTestDocKey(@"foo/ba") ]);
+  XCTAssertEqual(garbage, std::set<DocumentKey>({testutil::Key("foo/ba")}));
 
   [self removeMutationBatches:@[ batches[5] ]];
   garbage = [garbageCollector collectGarbage];
-  FSTAssertEqualSets(garbage, @[ FSTTestDocKey(@"bar/baz") ]);
+  XCTAssertEqual(garbage, std::set<DocumentKey>({testutil::Key("bar/baz")}));
 
   [self removeMutationBatches:@[ batches[2], batches[3] ]];
   garbage = [garbageCollector collectGarbage];
-  FSTAssertEqualSets(garbage, (@[ FSTTestDocKey(@"foo/bar"), FSTTestDocKey(@"foo/bar2") ]));
+  XCTAssertEqual(garbage,
+                 std::set<DocumentKey>({testutil::Key("foo/bar"), testutil::Key("foo/bar2")}));
 
   [batches addObject:[self addMutationBatchWithKey:@"foo/bar/suffix/baz"]];
   garbage = [garbageCollector collectGarbage];
-  FSTAssertEqualSets(garbage, @[]);
+  XCTAssertEqual(garbage, std::set<DocumentKey>({}));
 
   [self removeMutationBatches:@[ batches[4], batches[6] ]];
   garbage = [garbageCollector collectGarbage];
-  FSTAssertEqualSets(garbage, @[ FSTTestDocKey(@"foo/bar/suffix/baz") ]);
+  XCTAssertEqual(garbage, std::set<DocumentKey>({testutil::Key("foo/bar/suffix/baz")}));
 }
 
 - (void)testStreamToken {
diff --git a/Firestore/Example/Tests/Local/FSTQueryCacheTests.mm b/Firestore/Example/Tests/Local/FSTQueryCacheTests.mm
index 6ef927d..5dd4b5d 100644
--- a/Firestore/Example/Tests/Local/FSTQueryCacheTests.mm
+++ b/Firestore/Example/Tests/Local/FSTQueryCacheTests.mm
@@ -22,11 +22,16 @@
 #import "Firestore/Source/Local/FSTPersistence.h"
 #import "Firestore/Source/Local/FSTQueryData.h"
 #import "Firestore/Source/Local/FSTWriteGroup.h"
-#import "Firestore/Source/Model/FSTDocumentKey.h"
 
 #import "Firestore/Example/Tests/Util/FSTHelpers.h"
 #import "Firestore/third_party/Immutable/Tests/FSTImmutableSortedSet+Testing.h"
 
+#include "Firestore/core/src/firebase/firestore/model/document_key.h"
+#include "Firestore/core/test/firebase/firestore/testutil/testutil.h"
+
+namespace testutil = firebase::firestore::testutil;
+using firebase::firestore::model::DocumentKey;
+
 NS_ASSUME_NONNULL_BEGIN
 
 @implementation FSTQueryCacheTests {
@@ -164,8 +169,8 @@ NS_ASSUME_NONNULL_BEGIN
   FSTQueryData *rooms = [self queryDataWithQuery:_queryRooms];
   [self.queryCache addQueryData:rooms group:group];
 
-  FSTDocumentKey *key1 = FSTTestDocKey(@"rooms/foo");
-  FSTDocumentKey *key2 = FSTTestDocKey(@"rooms/bar");
+  DocumentKey key1 = testutil::Key("rooms/foo");
+  DocumentKey key2 = testutil::Key("rooms/bar");
   [self addMatchingKey:key1 forTargetID:rooms.targetID group:group];
   [self addMatchingKey:key2 forTargetID:rooms.targetID group:group];
 
@@ -182,7 +187,7 @@ NS_ASSUME_NONNULL_BEGIN
   if ([self isTestBaseClass]) return;
 
   FSTWriteGroup *group = [self.persistence startGroupWithAction:@"AddOrRemoveMatchingKeys"];
-  FSTDocumentKey *key = FSTTestDocKey(@"foo/bar");
+  DocumentKey key = testutil::Key("foo/bar");
 
   XCTAssertFalse([self.queryCache containsKey:key]);
 
@@ -204,9 +209,9 @@ NS_ASSUME_NONNULL_BEGIN
   if ([self isTestBaseClass]) return;
 
   FSTWriteGroup *group = [self.persistence startGroupWithAction:@"RemoveMatchingKeysForTargetID"];
-  FSTDocumentKey *key1 = FSTTestDocKey(@"foo/bar");
-  FSTDocumentKey *key2 = FSTTestDocKey(@"foo/baz");
-  FSTDocumentKey *key3 = FSTTestDocKey(@"foo/blah");
+  DocumentKey key1 = testutil::Key("foo/bar");
+  DocumentKey key2 = testutil::Key("foo/baz");
+  DocumentKey key3 = testutil::Key("foo/blah");
 
   [self addMatchingKey:key1 forTargetID:1 group:group];
   [self addMatchingKey:key2 forTargetID:1 group:group];
@@ -233,32 +238,32 @@ NS_ASSUME_NONNULL_BEGIN
   FSTWriteGroup *group = [self.persistence startGroupWithAction:@"RemoveEmitsGarbageEvents"];
   FSTEagerGarbageCollector *garbageCollector = [[FSTEagerGarbageCollector alloc] init];
   [garbageCollector addGarbageSource:self.queryCache];
-  FSTAssertEqualSets([garbageCollector collectGarbage], @[]);
+  XCTAssertEqual([garbageCollector collectGarbage], std::set<DocumentKey>({}));
 
   FSTQueryData *rooms = [self queryDataWithQuery:FSTTestQuery("rooms")];
-  FSTDocumentKey *room1 = FSTTestDocKey(@"rooms/bar");
-  FSTDocumentKey *room2 = FSTTestDocKey(@"rooms/foo");
+  DocumentKey room1 = testutil::Key("rooms/bar");
+  DocumentKey room2 = testutil::Key("rooms/foo");
   [self.queryCache addQueryData:rooms group:group];
   [self addMatchingKey:room1 forTargetID:rooms.targetID group:group];
   [self addMatchingKey:room2 forTargetID:rooms.targetID group:group];
 
   FSTQueryData *halls = [self queryDataWithQuery:FSTTestQuery("halls")];
-  FSTDocumentKey *hall1 = FSTTestDocKey(@"halls/bar");
-  FSTDocumentKey *hall2 = FSTTestDocKey(@"halls/foo");
+  DocumentKey hall1 = testutil::Key("halls/bar");
+  DocumentKey hall2 = testutil::Key("halls/foo");
   [self.queryCache addQueryData:halls group:group];
   [self addMatchingKey:hall1 forTargetID:halls.targetID group:group];
   [self addMatchingKey:hall2 forTargetID:halls.targetID group:group];
 
-  FSTAssertEqualSets([garbageCollector collectGarbage], @[]);
+  XCTAssertEqual([garbageCollector collectGarbage], std::set<DocumentKey>({}));
 
   [self removeMatchingKey:room1 forTargetID:rooms.targetID group:group];
-  FSTAssertEqualSets([garbageCollector collectGarbage], @[ room1 ]);
+  XCTAssertEqual([garbageCollector collectGarbage], std::set<DocumentKey>({room1}));
 
   [self.queryCache removeQueryData:rooms group:group];
-  FSTAssertEqualSets([garbageCollector collectGarbage], @[ room2 ]);
+  XCTAssertEqual([garbageCollector collectGarbage], std::set<DocumentKey>({room2}));
 
   [self.queryCache removeMatchingKeysForTargetID:halls.targetID group:group];
-  FSTAssertEqualSets([garbageCollector collectGarbage], (@[ hall1, hall2 ]));
+  XCTAssertEqual([garbageCollector collectGarbage], std::set<DocumentKey>({hall1, hall2}));
   [self.persistence commitGroup:group];
 }
 
@@ -266,9 +271,9 @@ NS_ASSUME_NONNULL_BEGIN
   if ([self isTestBaseClass]) return;
 
   FSTWriteGroup *group = [self.persistence startGroupWithAction:@"MatchingKeysForTargetID"];
-  FSTDocumentKey *key1 = FSTTestDocKey(@"foo/bar");
-  FSTDocumentKey *key2 = FSTTestDocKey(@"foo/baz");
-  FSTDocumentKey *key3 = FSTTestDocKey(@"foo/blah");
+  DocumentKey key1 = testutil::Key("foo/bar");
+  DocumentKey key2 = testutil::Key("foo/baz");
+  DocumentKey key3 = testutil::Key("foo/blah");
 
   [self addMatchingKey:key1 forTargetID:1 group:group];
   [self addMatchingKey:key2 forTargetID:1 group:group];
@@ -335,8 +340,8 @@ NS_ASSUME_NONNULL_BEGIN
                                                     targetID:1
                                         listenSequenceNumber:10
                                                      purpose:FSTQueryPurposeListen];
-  FSTDocumentKey *key1 = FSTTestDocKey(@"rooms/bar");
-  FSTDocumentKey *key2 = FSTTestDocKey(@"rooms/foo");
+  DocumentKey key1 = testutil::Key("rooms/bar");
+  DocumentKey key2 = testutil::Key("rooms/foo");
   [self.queryCache addQueryData:query1 group:group];
   [self addMatchingKey:key1 forTargetID:1 group:group];
   [self addMatchingKey:key2 forTargetID:1 group:group];
@@ -345,7 +350,7 @@ NS_ASSUME_NONNULL_BEGIN
                                                     targetID:2
                                         listenSequenceNumber:20
                                                      purpose:FSTQueryPurposeListen];
-  FSTDocumentKey *key3 = FSTTestDocKey(@"halls/foo");
+  DocumentKey key3 = testutil::Key("halls/foo");
   [self.queryCache addQueryData:query2 group:group];
   [self addMatchingKey:key3 forTargetID:2 group:group];
   XCTAssertEqual([self.queryCache highestTargetID], 2);
@@ -419,7 +424,7 @@ NS_ASSUME_NONNULL_BEGIN
                                  resumeToken:resumeToken];
 }
 
-- (void)addMatchingKey:(FSTDocumentKey *)key
+- (void)addMatchingKey:(cons DocumentKey &)key
            forTargetID:(FSTTargetID)targetID
                  group:(FSTWriteGroup *)group {
   FSTDocumentKeySet *keys = [FSTDocumentKeySet keySet];
@@ -427,7 +432,7 @@ NS_ASSUME_NONNULL_BEGIN
   [self.queryCache addMatchingKeys:keys forTargetID:targetID group:group];
 }
 
-- (void)removeMatchingKey:(FSTDocumentKey *)key
+- (void)removeMatchingKey:(const DocumentKey &)key
               forTargetID:(FSTTargetID)targetID
                     group:(FSTWriteGroup *)group {
   FSTDocumentKeySet *keys = [FSTDocumentKeySet keySet];
diff --git a/Firestore/Source/Core/FSTSyncEngine.mm b/Firestore/Source/Core/FSTSyncEngine.mm
index 61fac7d..d567878 100644
--- a/Firestore/Source/Core/FSTSyncEngine.mm
+++ b/Firestore/Source/Core/FSTSyncEngine.mm
@@ -43,10 +43,12 @@
 
 #include "Firestore/core/src/firebase/firestore/auth/user.h"
 #include "Firestore/core/src/firebase/firestore/core/target_id_generator.h"
+#include "Firestore/core/src/firebase/firestore/model/document_key.h"
 
 using firebase::firestore::auth::HashUser;
 using firebase::firestore::auth::User;
 using firebase::firestore::core::TargetIdGenerator;
+using firebase::firestore::model::DocumentKey;
 
 NS_ASSUME_NONNULL_BEGIN
 
@@ -509,9 +511,9 @@ static const FSTListenSequenceNumber kIrrelevantSequenceNumber = -1;
 
 /** Garbage collect the limbo documents that we no longer need to track. */
 - (void)garbageCollectLimboDocuments {
-  NSSet<FSTDocumentKey *> *garbage = [self.limboCollector collectGarbage];
-  for (FSTDocumentKey *key in garbage) {
-    FSTBoxedTargetID *limboTarget = self.limboTargetsByKey[key];
+  const std::set<DocumentKey> garbage = [self.limboCollector collectGarbage];
+  for (const DocumentKey &key : garbage) {
+    FSTBoxedTargetID *limboTarget = self.limboTargetsByKey[static_cast<FSTDocumentKey *>(key)];
     if (!limboTarget) {
       // This target already got removed, because the query failed.
       return;
diff --git a/Firestore/Source/Local/FSTDocumentReference.h b/Firestore/Source/Local/FSTDocumentReference.h
index 04b8416..e23905c 100644
--- a/Firestore/Source/Local/FSTDocumentReference.h
+++ b/Firestore/Source/Local/FSTDocumentReference.h
@@ -16,7 +16,7 @@
 
 #import <Foundation/Foundation.h>
 
-@class FSTDocumentKey;
+#include "Firestore/core/src/firebase/firestore/model/document_key.h"
 
 NS_ASSUME_NONNULL_BEGIN
 
@@ -32,12 +32,13 @@ NS_ASSUME_NONNULL_BEGIN
 @interface FSTDocumentReference : NSObject <NSCopying>
 
 /** Initializes the document reference with the given key and ID. */
-- (instancetype)initWithKey:(FSTDocumentKey *)key ID:(int32_t)ID NS_DESIGNATED_INITIALIZER;
+- (instancetype)initWithKey:(firebase::firestore::model::DocumentKey)key
+                         ID:(int32_t)ID NS_DESIGNATED_INITIALIZER;
 
 - (instancetype)init NS_UNAVAILABLE;
 
 /** The document key that's the target of this reference. */
-@property(nonatomic, strong, readonly) FSTDocumentKey *key;
+- (const firebase::firestore::model::DocumentKey &)key;
 
 /**
  * The targetID of a referring target or the batchID of a referring mutation batch. (Which this
diff --git a/Firestore/Source/Local/FSTDocumentReference.mm b/Firestore/Source/Local/FSTDocumentReference.mm
index 4310baa..3e755bc 100644
--- a/Firestore/Source/Local/FSTDocumentReference.mm
+++ b/Firestore/Source/Local/FSTDocumentReference.mm
@@ -16,20 +16,24 @@
 
 #import "Firestore/Source/Local/FSTDocumentReference.h"
 
-#include "Firestore/core/src/firebase/firestore/util/comparison.h"
+#include <utility>
 
-#import "Firestore/Source/Model/FSTDocumentKey.h"
+#include "Firestore/core/src/firebase/firestore/model/document_key.h"
+#include "Firestore/core/src/firebase/firestore/util/comparison.h"
 
+using firebase::firestore::model::DocumentKey;
 using firebase::firestore::util::WrapCompare;
 
 NS_ASSUME_NONNULL_BEGIN
 
-@implementation FSTDocumentReference
+@implementation FSTDocumentReference {
+  DocumentKey _key;
+}
 
-- (instancetype)initWithKey:(FSTDocumentKey *)key ID:(int32_t)ID {
+- (instancetype)initWithKey:(DocumentKey)key ID:(int32_t)ID {
   self = [super init];
   if (self) {
-    _key = key;
+    _key = std::move(key);
     _ID = ID;
   }
   return self;
@@ -51,7 +55,8 @@ NS_ASSUME_NONNULL_BEGIN
 }
 
 - (NSString *)description {
-  return [NSString stringWithFormat:@"<FSTDocumentReference: key=%@, ID=%d>", self.key, self.ID];
+  return [NSString stringWithFormat:@"<FSTDocumentReference: key=%s, ID=%d>",
+                                    self.key.ToString().c_str(), self.ID];
 }
 
 - (id)copyWithZone:(nullable NSZone *)zone {
@@ -59,6 +64,10 @@ NS_ASSUME_NONNULL_BEGIN
   return self;
 }
 
+- (const firebase::firestore::model::DocumentKey &)key {
+  return _key;
+}
+
 @end
 
 #pragma mark Comparators
diff --git a/Firestore/Source/Local/FSTEagerGarbageCollector.h b/Firestore/Source/Local/FSTEagerGarbageCollector.h
index 40815b7..72b7375 100644
--- a/Firestore/Source/Local/FSTEagerGarbageCollector.h
+++ b/Firestore/Source/Local/FSTEagerGarbageCollector.h
@@ -18,8 +18,6 @@
 
 #import "Firestore/Source/Local/FSTGarbageCollector.h"
 
-@class FSTDocumentKey;
-
 NS_ASSUME_NONNULL_BEGIN
 
 /**
diff --git a/Firestore/Source/Local/FSTEagerGarbageCollector.mm b/Firestore/Source/Local/FSTEagerGarbageCollector.mm
index 77a577e..f7c75dd 100644
--- a/Firestore/Source/Local/FSTEagerGarbageCollector.mm
+++ b/Firestore/Source/Local/FSTEagerGarbageCollector.mm
@@ -16,7 +16,11 @@
 
 #import "Firestore/Source/Local/FSTEagerGarbageCollector.h"
 
-#import "Firestore/Source/Model/FSTDocumentKey.h"
+#include <set>
+
+#include "Firestore/core/src/firebase/firestore/model/document_key.h"
+
+using firebase::firestore::model::DocumentKey;
 
 NS_ASSUME_NONNULL_BEGIN
 
@@ -27,18 +31,17 @@ NS_ASSUME_NONNULL_BEGIN
 /** The garbage collectible sources to double-check during garbage collection. */
 @property(nonatomic, strong, readonly) NSMutableArray<id<FSTGarbageSource>> *sources;
 
-/** A set of potentially garbage keys. */
-@property(nonatomic, strong, readonly) NSMutableSet<FSTDocumentKey *> *potentialGarbage;
-
 @end
 
-@implementation FSTEagerGarbageCollector
+@implementation FSTEagerGarbageCollector {
+  /** A set of potentially garbage keys. */
+  std::set<DocumentKey> _potentialGarbage;
+}
 
 - (instancetype)init {
   self = [super init];
   if (self) {
     _sources = [NSMutableArray array];
-    _potentialGarbage = [[NSMutableSet alloc] init];
   }
   return self;
 }
@@ -57,15 +60,15 @@ NS_ASSUME_NONNULL_BEGIN
   garbageSource.garbageCollector = nil;
 }
 
-- (void)addPotentialGarbageKey:(FSTDocumentKey *)key {
-  [self.potentialGarbage addObject:key];
+- (void)addPotentialGarbageKey:(const DocumentKey &)key {
+  _potentialGarbage.insert(key);
 }
 
-- (NSMutableSet<FSTDocumentKey *> *)collectGarbage {
+- (std::set<DocumentKey>)collectGarbage {
   NSMutableArray<id<FSTGarbageSource>> *sources = self.sources;
 
-  NSMutableSet<FSTDocumentKey *> *actualGarbage = [NSMutableSet set];
-  for (FSTDocumentKey *key in self.potentialGarbage) {
+  std::set<DocumentKey> actualGarbage;
+  for (const DocumentKey &key : _potentialGarbage) {
     BOOL isGarbage = YES;
     for (id<FSTGarbageSource> source in sources) {
       if ([source containsKey:key]) {
@@ -75,12 +78,12 @@ NS_ASSUME_NONNULL_BEGIN
     }
 
     if (isGarbage) {
-      [actualGarbage addObject:key];
+      actualGarbage.insert(key);
     }
   }
 
   // Clear locally retained potential keys and returned confirmed garbage.
-  [self.potentialGarbage removeAllObjects];
+  _potentialGarbage.clear();
   return actualGarbage;
 }
 
diff --git a/Firestore/Source/Local/FSTGarbageCollector.h b/Firestore/Source/Local/FSTGarbageCollector.h
index ff5116b..e69e2b2 100644
--- a/Firestore/Source/Local/FSTGarbageCollector.h
+++ b/Firestore/Source/Local/FSTGarbageCollector.h
@@ -16,9 +16,12 @@
 
 #import <Foundation/Foundation.h>
 
+#include <set>
+
 #import "Firestore/Source/Core/FSTTypes.h"
 
-@class FSTDocumentKey;
+#include "Firestore/core/src/firebase/firestore/model/document_key.h"
+
 @class FSTDocumentReference;
 @protocol FSTGarbageCollector;
 
@@ -41,7 +44,7 @@ NS_ASSUME_NONNULL_BEGIN
  * garbage collectors to double-check if a key exists in this collection when it was released
  * elsewhere.
  */
-- (BOOL)containsKey:(FSTDocumentKey *)key;
+- (BOOL)containsKey:(const firebase::firestore::model::DocumentKey&)key;
 
 @end
 
@@ -85,10 +88,10 @@ NS_ASSUME_NONNULL_BEGIN
  * matches any active targets. This behavior allows the client to avoid re-showing an old document
  * in the next latency-compensated view.
  */
-- (void)addPotentialGarbageKey:(FSTDocumentKey *)key;
+- (void)addPotentialGarbageKey:(const firebase::firestore::model::DocumentKey&)key;
 
 /** Returns the contents of the garbage bin and clears it. */
-- (NSSet<FSTDocumentKey *> *)collectGarbage;
+- (std::set<firebase::firestore::model::DocumentKey>)collectGarbage;
 
 @end
 
diff --git a/Firestore/Source/Local/FSTLevelDBMutationQueue.mm b/Firestore/Source/Local/FSTLevelDBMutationQueue.mm
index ac93ab1..221d56f 100644
--- a/Firestore/Source/Local/FSTLevelDBMutationQueue.mm
+++ b/Firestore/Source/Local/FSTLevelDBMutationQueue.mm
@@ -27,12 +27,12 @@
 #import "Firestore/Source/Local/FSTLevelDBKey.h"
 #import "Firestore/Source/Local/FSTLocalSerializer.h"
 #import "Firestore/Source/Local/FSTWriteGroup.h"
-#import "Firestore/Source/Model/FSTDocumentKey.h"
 #import "Firestore/Source/Model/FSTMutation.h"
 #import "Firestore/Source/Model/FSTMutationBatch.h"
 #import "Firestore/Source/Util/FSTAssert.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/resource_path.h"
 #include "Firestore/core/src/firebase/firestore/util/string_apple.h"
 #include "Firestore/core/src/firebase/firestore/util/string_util.h"
@@ -42,6 +42,7 @@ NS_ASSUME_NONNULL_BEGIN
 namespace util = firebase::firestore::util;
 using Firestore::StringView;
 using firebase::firestore::auth::User;
+using firebase::firestore::model::DocumentKey;
 using firebase::firestore::model::ResourcePath;
 using leveldb::DB;
 using leveldb::Iterator;
@@ -377,12 +378,12 @@ static ReadOptions StandardReadOptions() {
 }
 
 - (NSArray<FSTMutationBatch *> *)allMutationBatchesAffectingDocumentKey:
-    (FSTDocumentKey *)documentKey {
+    (const DocumentKey &)documentKey {
   NSString *userID = self.userID;
 
   // Scan the document-mutation index starting with a prefix starting with the given documentKey.
-  std::string indexPrefix =
-      [FSTLevelDBDocumentMutationKey keyPrefixWithUserID:self.userID resourcePath:documentKey.path];
+  std::string indexPrefix = [FSTLevelDBDocumentMutationKey keyPrefixWithUserID:self.userID
+                                                                  resourcePath:documentKey.path()];
   std::unique_ptr<Iterator> indexIterator(_db->NewIterator(StandardReadOptions()));
   indexIterator->Seek(indexPrefix);
 
@@ -403,7 +404,7 @@ static ReadOptions StandardReadOptions() {
     // occur before any rows for documents nested in a subcollection beneath documentKey so we can
     // stop as soon as we hit any such row.
     if (!indexKey.starts_with(indexPrefix) || ![rowKey decodeKey:indexKey] ||
-        ![rowKey.documentKey isEqualToKey:documentKey]) {
+        DocumentKey{rowKey.documentKey} != documentKey) {
       break;
     }
 
@@ -619,9 +620,9 @@ static ReadOptions StandardReadOptions() {
 
 #pragma mark - FSTGarbageSource implementation
 
-- (BOOL)containsKey:(FSTDocumentKey *)documentKey {
-  std::string indexPrefix =
-      [FSTLevelDBDocumentMutationKey keyPrefixWithUserID:self.userID resourcePath:documentKey.path];
+- (BOOL)containsKey:(const DocumentKey &)documentKey {
+  std::string indexPrefix = [FSTLevelDBDocumentMutationKey keyPrefixWithUserID:self.userID
+                                                                  resourcePath:documentKey.path()];
   std::unique_ptr<Iterator> indexIterator(_db->NewIterator(StandardReadOptions()));
   indexIterator->Seek(indexPrefix);
 
@@ -632,7 +633,7 @@ static ReadOptions StandardReadOptions() {
     // Check both that the key prefix matches and that the decoded document key is exactly the key
     // we're looking for.
     if (iteratorKey.starts_with(indexPrefix) && [rowKey decodeKey:iteratorKey] &&
-        [rowKey.documentKey isEqualToKey:documentKey]) {
+        DocumentKey{rowKey.documentKey} == documentKey) {
       return YES;
     }
   }
diff --git a/Firestore/Source/Local/FSTLevelDBQueryCache.mm b/Firestore/Source/Local/FSTLevelDBQueryCache.mm
index c7eecc4..09440e9 100644
--- a/Firestore/Source/Local/FSTLevelDBQueryCache.mm
+++ b/Firestore/Source/Local/FSTLevelDBQueryCache.mm
@@ -23,14 +23,16 @@
 #import "Firestore/Source/Local/FSTLocalSerializer.h"
 #import "Firestore/Source/Local/FSTQueryData.h"
 #import "Firestore/Source/Local/FSTWriteGroup.h"
-#import "Firestore/Source/Model/FSTDocumentKey.h"
 #import "Firestore/Source/Util/FSTAssert.h"
 #include "absl/strings/match.h"
 
+#include "Firestore/core/src/firebase/firestore/model/document_key.h"
+
 NS_ASSUME_NONNULL_BEGIN
 
 using firebase::firestore::local::LevelDbTransaction;
 using Firestore::StringView;
+using firebase::firestore::model::DocumentKey;
 using leveldb::DB;
 using leveldb::Slice;
 using leveldb::Status;
@@ -323,7 +325,7 @@ using leveldb::Status;
     if (![rowKey decodeKey:indexKey] || rowKey.targetID != targetID) {
       break;
     }
-    FSTDocumentKey *documentKey = rowKey.documentKey;
+    const DocumentKey &documentKey = rowKey.documentKey;
 
     // Delete both index rows
     [group removeMessageForKey:indexKey];
@@ -356,14 +358,14 @@ using leveldb::Status;
 
 #pragma mark - FSTGarbageSource implementation
 
-- (BOOL)containsKey:(FSTDocumentKey *)key {
-  std::string indexPrefix = [FSTLevelDBDocumentTargetKey keyPrefixWithResourcePath:key.path];
+- (BOOL)containsKey:(const DocumentKey &)key {
+  std::string indexPrefix = [FSTLevelDBDocumentTargetKey keyPrefixWithResourcePath:key.path()];
   auto indexIterator = _db.currentTransaction->NewIterator();
   indexIterator->Seek(indexPrefix);
 
   if (indexIterator->Valid()) {
     FSTLevelDBDocumentTargetKey *rowKey = [[FSTLevelDBDocumentTargetKey alloc] init];
-    if ([rowKey decodeKey:indexIterator->key()] && [rowKey.documentKey isEqualToKey:key]) {
+    if ([rowKey decodeKey:indexIterator->key()] && DocumentKey{rowKey.documentKey} == key) {
       return YES;
     }
   }
diff --git a/Firestore/Source/Local/FSTLevelDBRemoteDocumentCache.mm b/Firestore/Source/Local/FSTLevelDBRemoteDocumentCache.mm
index 97f0ef1..cd14390 100644
--- a/Firestore/Source/Local/FSTLevelDBRemoteDocumentCache.mm
+++ b/Firestore/Source/Local/FSTLevelDBRemoteDocumentCache.mm
@@ -27,12 +27,14 @@
 #import "Firestore/Source/Local/FSTWriteGroup.h"
 #import "Firestore/Source/Model/FSTDocument.h"
 #import "Firestore/Source/Model/FSTDocumentDictionary.h"
-#import "Firestore/Source/Model/FSTDocumentKey.h"
 #import "Firestore/Source/Model/FSTDocumentSet.h"
 #import "Firestore/Source/Util/FSTAssert.h"
 
+#include "Firestore/core/src/firebase/firestore/model/document_key.h"
+
 NS_ASSUME_NONNULL_BEGIN
 
+using firebase::firestore::model::DocumentKey;
 using leveldb::DB;
 using leveldb::Iterator;
 using leveldb::ReadOptions;
@@ -79,12 +81,12 @@ static ReadOptions StandardReadOptions() {
   [group setMessage:[self.serializer encodedMaybeDocument:document] forKey:key];
 }
 
-- (void)removeEntryForKey:(FSTDocumentKey *)documentKey group:(FSTWriteGroup *)group {
+- (void)removeEntryForKey:(const DocumentKey &)documentKey group:(FSTWriteGroup *)group {
   std::string key = [self remoteDocumentKey:documentKey];
   [group removeMessageForKey:key];
 }
 
-- (nullable FSTMaybeDocument *)entryForKey:(FSTDocumentKey *)documentKey {
+- (nullable FSTMaybeDocument *)entryForKey:(const DocumentKey &)documentKey {
   std::string key = [FSTLevelDBRemoteDocumentKey keyWithDocumentKey:documentKey];
   std::string value;
   Status status = _db->Get(StandardReadOptions(), key, &value);
@@ -93,7 +95,7 @@ static ReadOptions StandardReadOptions() {
   } else if (status.ok()) {
     return [self decodedMaybeDocument:value withKey:documentKey];
   } else {
-    FSTFail(@"Fetch document for key (%@) failed with status: %s", documentKey,
+    FSTFail(@"Fetch document for key (%s) failed with status: %s", documentKey.ToString().c_str(),
             status.ToString().c_str());
   }
 }
@@ -127,11 +129,11 @@ static ReadOptions StandardReadOptions() {
   return results;
 }
 
-- (std::string)remoteDocumentKey:(FSTDocumentKey *)key {
+- (std::string)remoteDocumentKey:(const DocumentKey &)key {
   return [FSTLevelDBRemoteDocumentKey keyWithDocumentKey:key];
 }
 
-- (FSTMaybeDocument *)decodedMaybeDocument:(Slice)slice withKey:(FSTDocumentKey *)documentKey {
+- (FSTMaybeDocument *)decodedMaybeDocument:(Slice)slice withKey:(const DocumentKey &)documentKey {
   NSData *data =
       [[NSData alloc] initWithBytesNoCopy:(void *)slice.data() length:slice.size() freeWhenDone:NO];
 
@@ -143,8 +145,8 @@ static ReadOptions StandardReadOptions() {
 
   FSTMaybeDocument *maybeDocument = [self.serializer decodedMaybeDocument:proto];
   FSTAssert([maybeDocument.key isEqualToKey:documentKey],
-            @"Read document has key (%s) instead of expected key (%@).",
-            maybeDocument.key.ToString().c_str(), documentKey);
+            @"Read document has key (%s) instead of expected key (%s).",
+            maybeDocument.key.ToString().c_str(), documentKey.ToString().c_str());
   return maybeDocument;
 }
 
diff --git a/Firestore/Source/Local/FSTLocalDocumentsView.h b/Firestore/Source/Local/FSTLocalDocumentsView.h
index 46830e7..e75e0f3 100644
--- a/Firestore/Source/Local/FSTLocalDocumentsView.h
+++ b/Firestore/Source/Local/FSTLocalDocumentsView.h
@@ -19,7 +19,8 @@
 #import "Firestore/Source/Model/FSTDocumentDictionary.h"
 #import "Firestore/Source/Model/FSTDocumentKeySet.h"
 
-@class FSTDocumentKey;
+#include "Firestore/core/src/firebase/firestore/model/document_key.h"
+
 @class FSTMaybeDocument;
 @class FSTQuery;
 @protocol FSTMutationQueue;
@@ -44,7 +45,7 @@ NS_ASSUME_NONNULL_BEGIN
  *
  * @return Local view of the document or nil if we don't have any cached state for it.
  */
-- (nullable FSTMaybeDocument *)documentForKey:(FSTDocumentKey *)key;
+- (nullable FSTMaybeDocument *)documentForKey:(const firebase::firestore::model::DocumentKey &)key;
 
 /**
  * Gets the local view of the documents identified by `keys`.
diff --git a/Firestore/Source/Local/FSTLocalDocumentsView.mm b/Firestore/Source/Local/FSTLocalDocumentsView.mm
index c059146..e9b9423 100644
--- a/Firestore/Source/Local/FSTLocalDocumentsView.mm
+++ b/Firestore/Source/Local/FSTLocalDocumentsView.mm
@@ -22,13 +22,14 @@
 #import "Firestore/Source/Local/FSTRemoteDocumentCache.h"
 #import "Firestore/Source/Model/FSTDocument.h"
 #import "Firestore/Source/Model/FSTDocumentDictionary.h"
-#import "Firestore/Source/Model/FSTDocumentKey.h"
 #import "Firestore/Source/Model/FSTMutation.h"
 #import "Firestore/Source/Model/FSTMutationBatch.h"
 #import "Firestore/Source/Util/FSTAssert.h"
 
+#include "Firestore/core/src/firebase/firestore/model/document_key.h"
 #include "Firestore/core/src/firebase/firestore/model/resource_path.h"
 
+using firebase::firestore::model::DocumentKey;
 using firebase::firestore::model::ResourcePath;
 
 NS_ASSUME_NONNULL_BEGIN
@@ -58,7 +59,7 @@ NS_ASSUME_NONNULL_BEGIN
   return self;
 }
 
-- (nullable FSTMaybeDocument *)documentForKey:(FSTDocumentKey *)key {
+- (nullable FSTMaybeDocument *)documentForKey:(const DocumentKey &)key {
   FSTMaybeDocument *_Nullable remoteDoc = [self.remoteDocumentCache entryForKey:key];
   return [self localDocument:remoteDoc key:key];
 }
@@ -78,7 +79,7 @@ NS_ASSUME_NONNULL_BEGIN
 }
 
 - (FSTDocumentDictionary *)documentsMatchingQuery:(FSTQuery *)query {
-  if ([FSTDocumentKey isDocumentKey:query.path]) {
+  if (DocumentKey::IsDocumentKey(query.path)) {
     return [self documentsMatchingDocumentQuery:query.path];
   } else {
     return [self documentsMatchingCollectionQuery:query];
@@ -88,7 +89,7 @@ NS_ASSUME_NONNULL_BEGIN
 - (FSTDocumentDictionary *)documentsMatchingDocumentQuery:(const ResourcePath &)docPath {
   FSTDocumentDictionary *result = [FSTDocumentDictionary documentDictionary];
   // Just do a simple document lookup.
-  FSTMaybeDocument *doc = [self documentForKey:[FSTDocumentKey keyWithPath:docPath]];
+  FSTMaybeDocument *doc = [self documentForKey:DocumentKey{docPath}];
   if ([doc isKindOfClass:[FSTDocument class]]) {
     result = [result dictionaryBySettingObject:(FSTDocument *)doc forKey:doc.key];
   }
@@ -148,7 +149,7 @@ NS_ASSUME_NONNULL_BEGIN
  * @param documentKey The key of the document (necessary when remoteDocument is nil).
  */
 - (nullable FSTMaybeDocument *)localDocument:(nullable FSTMaybeDocument *)document
-                                         key:(FSTDocumentKey *)documentKey {
+                                         key:(const DocumentKey &)documentKey {
   NSArray<FSTMutationBatch *> *batches =
       [self.mutationQueue allMutationBatchesAffectingDocumentKey:documentKey];
   for (FSTMutationBatch *batch in batches) {
diff --git a/Firestore/Source/Local/FSTLocalSerializer.mm b/Firestore/Source/Local/FSTLocalSerializer.mm
index ec70ca0..1d9455d 100644
--- a/Firestore/Source/Local/FSTLocalSerializer.mm
+++ b/Firestore/Source/Local/FSTLocalSerializer.mm
@@ -30,6 +30,10 @@
 #import "Firestore/Source/Remote/FSTSerializerBeta.h"
 #import "Firestore/Source/Util/FSTAssert.h"
 
+#include "Firestore/core/src/firebase/firestore/model/document_key.h"
+
+using firebase::firestore::model::DocumentKey;
+
 @interface FSTLocalSerializer ()
 
 @property(nonatomic, strong, readonly) FSTSerializerBeta *remoteSerializer;
@@ -95,7 +99,7 @@
   FSTSerializerBeta *remoteSerializer = self.remoteSerializer;
 
   FSTObjectValue *data = [remoteSerializer decodedFields:document.fields];
-  FSTDocumentKey *key = [remoteSerializer decodedDocumentKey:document.name];
+  const DocumentKey key = [remoteSerializer decodedDocumentKey:document.name];
   FSTSnapshotVersion *version = [remoteSerializer decodedVersion:document.updateTime];
   return [FSTDocument documentWithData:data key:key version:version hasLocalMutations:NO];
 }
@@ -114,7 +118,7 @@
 - (FSTDeletedDocument *)decodedDeletedDocument:(FSTPBNoDocument *)proto {
   FSTSerializerBeta *remoteSerializer = self.remoteSerializer;
 
-  FSTDocumentKey *key = [remoteSerializer decodedDocumentKey:proto.name];
+  const DocumentKey key = [remoteSerializer decodedDocumentKey:proto.name];
   FSTSnapshotVersion *version = [remoteSerializer decodedVersion:proto.readTime];
   return [FSTDeletedDocument documentWithKey:key version:version];
 }
diff --git a/Firestore/Source/Local/FSTLocalStore.h b/Firestore/Source/Local/FSTLocalStore.h
index 4ec23fd..e2134dc 100644
--- a/Firestore/Source/Local/FSTLocalStore.h
+++ b/Firestore/Source/Local/FSTLocalStore.h
@@ -22,6 +22,7 @@
 #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"
 
 @class FSTLocalViewChanges;
 @class FSTLocalWriteResult;
@@ -104,7 +105,7 @@ NS_ASSUME_NONNULL_BEGIN
 - (FSTLocalWriteResult *)locallyWriteMutations:(NSArray<FSTMutation *> *)mutations;
 
 /** Returns the current value of a document with a given key, or nil if not found. */
-- (nullable FSTMaybeDocument *)readDocument:(FSTDocumentKey *)key;
+- (nullable FSTMaybeDocument *)readDocument:(const firebase::firestore::model::DocumentKey &)key;
 
 /**
  * Acknowledges the given batch.
diff --git a/Firestore/Source/Local/FSTLocalStore.mm b/Firestore/Source/Local/FSTLocalStore.mm
index f61a20f..922f281 100644
--- a/Firestore/Source/Local/FSTLocalStore.mm
+++ b/Firestore/Source/Local/FSTLocalStore.mm
@@ -34,7 +34,6 @@
 #import "Firestore/Source/Local/FSTWriteGroup.h"
 #import "Firestore/Source/Model/FSTDocument.h"
 #import "Firestore/Source/Model/FSTDocumentDictionary.h"
-#import "Firestore/Source/Model/FSTDocumentKey.h"
 #import "Firestore/Source/Model/FSTMutation.h"
 #import "Firestore/Source/Model/FSTMutationBatch.h"
 #import "Firestore/Source/Remote/FSTRemoteEvent.h"
@@ -43,8 +42,10 @@
 
 #include "Firestore/core/src/firebase/firestore/auth/user.h"
 #include "Firestore/core/src/firebase/firestore/core/target_id_generator.h"
+#include "Firestore/core/src/firebase/firestore/model/document_key.h"
 
 using firebase::firestore::auth::User;
+using firebase::firestore::model::DocumentKey;
 using firebase::firestore::core::TargetIdGenerator;
 
 NS_ASSUME_NONNULL_BEGIN
@@ -387,7 +388,7 @@ NS_ASSUME_NONNULL_BEGIN
   return [self.mutationQueue nextMutationBatchAfterBatchID:batchID];
 }
 
-- (nullable FSTMaybeDocument *)readDocument:(FSTDocumentKey *)key {
+- (nullable FSTMaybeDocument *)readDocument:(const DocumentKey &)key {
   return [self.localDocuments documentForKey:key];
 }
 
@@ -458,9 +459,9 @@ NS_ASSUME_NONNULL_BEGIN
   FSTWriteGroup *group = [self.persistence startGroupWithAction:@"Garbage Collection"];
   // Call collectGarbage regardless of whether isGCEnabled so the referenceSet doesn't continue to
   // accumulate the garbage keys.
-  NSSet<FSTDocumentKey *> *garbage = [self.garbageCollector collectGarbage];
-  if (garbage.count > 0) {
-    for (FSTDocumentKey *key in garbage) {
+  std::set<DocumentKey> garbage = [self.garbageCollector collectGarbage];
+  if (garbage.size() > 0) {
+    for (const DocumentKey &key : garbage) {
       [self.remoteDocumentCache removeEntryForKey:key group:group];
     }
   }
@@ -527,7 +528,7 @@ NS_ASSUME_NONNULL_BEGIN
 
   for (FSTMutationBatch *batch in batches) {
     for (FSTMutation *mutation in batch.mutations) {
-      FSTDocumentKey *key = mutation.key;
+      const DocumentKey &key = mutation.key;
       affectedDocs = [affectedDocs setByAddingObject:key];
     }
   }
diff --git a/Firestore/Source/Local/FSTLocalViewChanges.h b/Firestore/Source/Local/FSTLocalViewChanges.h
index e391472..eb84642 100644
--- a/Firestore/Source/Local/FSTLocalViewChanges.h
+++ b/Firestore/Source/Local/FSTLocalViewChanges.h
@@ -18,7 +18,6 @@
 
 #import "Firestore/Source/Model/FSTDocumentKeySet.h"
 
-@class FSTDocumentKey;
 @class FSTDocumentSet;
 @class FSTMutation;
 @class FSTQuery;
diff --git a/Firestore/Source/Local/FSTMemoryMutationQueue.mm b/Firestore/Source/Local/FSTMemoryMutationQueue.mm
index 8375efd..183a806 100644
--- a/Firestore/Source/Local/FSTMemoryMutationQueue.mm
+++ b/Firestore/Source/Local/FSTMemoryMutationQueue.mm
@@ -18,13 +18,14 @@
 
 #import "Firestore/Source/Core/FSTQuery.h"
 #import "Firestore/Source/Local/FSTDocumentReference.h"
-#import "Firestore/Source/Model/FSTDocumentKey.h"
 #import "Firestore/Source/Model/FSTMutation.h"
 #import "Firestore/Source/Model/FSTMutationBatch.h"
 #import "Firestore/Source/Util/FSTAssert.h"
 
+#include "Firestore/core/src/firebase/firestore/model/document_key.h"
 #include "Firestore/core/src/firebase/firestore/model/resource_path.h"
 
+using firebase::firestore::model::DocumentKey;
 using firebase::firestore::model::ResourcePath;
 
 NS_ASSUME_NONNULL_BEGIN
@@ -230,7 +231,7 @@ static const NSComparator NumberComparator = ^NSComparisonResult(NSNumber *left,
 }
 
 - (NSArray<FSTMutationBatch *> *)allMutationBatchesAffectingDocumentKey:
-    (FSTDocumentKey *)documentKey {
+    (const DocumentKey &)documentKey {
   FSTDocumentReference *start = [[FSTDocumentReference alloc] initWithKey:documentKey ID:0];
 
   NSMutableArray<FSTMutationBatch *> *result = [NSMutableArray array];
@@ -258,17 +259,17 @@ static const NSComparator NumberComparator = ^NSComparisonResult(NSNumber *left,
   // key in this reference must have an even number of segments. The empty segment can be used as
   // a suffix of the query path because it precedes all other segments in an ordered traversal.
   ResourcePath startPath = query.path;
-  if (![FSTDocumentKey isDocumentKey:startPath]) {
+  if (!DocumentKey::IsDocumentKey(startPath)) {
     startPath = startPath.Append("");
   }
   FSTDocumentReference *start =
-      [[FSTDocumentReference alloc] initWithKey:[FSTDocumentKey keyWithPath:startPath] ID:0];
+      [[FSTDocumentReference alloc] initWithKey:DocumentKey{startPath} ID:0];
 
   // Find unique batchIDs referenced by all documents potentially matching the query.
   __block FSTImmutableSortedSet<NSNumber *> *uniqueBatchIDs =
       [FSTImmutableSortedSet setWithComparator:NumberComparator];
   FSTDocumentReferenceBlock block = ^(FSTDocumentReference *reference, BOOL *stop) {
-    const ResourcePath &rowKeyPath = reference.key.path;
+    const ResourcePath &rowKeyPath = reference.key.path();
     if (!prefix.IsPrefixOf(rowKeyPath)) {
       *stop = YES;
       return;
@@ -354,7 +355,7 @@ static const NSComparator NumberComparator = ^NSComparisonResult(NSNumber *left,
   for (FSTMutationBatch *batch in batches) {
     FSTBatchID batchID = batch.batchID;
     for (FSTMutation *mutation in batch.mutations) {
-      FSTDocumentKey *key = mutation.key;
+      const DocumentKey &key = mutation.key;
       [garbageCollector addPotentialGarbageKey:key];
 
       FSTDocumentReference *reference = [[FSTDocumentReference alloc] initWithKey:key ID:batchID];
@@ -373,15 +374,15 @@ static const NSComparator NumberComparator = ^NSComparisonResult(NSNumber *left,
 
 #pragma mark - FSTGarbageSource implementation
 
-- (BOOL)containsKey:(FSTDocumentKey *)key {
+- (BOOL)containsKey:(const DocumentKey &)key {
   // Create a reference with a zero ID as the start position to find any document reference with
   // this key.
   FSTDocumentReference *reference = [[FSTDocumentReference alloc] initWithKey:key ID:0];
 
   NSEnumerator<FSTDocumentReference *> *enumerator =
       [self.batchesByDocumentKey objectEnumeratorFrom:reference];
-  FSTDocumentKey *_Nullable firstKey = [enumerator nextObject].key;
-  return [firstKey isEqual:key];
+  FSTDocumentReference *_Nullable firstReference = [enumerator nextObject];
+  return firstReference && firstReference.key == reference.key;
 }
 
 #pragma mark - Helpers
diff --git a/Firestore/Source/Local/FSTMemoryQueryCache.mm b/Firestore/Source/Local/FSTMemoryQueryCache.mm
index 56d5699..10d04e0 100644
--- a/Firestore/Source/Local/FSTMemoryQueryCache.mm
+++ b/Firestore/Source/Local/FSTMemoryQueryCache.mm
@@ -21,6 +21,8 @@
 #import "Firestore/Source/Local/FSTQueryData.h"
 #import "Firestore/Source/Local/FSTReferenceSet.h"
 
+#include "Firestore/core/src/firebase/firestore/model/document_key.h"
+
 NS_ASSUME_NONNULL_BEGIN
 
 @interface FSTMemoryQueryCache ()
@@ -145,7 +147,7 @@ NS_ASSUME_NONNULL_BEGIN
   self.references.garbageCollector = garbageCollector;
 }
 
-- (BOOL)containsKey:(FSTDocumentKey *)key {
+- (BOOL)containsKey:(const firebase::firestore::model::DocumentKey &)key {
   return [self.references containsKey:key];
 }
 
diff --git a/Firestore/Source/Local/FSTMemoryRemoteDocumentCache.mm b/Firestore/Source/Local/FSTMemoryRemoteDocumentCache.mm
index fed608c..c35c23d 100644
--- a/Firestore/Source/Local/FSTMemoryRemoteDocumentCache.mm
+++ b/Firestore/Source/Local/FSTMemoryRemoteDocumentCache.mm
@@ -19,7 +19,10 @@
 #import "Firestore/Source/Core/FSTQuery.h"
 #import "Firestore/Source/Model/FSTDocument.h"
 #import "Firestore/Source/Model/FSTDocumentDictionary.h"
-#import "Firestore/Source/Model/FSTDocumentKey.h"
+
+#include "Firestore/core/src/firebase/firestore/model/document_key.h"
+
+using firebase::firestore::model::DocumentKey;
 
 NS_ASSUME_NONNULL_BEGIN
 
@@ -46,12 +49,12 @@ NS_ASSUME_NONNULL_BEGIN
   self.docs = [self.docs dictionaryBySettingObject:document forKey:document.key];
 }
 
-- (void)removeEntryForKey:(FSTDocumentKey *)key group:(FSTWriteGroup *)group {
+- (void)removeEntryForKey:(const DocumentKey &)key group:(FSTWriteGroup *)group {
   self.docs = [self.docs dictionaryByRemovingObjectForKey:key];
 }
 
-- (nullable FSTMaybeDocument *)entryForKey:(FSTDocumentKey *)key {
-  return self.docs[key];
+- (nullable FSTMaybeDocument *)entryForKey:(const DocumentKey &)key {
+  return self.docs[static_cast<FSTDocumentKey *>(key)];
 }
 
 - (FSTDocumentDictionary *)documentsMatchingQuery:(FSTQuery *)query {
diff --git a/Firestore/Source/Local/FSTMutationQueue.h b/Firestore/Source/Local/FSTMutationQueue.h
index 12f3284..c2e2548 100644
--- a/Firestore/Source/Local/FSTMutationQueue.h
+++ b/Firestore/Source/Local/FSTMutationQueue.h
@@ -19,7 +19,8 @@
 #import "Firestore/Source/Core/FSTTypes.h"
 #import "Firestore/Source/Local/FSTGarbageCollector.h"
 
-@class FSTDocumentKey;
+#include "Firestore/core/src/firebase/firestore/model/document_key.h"
+
 @class FSTMutation;
 @class FSTMutationBatch;
 @class FSTQuery;
@@ -123,7 +124,7 @@ NS_ASSUME_NONNULL_BEGIN
 // TODO(mcg): This should really return an NSEnumerator
 // also for b/32992024, all backing stores should really index by document key
 - (NSArray<FSTMutationBatch *> *)allMutationBatchesAffectingDocumentKey:
-    (FSTDocumentKey *)documentKey;
+    (const firebase::firestore::model::DocumentKey &)documentKey;
 
 /**
  * Finds all mutation batches that could affect the results for the given query. Not all
diff --git a/Firestore/Source/Local/FSTNoOpGarbageCollector.h b/Firestore/Source/Local/FSTNoOpGarbageCollector.h
index c9b5862..f58c5a0 100644
--- a/Firestore/Source/Local/FSTNoOpGarbageCollector.h
+++ b/Firestore/Source/Local/FSTNoOpGarbageCollector.h
@@ -18,8 +18,6 @@
 
 #import "Firestore/Source/Local/FSTGarbageCollector.h"
 
-@class FSTDocumentKey;
-
 NS_ASSUME_NONNULL_BEGIN
 
 /**
diff --git a/Firestore/Source/Local/FSTNoOpGarbageCollector.mm b/Firestore/Source/Local/FSTNoOpGarbageCollector.mm
index e03b599..451cde2 100644
--- a/Firestore/Source/Local/FSTNoOpGarbageCollector.mm
+++ b/Firestore/Source/Local/FSTNoOpGarbageCollector.mm
@@ -16,6 +16,10 @@
 
 #import "Firestore/Source/Local/FSTNoOpGarbageCollector.h"
 
+#include "Firestore/core/src/firebase/firestore/model/document_key.h"
+
+using firebase::firestore::model::DocumentKey;
+
 NS_ASSUME_NONNULL_BEGIN
 
 @implementation FSTNoOpGarbageCollector
@@ -32,12 +36,12 @@ NS_ASSUME_NONNULL_BEGIN
   // Not tracking garbage so don't track sources.
 }
 
-- (void)addPotentialGarbageKey:(FSTDocumentKey *)key {
+- (void)addPotentialGarbageKey:(const DocumentKey&)key {
   // Not tracking garbage so ignore.
 }
 
-- (NSSet<FSTDocumentKey *> *)collectGarbage {
-  return [NSSet set];
+- (std::set<DocumentKey>)collectGarbage {
+  return {};
 }
 
 @end
diff --git a/Firestore/Source/Local/FSTQueryCache.h b/Firestore/Source/Local/FSTQueryCache.h
index 5c43de4..1b72290 100644
--- a/Firestore/Source/Local/FSTQueryCache.h
+++ b/Firestore/Source/Local/FSTQueryCache.h
@@ -20,7 +20,6 @@
 #import "Firestore/Source/Local/FSTGarbageCollector.h"
 #import "Firestore/Source/Model/FSTDocumentKeySet.h"
 
-@class FSTDocumentKey;
 @class FSTDocumentSet;
 @class FSTMaybeDocument;
 @class FSTQuery;
diff --git a/Firestore/Source/Local/FSTReferenceSet.h b/Firestore/Source/Local/FSTReferenceSet.h
index 66285d9..9d842cb 100644
--- a/Firestore/Source/Local/FSTReferenceSet.h
+++ b/Firestore/Source/Local/FSTReferenceSet.h
@@ -20,8 +20,6 @@
 #import "Firestore/Source/Local/FSTGarbageCollector.h"
 #import "Firestore/Source/Model/FSTDocumentKeySet.h"
 
-@class FSTDocumentKey;
-
 NS_ASSUME_NONNULL_BEGIN
 
 /**
@@ -46,13 +44,13 @@ NS_ASSUME_NONNULL_BEGIN
 - (BOOL)isEmpty;
 
 /** Adds a reference to the given document key for the given ID. */
-- (void)addReferenceToKey:(FSTDocumentKey *)key forID:(int)ID;
+- (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;
 
 /** Removes a reference to the given document key for the given ID. */
-- (void)removeReferenceToKey:(FSTDocumentKey *)key forID:(int)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;
diff --git a/Firestore/Source/Local/FSTReferenceSet.mm b/Firestore/Source/Local/FSTReferenceSet.mm
index 15bb033..14f5d47 100644
--- a/Firestore/Source/Local/FSTReferenceSet.mm
+++ b/Firestore/Source/Local/FSTReferenceSet.mm
@@ -17,7 +17,10 @@
 #import "Firestore/Source/Local/FSTReferenceSet.h"
 
 #import "Firestore/Source/Local/FSTDocumentReference.h"
-#import "Firestore/Source/Model/FSTDocumentKey.h"
+
+#include "Firestore/core/src/firebase/firestore/model/document_key.h"
+
+using firebase::firestore::model::DocumentKey;
 
 NS_ASSUME_NONNULL_BEGIN
 
@@ -59,7 +62,7 @@ NS_ASSUME_NONNULL_BEGIN
 
 #pragma mark - Public methods
 
-- (void)addReferenceToKey:(FSTDocumentKey *)key forID:(int)ID {
+- (void)addReferenceToKey:(const DocumentKey &)key forID:(int)ID {
   FSTDocumentReference *reference = [[FSTDocumentReference alloc] initWithKey:key ID:ID];
   self.referencesByKey = [self.referencesByKey setByAddingObject:reference];
   self.referencesByID = [self.referencesByID setByAddingObject:reference];
@@ -71,7 +74,7 @@ NS_ASSUME_NONNULL_BEGIN
   }];
 }
 
-- (void)removeReferenceToKey:(FSTDocumentKey *)key forID:(int)ID {
+- (void)removeReferenceToKey:(const DocumentKey &)key forID:(int)ID {
   [self removeReference:[[FSTDocumentReference alloc] initWithKey:key ID:ID]];
 }
 
@@ -82,9 +85,10 @@ NS_ASSUME_NONNULL_BEGIN
 }
 
 - (void)removeReferencesForID:(int)ID {
-  FSTDocumentKey *emptyKey = [FSTDocumentKey keyWithSegments:{}];
-  FSTDocumentReference *start = [[FSTDocumentReference alloc] initWithKey:emptyKey ID:ID];
-  FSTDocumentReference *end = [[FSTDocumentReference alloc] initWithKey:emptyKey ID:(ID + 1)];
+  FSTDocumentReference *start =
+      [[FSTDocumentReference alloc] initWithKey:DocumentKey::Empty() ID:ID];
+  FSTDocumentReference *end =
+      [[FSTDocumentReference alloc] initWithKey:DocumentKey::Empty() ID:(ID + 1)];
 
   [self.referencesByID enumerateObjectsFrom:start
                                          to:end
@@ -106,9 +110,10 @@ NS_ASSUME_NONNULL_BEGIN
 }
 
 - (FSTDocumentKeySet *)referencedKeysForID:(int)ID {
-  FSTDocumentKey *emptyKey = [FSTDocumentKey keyWithSegments:{}];
-  FSTDocumentReference *start = [[FSTDocumentReference alloc] initWithKey:emptyKey ID:ID];
-  FSTDocumentReference *end = [[FSTDocumentReference alloc] initWithKey:emptyKey ID:(ID + 1)];
+  FSTDocumentReference *start =
+      [[FSTDocumentReference alloc] initWithKey:DocumentKey::Empty() ID:ID];
+  FSTDocumentReference *end =
+      [[FSTDocumentReference alloc] initWithKey:DocumentKey::Empty() ID:(ID + 1)];
 
   __block FSTDocumentKeySet *keys = [FSTDocumentKeySet keySet];
   [self.referencesByID enumerateObjectsFrom:start
@@ -119,15 +124,15 @@ NS_ASSUME_NONNULL_BEGIN
   return keys;
 }
 
-- (BOOL)containsKey:(FSTDocumentKey *)key {
+- (BOOL)containsKey:(const DocumentKey &)key {
   // Create a reference with a zero ID as the start position to find any document reference with
   // this key.
   FSTDocumentReference *reference = [[FSTDocumentReference alloc] initWithKey:key ID:0];
 
   NSEnumerator<FSTDocumentReference *> *enumerator =
       [self.referencesByKey objectEnumeratorFrom:reference];
-  FSTDocumentKey *_Nullable firstKey = [enumerator nextObject].key;
-  return [firstKey isEqual:reference.key];
+  FSTDocumentReference *_Nullable firstReference = [enumerator nextObject];
+  return firstReference && firstReference.key == reference.key;
 }
 
 @end
diff --git a/Firestore/Source/Local/FSTRemoteDocumentCache.h b/Firestore/Source/Local/FSTRemoteDocumentCache.h
index fa42ce5..e2ad988 100644
--- a/Firestore/Source/Local/FSTRemoteDocumentCache.h
+++ b/Firestore/Source/Local/FSTRemoteDocumentCache.h
@@ -18,7 +18,8 @@
 
 #import "Firestore/Source/Model/FSTDocumentDictionary.h"
 
-@class FSTDocumentKey;
+#include "Firestore/core/src/firebase/firestore/model/document_key.h"
+
 @class FSTMaybeDocument;
 @class FSTQuery;
 @class FSTWriteGroup;
@@ -48,7 +49,8 @@ NS_ASSUME_NONNULL_BEGIN
 - (void)addEntry:(FSTMaybeDocument *)maybeDocument group:(FSTWriteGroup *)group;
 
 /** Removes the cached entry for the given key (no-op if no entry exists). */
-- (void)removeEntryForKey:(FSTDocumentKey *)documentKey group:(FSTWriteGroup *)group;
+- (void)removeEntryForKey:(const firebase::firestore::model::DocumentKey &)documentKey
+                    group:(FSTWriteGroup *)group;
 
 /**
  * Looks up an entry in the cache.
@@ -56,7 +58,8 @@ NS_ASSUME_NONNULL_BEGIN
  * @param documentKey The key of the entry to look up.
  * @return The cached FSTDocument or FSTDeletedDocument entry, or nil if we have nothing cached.
  */
-- (nullable FSTMaybeDocument *)entryForKey:(FSTDocumentKey *)documentKey;
+- (nullable FSTMaybeDocument *)entryForKey:
+    (const firebase::firestore::model::DocumentKey &)documentKey;
 
 /**
  * Executes a query against the cached FSTDocument entries
diff --git a/Firestore/Source/Local/FSTRemoteDocumentChangeBuffer.h b/Firestore/Source/Local/FSTRemoteDocumentChangeBuffer.h
index be0d609..cfdf55b 100644
--- a/Firestore/Source/Local/FSTRemoteDocumentChangeBuffer.h
+++ b/Firestore/Source/Local/FSTRemoteDocumentChangeBuffer.h
@@ -16,11 +16,12 @@
 
 #import <Foundation/Foundation.h>
 
+#include "Firestore/core/src/firebase/firestore/model/document_key.h"
+
 NS_ASSUME_NONNULL_BEGIN
 
 @protocol FSTRemoteDocumentCache;
 @class FSTMaybeDocument;
-@class FSTDocumentKey;
 @class FSTWriteGroup;
 
 /**
@@ -53,7 +54,8 @@ NS_ASSUME_NONNULL_BEGIN
  * @param documentKey The key of the entry to look up.
  * @return The cached FSTDocument or FSTDeletedDocument entry, or nil if we have nothing cached.
  */
-- (nullable FSTMaybeDocument *)entryForKey:(FSTDocumentKey *)documentKey;
+- (nullable FSTMaybeDocument *)entryForKey:
+    (const firebase::firestore::model::DocumentKey &)documentKey;
 
 /**
  * Applies buffered changes to the underlying FSTRemoteDocumentCache, using the provided
diff --git a/Firestore/Source/Local/FSTRemoteDocumentChangeBuffer.mm b/Firestore/Source/Local/FSTRemoteDocumentChangeBuffer.mm
index b228461..4f29bdf 100644
--- a/Firestore/Source/Local/FSTRemoteDocumentChangeBuffer.mm
+++ b/Firestore/Source/Local/FSTRemoteDocumentChangeBuffer.mm
@@ -16,12 +16,17 @@
 
 #import "Firestore/Source/Local/FSTRemoteDocumentChangeBuffer.h"
 
+#include <map>
+#include <memory>
+
 #import "Firestore/Source/Local/FSTRemoteDocumentCache.h"
 #import "Firestore/Source/Model/FSTDocument.h"
-#import "Firestore/Source/Model/FSTDocumentKey.h"
 #import "Firestore/Source/Util/FSTAssert.h"
 
 #include "Firestore/core/src/firebase/firestore/model/document_key.h"
+#include "absl/memory/memory.h"
+
+using firebase::firestore::model::DocumentKey;
 
 NS_ASSUME_NONNULL_BEGIN
 
@@ -32,13 +37,12 @@ NS_ASSUME_NONNULL_BEGIN
 /** The underlying cache we're buffering changes for. */
 @property(nonatomic, strong, nonnull) id<FSTRemoteDocumentCache> remoteDocumentCache;
 
-/** The buffered changes, stored as a dictionary for easy lookups. */
-@property(nonatomic, strong, nullable)
-    NSMutableDictionary<FSTDocumentKey *, FSTMaybeDocument *> *changes;
-
 @end
 
-@implementation FSTRemoteDocumentChangeBuffer
+@implementation FSTRemoteDocumentChangeBuffer {
+  /** The buffered changes, stored as a dictionary for easy lookups. */
+  std::unique_ptr<std::map<DocumentKey, FSTMaybeDocument *>> _changes;
+}
 
 + (instancetype)changeBufferWithCache:(id<FSTRemoteDocumentCache>)cache {
   return [[FSTRemoteDocumentChangeBuffer alloc] initWithCache:cache];
@@ -47,7 +51,7 @@ NS_ASSUME_NONNULL_BEGIN
 - (instancetype)initWithCache:(id<FSTRemoteDocumentCache>)cache {
   if (self = [super init]) {
     _remoteDocumentCache = cache;
-    _changes = [NSMutableDictionary dictionary];
+    _changes = absl::make_unique<std::map<DocumentKey, FSTMaybeDocument *>>();
   }
   return self;
 }
@@ -55,34 +59,33 @@ NS_ASSUME_NONNULL_BEGIN
 - (void)addEntry:(FSTMaybeDocument *)maybeDocument {
   [self assertValid];
 
-  self.changes[(FSTDocumentKey *)maybeDocument.key] = maybeDocument;
+  (*_changes)[maybeDocument.key] = maybeDocument;
 }
 
-- (nullable FSTMaybeDocument *)entryForKey:(FSTDocumentKey *)documentKey {
+- (nullable FSTMaybeDocument *)entryForKey:(const DocumentKey &)documentKey {
   [self assertValid];
 
-  FSTMaybeDocument *bufferedEntry = self.changes[documentKey];
-  if (bufferedEntry) {
-    return bufferedEntry;
-  } else {
+  const auto iter = _changes->find(documentKey);
+  if (iter == _changes->end()) {
     return [self.remoteDocumentCache entryForKey:documentKey];
+  } else {
+    return iter->second;
   }
 }
 
 - (void)applyToWriteGroup:(FSTWriteGroup *)group {
   [self assertValid];
 
-  [self.changes enumerateKeysAndObjectsUsingBlock:^(FSTDocumentKey *key, FSTMaybeDocument *value,
-                                                    BOOL *stop) {
-    [self.remoteDocumentCache addEntry:value group:group];
-  }];
+  for (const auto &kv : *_changes) {
+    [self.remoteDocumentCache addEntry:kv.second group:group];
+  }
 
   // We should not be used to buffer any more changes.
-  self.changes = nil;
+  _changes.reset();
 }
 
 - (void)assertValid {
-  FSTAssert(self.changes, @"Changes have already been applied.");
+  FSTAssert(_changes, @"Changes have already been applied.");
 }
 
 @end