port C++ DocumentKey to the rest of Firestore code (#977)

* port C++ DocumentKey to API's and Core's

* address changes

* address changes

* fix Hash return types

30 files changed, 254 insertions, 193 deletions

diff --git a/Firestore/Example/Tests/Core/FSTSyncEngine+Testing.h b/Firestore/Example/Tests/Core/FSTSyncEngine+Testing.h
index ab0697b..6e7c45b 100644
--- a/Firestore/Example/Tests/Core/FSTSyncEngine+Testing.h
+++ b/Firestore/Example/Tests/Core/FSTSyncEngine+Testing.h
@@ -18,14 +18,15 @@
 
 #import "Firestore/Source/Core/FSTSyncEngine.h"
 
-@class FSTDocumentKey;
+#include "Firestore/core/src/firebase/firestore/model/document_key.h"
 
 NS_ASSUME_NONNULL_BEGIN
 
 @interface FSTSyncEngine (Testing)
 
 /** Returns the current set of limbo document keys and their associated target IDs. */
-- (NSDictionary<FSTDocumentKey *, FSTBoxedTargetID *> *)currentLimboDocuments;
+- (std::map<firebase::firestore::model::DocumentKey, firebase::firestore::model::TargetId>)
+    currentLimboDocuments;
 
 @end
 
diff --git a/Firestore/Example/Tests/Integration/FSTDatastoreTests.mm b/Firestore/Example/Tests/Integration/FSTDatastoreTests.mm
index 9ce3d30..430366f 100644
--- a/Firestore/Example/Tests/Integration/FSTDatastoreTests.mm
+++ b/Firestore/Example/Tests/Integration/FSTDatastoreTests.mm
@@ -48,6 +48,7 @@ namespace util = firebase::firestore::util;
 using firebase::firestore::auth::EmptyCredentialsProvider;
 using firebase::firestore::core::DatabaseInfo;
 using firebase::firestore::model::DatabaseId;
+using firebase::firestore::model::TargetId;
 
 NS_ASSUME_NONNULL_BEGIN
 
@@ -126,7 +127,7 @@ NS_ASSUME_NONNULL_BEGIN
   [expectation fulfill];
 }
 
-- (void)rejectListenWithTargetID:(FSTBoxedTargetID *)targetID error:(NSError *)error {
+- (void)rejectListenWithTargetID:(const TargetId)targetID error:(NSError *)error {
   FSTFail(@"Not implemented");
 }
 
diff --git a/Firestore/Example/Tests/SpecTests/FSTSpecTests.mm b/Firestore/Example/Tests/SpecTests/FSTSpecTests.mm
index 0c3d9a1..f66e6c7 100644
--- a/Firestore/Example/Tests/SpecTests/FSTSpecTests.mm
+++ b/Firestore/Example/Tests/SpecTests/FSTSpecTests.mm
@@ -16,6 +16,7 @@
 
 #import "Firestore/Example/Tests/SpecTests/FSTSpecTests.h"
 
+#include <map>
 #include <utility>
 
 #import <FirebaseFirestore/FIRFirestoreErrors.h>
@@ -44,10 +45,13 @@
 #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/src/firebase/firestore/util/string_apple.h"
 
 namespace util = firebase::firestore::util;
 using firebase::firestore::auth::User;
+using firebase::firestore::model::DocumentKey;
+using firebase::firestore::model::TargetId;
 
 NS_ASSUME_NONNULL_BEGIN
 
@@ -561,22 +565,22 @@ static NSString *const kNoIOSTag = @"no-ios";
 
 - (void)validateLimboDocuments {
   // Make a copy so it can modified while checking against the expected limbo docs.
-  NSMutableDictionary<FSTDocumentKey *, FSTBoxedTargetID *> *actualLimboDocs =
-      [NSMutableDictionary dictionaryWithDictionary:self.driver.currentLimboDocuments];
+  std::map<DocumentKey, TargetId> actualLimboDocs = self.driver.currentLimboDocuments;
 
   // Validate that each limbo doc has an expected active target
-  [actualLimboDocs enumerateKeysAndObjectsUsingBlock:^(FSTDocumentKey *key,
-                                                       FSTBoxedTargetID *targetID, BOOL *stop) {
-    XCTAssertNotNil(self.driver.expectedActiveTargets[targetID],
+  for (const auto &kv : actualLimboDocs) {
+    XCTAssertNotNil(self.driver.expectedActiveTargets[@(kv.second)],
                     @"Found limbo doc without an expected active target");
-  }];
+  }
 
   for (FSTDocumentKey *expectedLimboDoc in self.driver.expectedLimboDocuments) {
-    XCTAssertNotNil(actualLimboDocs[expectedLimboDoc],
-                    @"Expected doc to be in limbo, but was not: %@", expectedLimboDoc);
-    [actualLimboDocs removeObjectForKey:expectedLimboDoc];
+    XCTAssert(actualLimboDocs.find(expectedLimboDoc) != actualLimboDocs.end(),
+              @"Expected doc to be in limbo, but was not: %@", expectedLimboDoc);
+    actualLimboDocs.erase(expectedLimboDoc);
   }
-  XCTAssertTrue(actualLimboDocs.count == 0, "Unexpected docs in limbo: %@", actualLimboDocs);
+  XCTAssertTrue(actualLimboDocs.empty(), "%lu Unexpected docs in limbo, the first one is <%s, %d>",
+                actualLimboDocs.size(), actualLimboDocs.begin()->first.ToString().c_str(),
+                actualLimboDocs.begin()->second);
 }
 
 - (void)validateActiveTargets {
diff --git a/Firestore/Example/Tests/SpecTests/FSTSyncEngineTestDriver.h b/Firestore/Example/Tests/SpecTests/FSTSyncEngineTestDriver.h
index f3d9a5d..ac44cb5 100644
--- a/Firestore/Example/Tests/SpecTests/FSTSyncEngineTestDriver.h
+++ b/Firestore/Example/Tests/SpecTests/FSTSyncEngineTestDriver.h
@@ -16,6 +16,7 @@
 
 #import <Foundation/Foundation.h>
 
+#include <map>
 #include <unordered_map>
 
 #import "Firestore/Source/Core/FSTTypes.h"
@@ -23,6 +24,7 @@
 #import "Firestore/Source/Util/FSTDispatchQueue.h"
 
 #include "Firestore/core/src/firebase/firestore/auth/user.h"
+#include "Firestore/core/src/firebase/firestore/model/document_key.h"
 
 @class FSTDocumentKey;
 @class FSTMutation;
@@ -242,6 +244,10 @@ typedef std::unordered_map<firebase::firestore::auth::User,
  */
 - (NSArray<FSTQueryEvent *> *)capturedEventsSinceLastCall;
 
+/** The current set of documents in limbo. */
+- (std::map<firebase::firestore::model::DocumentKey, firebase::firestore::model::TargetId>)
+    currentLimboDocuments;
+
 /**
  * The writes that have been sent to the FSTSyncEngine via writeUserMutation: but not yet
  * acknowledged by calling receiveWriteAck/Error:. They are tracked per-user.
@@ -263,10 +269,6 @@ typedef std::unordered_map<firebase::firestore::auth::User,
 /** The current user for the FSTSyncEngine; determines which mutation queue is active. */
 @property(nonatomic, assign, readonly) const firebase::firestore::auth::User &currentUser;
 
-/** The current set of documents in limbo. */
-@property(nonatomic, strong, readonly)
-    NSDictionary<FSTDocumentKey *, FSTBoxedTargetID *> *currentLimboDocuments;
-
 /** The expected set of documents in limbo. */
 @property(nonatomic, strong, readwrite) NSSet<FSTDocumentKey *> *expectedLimboDocuments;
 
diff --git a/Firestore/Example/Tests/SpecTests/FSTSyncEngineTestDriver.mm b/Firestore/Example/Tests/SpecTests/FSTSyncEngineTestDriver.mm
index bfcd4dd..42a0ac0 100644
--- a/Firestore/Example/Tests/SpecTests/FSTSyncEngineTestDriver.mm
+++ b/Firestore/Example/Tests/SpecTests/FSTSyncEngineTestDriver.mm
@@ -16,6 +16,7 @@
 
 #import "Firestore/Example/Tests/SpecTests/FSTSyncEngineTestDriver.h"
 
+#include <map>
 #include <unordered_map>
 
 #import <FirebaseFirestore/FIRFirestoreErrors.h>
@@ -40,12 +41,15 @@
 #include "Firestore/core/src/firebase/firestore/auth/user.h"
 #include "Firestore/core/src/firebase/firestore/core/database_info.h"
 #include "Firestore/core/src/firebase/firestore/model/database_id.h"
+#include "Firestore/core/src/firebase/firestore/model/document_key.h"
 
 using firebase::firestore::auth::EmptyCredentialsProvider;
 using firebase::firestore::auth::HashUser;
 using firebase::firestore::auth::User;
 using firebase::firestore::core::DatabaseInfo;
 using firebase::firestore::model::DatabaseId;
+using firebase::firestore::model::DocumentKey;
+using firebase::firestore::model::TargetId;
 
 NS_ASSUME_NONNULL_BEGIN
 
@@ -349,7 +353,7 @@ NS_ASSUME_NONNULL_BEGIN
   }];
 }
 
-- (NSDictionary<FSTDocumentKey *, FSTBoxedTargetID *> *)currentLimboDocuments {
+- (std::map<DocumentKey, TargetId>)currentLimboDocuments {
   return [self.syncEngine currentLimboDocuments];
 }
 
diff --git a/Firestore/Source/API/FIRCollectionReference.mm b/Firestore/Source/API/FIRCollectionReference.mm
index dfd06c9..3f1559e 100644
--- a/Firestore/Source/API/FIRCollectionReference.mm
+++ b/Firestore/Source/API/FIRCollectionReference.mm
@@ -23,14 +23,15 @@
 #import "Firestore/Source/API/FIRQuery+Internal.h"
 #import "Firestore/Source/API/FIRQuery_Init.h"
 #import "Firestore/Source/Core/FSTQuery.h"
-#import "Firestore/Source/Model/FSTDocumentKey.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/resource_path.h"
 #include "Firestore/core/src/firebase/firestore/util/string_apple.h"
 
 namespace util = firebase::firestore::util;
+using firebase::firestore::model::DocumentKey;
 using firebase::firestore::model::ResourcePath;
 using firebase::firestore::util::CreateAutoId;
 
@@ -99,7 +100,7 @@ NS_ASSUME_NONNULL_BEGIN
   if (parentPath.empty()) {
     return nil;
   } else {
-    FSTDocumentKey *key = [FSTDocumentKey keyWithPath:parentPath];
+    DocumentKey key{parentPath};
     return [FIRDocumentReference referenceWithKey:key firestore:self.firestore];
   }
 }
@@ -130,7 +131,7 @@ NS_ASSUME_NONNULL_BEGIN
 }
 
 - (FIRDocumentReference *)documentWithAutoID {
-  FSTDocumentKey *key = [FSTDocumentKey keyWithPath:self.query.path.Append(CreateAutoId())];
+  const DocumentKey key{self.query.path.Append(CreateAutoId())};
   return [FIRDocumentReference referenceWithKey:key firestore:self.firestore];
 }
 
diff --git a/Firestore/Source/API/FIRDocumentReference+Internal.h b/Firestore/Source/API/FIRDocumentReference+Internal.h
index 706e8db..eb078ca 100644
--- a/Firestore/Source/API/FIRDocumentReference+Internal.h
+++ b/Firestore/Source/API/FIRDocumentReference+Internal.h
@@ -16,20 +16,20 @@
 
 #import "FIRDocumentReference.h"
 
+#include "Firestore/core/src/firebase/firestore/model/document_key.h"
 #include "Firestore/core/src/firebase/firestore/model/resource_path.h"
 
 NS_ASSUME_NONNULL_BEGIN
 
-@class FSTDocumentKey;
-
 /** Internal FIRDocumentReference API we don't want exposed in our public header files. */
 @interface FIRDocumentReference (Internal)
 
 + (instancetype)referenceWithPath:(const firebase::firestore::model::ResourcePath &)path
                         firestore:(FIRFirestore *)firestore;
-+ (instancetype)referenceWithKey:(FSTDocumentKey *)key firestore:(FIRFirestore *)firestore;
++ (instancetype)referenceWithKey:(firebase::firestore::model::DocumentKey)key
+                       firestore:(FIRFirestore *)firestore;
 
-@property(nonatomic, strong, readonly) FSTDocumentKey *key;
+- (const firebase::firestore::model::DocumentKey &)key;
 
 @end
 
diff --git a/Firestore/Source/API/FIRDocumentReference.mm b/Firestore/Source/API/FIRDocumentReference.mm
index 5968fb2..cc52d45 100644
--- a/Firestore/Source/API/FIRDocumentReference.mm
+++ b/Firestore/Source/API/FIRDocumentReference.mm
@@ -16,6 +16,9 @@
 
 #import "FIRDocumentReference.h"
 
+#include <memory>
+#include <utility>
+
 #import <GRPCClient/GRPCCall.h>
 
 #import "FIRFirestoreErrors.h"
@@ -30,7 +33,6 @@
 #import "Firestore/Source/Core/FSTEventManager.h"
 #import "Firestore/Source/Core/FSTFirestoreClient.h"
 #import "Firestore/Source/Core/FSTQuery.h"
-#import "Firestore/Source/Model/FSTDocumentKey.h"
 #import "Firestore/Source/Model/FSTDocumentSet.h"
 #import "Firestore/Source/Model/FSTFieldValue.h"
 #import "Firestore/Source/Model/FSTMutation.h"
@@ -38,10 +40,12 @@
 #import "Firestore/Source/Util/FSTAsyncQueryListener.h"
 #import "Firestore/Source/Util/FSTUsageValidation.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"
 
 namespace util = firebase::firestore::util;
+using firebase::firestore::model::DocumentKey;
 using firebase::firestore::model::ResourcePath;
 
 NS_ASSUME_NONNULL_BEGIN
@@ -83,35 +87,17 @@ NS_ASSUME_NONNULL_BEGIN
 #pragma mark - FIRDocumentReference
 
 @interface FIRDocumentReference ()
-- (instancetype)initWithKey:(FSTDocumentKey *)key
+- (instancetype)initWithKey:(DocumentKey)key
                   firestore:(FIRFirestore *)firestore NS_DESIGNATED_INITIALIZER;
-@property(nonatomic, strong, readonly) FSTDocumentKey *key;
 @end
 
-@implementation FIRDocumentReference (Internal)
-
-+ (instancetype)referenceWithPath:(const ResourcePath &)path firestore:(FIRFirestore *)firestore {
-  if (path.size() % 2 != 0) {
-    FSTThrowInvalidArgument(
-        @"Invalid document reference. Document references must have an even "
-         "number of segments, but %s has %zu",
-        path.CanonicalString().c_str(), path.size());
-  }
-  return
-      [FIRDocumentReference referenceWithKey:[FSTDocumentKey keyWithPath:path] firestore:firestore];
-}
-
-+ (instancetype)referenceWithKey:(FSTDocumentKey *)key firestore:(FIRFirestore *)firestore {
-  return [[FIRDocumentReference alloc] initWithKey:key firestore:firestore];
+@implementation FIRDocumentReference {
+  DocumentKey _key;
 }
 
-@end
-
-@implementation FIRDocumentReference
-
-- (instancetype)initWithKey:(FSTDocumentKey *)key firestore:(FIRFirestore *)firestore {
+- (instancetype)initWithKey:(DocumentKey)key firestore:(FIRFirestore *)firestore {
   if (self = [super init]) {
-    _key = key;
+    _key = std::move(key);
     _firestore = firestore;
   }
   return self;
@@ -129,28 +115,28 @@ NS_ASSUME_NONNULL_BEGIN
 - (BOOL)isEqualToReference:(nullable FIRDocumentReference *)reference {
   if (self == reference) return YES;
   if (reference == nil) return NO;
-  return [self.firestore isEqual:reference.firestore] && [self.key isEqualToKey:reference.key];
+  return [self.firestore isEqual:reference.firestore] && self.key == reference.key;
 }
 
 - (NSUInteger)hash {
   NSUInteger hash = [self.firestore hash];
-  hash = hash * 31u + [self.key hash];
+  hash = hash * 31u + self.key.Hash();
   return hash;
 }
 
 #pragma mark - Public Methods
 
 - (NSString *)documentID {
-  return util::WrapNSString(self.key.path.last_segment());
+  return util::WrapNSString(self.key.path().last_segment());
 }
 
 - (FIRCollectionReference *)parent {
   return
-      [FIRCollectionReference referenceWithPath:self.key.path.PopLast() firestore:self.firestore];
+      [FIRCollectionReference referenceWithPath:self.key.path().PopLast() firestore:self.firestore];
 }
 
 - (NSString *)path {
-  return util::WrapNSString(self.key.path.CanonicalString());
+  return util::WrapNSString(self.key.path().CanonicalString());
 }
 
 - (FIRCollectionReference *)collectionWithPath:(NSString *)collectionPath {
@@ -158,7 +144,7 @@ NS_ASSUME_NONNULL_BEGIN
     FSTThrowInvalidArgument(@"Collection path cannot be nil.");
   }
   const ResourcePath subPath = ResourcePath::FromString(util::MakeStringView(collectionPath));
-  const ResourcePath path = self.key.path.Append(subPath);
+  const ResourcePath path = self.key.path().Append(subPath);
   return [FIRCollectionReference referenceWithPath:path firestore:self.firestore];
 }
 
@@ -271,8 +257,8 @@ NS_ASSUME_NONNULL_BEGIN
 addSnapshotListenerInternalWithOptions:(FSTListenOptions *)internalOptions
                               listener:(FIRDocumentSnapshotBlock)listener {
   FIRFirestore *firestore = self.firestore;
-  FSTQuery *query = [FSTQuery queryWithPath:self.key.path];
-  FSTDocumentKey *key = self.key;
+  FSTQuery *query = [FSTQuery queryWithPath:self.key.path()];
+  const DocumentKey key = self.key;
 
   FSTViewSnapshotHandler snapshotHandler = ^(FSTViewSnapshot *snapshot, NSError *error) {
     if (error) {
@@ -313,4 +299,28 @@ addSnapshotListenerInternalWithOptions:(FSTListenOptions *)internalOptions
 
 @end
 
+#pragma mark - FIRDocumentReference (Internal)
+
+@implementation FIRDocumentReference (Internal)
+
++ (instancetype)referenceWithPath:(const ResourcePath &)path firestore:(FIRFirestore *)firestore {
+  if (path.size() % 2 != 0) {
+    FSTThrowInvalidArgument(
+        @"Invalid document reference. Document references must have an even "
+         "number of segments, but %s has %zu",
+        path.CanonicalString().c_str(), path.size());
+  }
+  return [FIRDocumentReference referenceWithKey:DocumentKey{path} firestore:firestore];
+}
+
++ (instancetype)referenceWithKey:(DocumentKey)key firestore:(FIRFirestore *)firestore {
+  return [[FIRDocumentReference alloc] initWithKey:std::move(key) firestore:firestore];
+}
+
+- (const firebase::firestore::model::DocumentKey &)key {
+  return _key;
+}
+
+@end
+
 NS_ASSUME_NONNULL_END
diff --git a/Firestore/Source/API/FIRDocumentSnapshot+Internal.h b/Firestore/Source/API/FIRDocumentSnapshot+Internal.h
index f2776f0..f4cef9a 100644
--- a/Firestore/Source/API/FIRDocumentSnapshot+Internal.h
+++ b/Firestore/Source/API/FIRDocumentSnapshot+Internal.h
@@ -16,9 +16,10 @@
 
 #import "FIRDocumentSnapshot.h"
 
+#include "Firestore/core/src/firebase/firestore/model/document_key.h"
+
 @class FIRFirestore;
 @class FSTDocument;
-@class FSTDocumentKey;
 
 NS_ASSUME_NONNULL_BEGIN
 
@@ -26,7 +27,7 @@ NS_ASSUME_NONNULL_BEGIN
 @interface FIRDocumentSnapshot (Internal)
 
 + (instancetype)snapshotWithFirestore:(FIRFirestore *)firestore
-                          documentKey:(FSTDocumentKey *)documentKey
+                          documentKey:(firebase::firestore::model::DocumentKey)documentKey
                              document:(nullable FSTDocument *)document
                             fromCache:(BOOL)fromCache;
 
diff --git a/Firestore/Source/API/FIRDocumentSnapshot.mm b/Firestore/Source/API/FIRDocumentSnapshot.mm
index 4d82986..39d6af1 100644
--- a/Firestore/Source/API/FIRDocumentSnapshot.mm
+++ b/Firestore/Source/API/FIRDocumentSnapshot.mm
@@ -16,34 +16,38 @@
 
 #import "FIRDocumentSnapshot.h"
 
+#include <utility>
+
 #import "Firestore/Source/API/FIRDocumentReference+Internal.h"
 #import "Firestore/Source/API/FIRFieldPath+Internal.h"
 #import "Firestore/Source/API/FIRFirestore+Internal.h"
 #import "Firestore/Source/API/FIRSnapshotMetadata+Internal.h"
 #import "Firestore/Source/API/FIRSnapshotOptions+Internal.h"
 #import "Firestore/Source/Model/FSTDocument.h"
-#import "Firestore/Source/Model/FSTDocumentKey.h"
 #import "Firestore/Source/Model/FSTFieldValue.h"
 #import "Firestore/Source/Util/FSTAssert.h"
 #import "Firestore/Source/Util/FSTUsageValidation.h"
 
 #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/util/string_apple.h"
 
 namespace util = firebase::firestore::util;
 using firebase::firestore::model::DatabaseId;
+using firebase::firestore::model::DocumentKey;
 
 NS_ASSUME_NONNULL_BEGIN
 
 @interface FIRDocumentSnapshot ()
 
 - (instancetype)initWithFirestore:(FIRFirestore *)firestore
-                      documentKey:(FSTDocumentKey *)documentKey
+                      documentKey:(DocumentKey)documentKey
                          document:(nullable FSTDocument *)document
                         fromCache:(BOOL)fromCache NS_DESIGNATED_INITIALIZER;
 
+- (const DocumentKey &)internalKey;
+
 @property(nonatomic, strong, readonly) FIRFirestore *firestore;
-@property(nonatomic, strong, readonly) FSTDocumentKey *internalKey;
 @property(nonatomic, strong, readonly, nullable) FSTDocument *internalDocument;
 @property(nonatomic, assign, readonly) BOOL fromCache;
 
@@ -52,11 +56,11 @@ NS_ASSUME_NONNULL_BEGIN
 @implementation FIRDocumentSnapshot (Internal)
 
 + (instancetype)snapshotWithFirestore:(FIRFirestore *)firestore
-                          documentKey:(FSTDocumentKey *)documentKey
+                          documentKey:(DocumentKey)documentKey
                              document:(nullable FSTDocument *)document
                             fromCache:(BOOL)fromCache {
   return [[[self class] alloc] initWithFirestore:firestore
-                                     documentKey:documentKey
+                                     documentKey:std::move(documentKey)
                                         document:document
                                        fromCache:fromCache];
 }
@@ -65,23 +69,28 @@ NS_ASSUME_NONNULL_BEGIN
 
 @implementation FIRDocumentSnapshot {
   FIRSnapshotMetadata *_cachedMetadata;
+  DocumentKey _internalKey;
 }
 
 @dynamic metadata;
 
 - (instancetype)initWithFirestore:(FIRFirestore *)firestore
-                      documentKey:(FSTDocumentKey *)documentKey
+                      documentKey:(DocumentKey)documentKey
                          document:(nullable FSTDocument *)document
                         fromCache:(BOOL)fromCache {
   if (self = [super init]) {
     _firestore = firestore;
-    _internalKey = documentKey;
+    _internalKey = std::move(documentKey);
     _internalDocument = document;
     _fromCache = fromCache;
   }
   return self;
 }
 
+- (const DocumentKey &)internalKey {
+  return _internalKey;
+}
+
 // NSObject Methods
 - (BOOL)isEqual:(nullable id)other {
   if (other == self) return YES;
@@ -95,8 +104,7 @@ NS_ASSUME_NONNULL_BEGIN
   if (self == snapshot) return YES;
   if (snapshot == nil) return NO;
 
-  return [self.firestore isEqual:snapshot.firestore] &&
-         [self.internalKey isEqual:snapshot.internalKey] &&
+  return [self.firestore isEqual:snapshot.firestore] && self.internalKey == snapshot.internalKey &&
          (self.internalDocument == snapshot.internalDocument ||
           [self.internalDocument isEqual:snapshot.internalDocument]) &&
          self.fromCache == snapshot.fromCache;
@@ -104,7 +112,7 @@ NS_ASSUME_NONNULL_BEGIN
 
 - (NSUInteger)hash {
   NSUInteger hash = [self.firestore hash];
-  hash = hash * 31u + [self.internalKey hash];
+  hash = hash * 31u + self.internalKey.Hash();
   hash = hash * 31u + [self.internalDocument hash];
   hash = hash * 31u + (self.fromCache ? 1 : 0);
   return hash;
@@ -121,7 +129,7 @@ NS_ASSUME_NONNULL_BEGIN
 }
 
 - (NSString *)documentID {
-  return util::WrapNSString(self.internalKey.path.last_segment());
+  return util::WrapNSString(self.internalKey.path().last_segment());
 }
 
 - (FIRSnapshotMetadata *)metadata {
@@ -221,7 +229,7 @@ NS_ASSUME_NONNULL_BEGIN
 @interface FIRQueryDocumentSnapshot ()
 
 - (instancetype)initWithFirestore:(FIRFirestore *)firestore
-                      documentKey:(FSTDocumentKey *)documentKey
+                      documentKey:(DocumentKey)documentKey
                          document:(FSTDocument *)document
                         fromCache:(BOOL)fromCache NS_DESIGNATED_INITIALIZER;
 
@@ -230,11 +238,11 @@ NS_ASSUME_NONNULL_BEGIN
 @implementation FIRQueryDocumentSnapshot
 
 - (instancetype)initWithFirestore:(FIRFirestore *)firestore
-                      documentKey:(FSTDocumentKey *)documentKey
+                      documentKey:(DocumentKey)documentKey
                          document:(FSTDocument *)document
                         fromCache:(BOOL)fromCache {
   self = [super initWithFirestore:firestore
-                      documentKey:documentKey
+                      documentKey:std::move(documentKey)
                          document:document
                         fromCache:fromCache];
   return self;
diff --git a/Firestore/Source/API/FIRFirestore.mm b/Firestore/Source/API/FIRFirestore.mm
index f3769ed..f04bd8b 100644
--- a/Firestore/Source/API/FIRFirestore.mm
+++ b/Firestore/Source/API/FIRFirestore.mm
@@ -33,7 +33,6 @@
 #import "Firestore/Source/API/FSTUserDataConverter.h"
 
 #import "Firestore/Source/Core/FSTFirestoreClient.h"
-#import "Firestore/Source/Model/FSTDocumentKey.h"
 #import "Firestore/Source/Util/FSTAssert.h"
 #import "Firestore/Source/Util/FSTDispatchQueue.h"
 #import "Firestore/Source/Util/FSTLogger.h"
diff --git a/Firestore/Source/API/FIRQuery.mm b/Firestore/Source/API/FIRQuery.mm
index 07dac39..9cdc572 100644
--- a/Firestore/Source/API/FIRQuery.mm
+++ b/Firestore/Source/API/FIRQuery.mm
@@ -31,17 +31,18 @@
 #import "Firestore/Source/Core/FSTFirestoreClient.h"
 #import "Firestore/Source/Core/FSTQuery.h"
 #import "Firestore/Source/Model/FSTDocument.h"
-#import "Firestore/Source/Model/FSTDocumentKey.h"
 #import "Firestore/Source/Model/FSTFieldValue.h"
 #import "Firestore/Source/Util/FSTAssert.h"
 #import "Firestore/Source/Util/FSTAsyncQueryListener.h"
 #import "Firestore/Source/Util/FSTUsageValidation.h"
 
+#include "Firestore/core/src/firebase/firestore/model/document_key.h"
 #include "Firestore/core/src/firebase/firestore/model/field_path.h"
 #include "Firestore/core/src/firebase/firestore/model/resource_path.h"
 #include "Firestore/core/src/firebase/firestore/util/string_apple.h"
 
 namespace util = firebase::firestore::util;
+using firebase::firestore::model::DocumentKey;
 using firebase::firestore::model::FieldPath;
 using firebase::firestore::model::ResourcePath;
 
@@ -468,8 +469,8 @@ addSnapshotListenerInternalWithOptions:(FSTListenOptions *)internalOptions
              "a valid document ID, but it was an empty string.");
       }
       ResourcePath path = self.query.path.Append([documentKey UTF8String]);
-      fieldValue = [FSTReferenceValue referenceValue:[FSTDocumentKey keyWithPath:path]
-                                          databaseID:self.firestore.databaseID];
+      fieldValue =
+          [FSTReferenceValue referenceValue:DocumentKey{path} databaseID:self.firestore.databaseID];
     } else if ([value isKindOfClass:[FIRDocumentReference class]]) {
       FIRDocumentReference *ref = (FIRDocumentReference *)value;
       fieldValue = [FSTReferenceValue referenceValue:ref.key databaseID:self.firestore.databaseID];
@@ -615,8 +616,7 @@ addSnapshotListenerInternalWithOptions:(FSTListenOptions *)internalOptions
         FSTThrowInvalidUsage(@"InvalidQueryException",
                              @"Invalid query. Document ID '%@' contains a slash.", documentID);
       }
-      FSTDocumentKey *key =
-          [FSTDocumentKey keyWithPath:self.query.path.Append([documentID UTF8String])];
+      const DocumentKey key{self.query.path.Append([documentID UTF8String])};
       [components
           addObject:[FSTReferenceValue referenceValue:key databaseID:self.firestore.databaseID]];
     } else {
diff --git a/Firestore/Source/API/FSTUserDataConverter.h b/Firestore/Source/API/FSTUserDataConverter.h
index 1058848..3b178be 100644
--- a/Firestore/Source/API/FSTUserDataConverter.h
+++ b/Firestore/Source/API/FSTUserDataConverter.h
@@ -17,9 +17,9 @@
 #import <Foundation/Foundation.h>
 
 #include "Firestore/core/src/firebase/firestore/model/database_id.h"
+#include "Firestore/core/src/firebase/firestore/model/document_key.h"
 
 @class FIRSetOptions;
-@class FSTDocumentKey;
 @class FSTObjectValue;
 @class FSTFieldMask;
 @class FSTFieldValue;
@@ -48,7 +48,7 @@ NS_ASSUME_NONNULL_BEGIN
  * Converts the parsed document data into 1 or 2 mutations (depending on whether there are any
  * field transforms) using the specified document key and precondition.
  */
-- (NSArray<FSTMutation *> *)mutationsWithKey:(FSTDocumentKey *)key
+- (NSArray<FSTMutation *> *)mutationsWithKey:(const firebase::firestore::model::DocumentKey &)key
                                 precondition:(FSTPrecondition *)precondition;
 
 @end
@@ -71,7 +71,7 @@ NS_ASSUME_NONNULL_BEGIN
  * Converts the parsed update data into 1 or 2 mutations (depending on whether there are any
  * field transforms) using the specified document key and precondition.
  */
-- (NSArray<FSTMutation *> *)mutationsWithKey:(FSTDocumentKey *)key
+- (NSArray<FSTMutation *> *)mutationsWithKey:(const firebase::firestore::model::DocumentKey &)key
                                 precondition:(FSTPrecondition *)precondition;
 
 @end
@@ -81,17 +81,18 @@ NS_ASSUME_NONNULL_BEGIN
  * This is necessary because keys assume a database from context (usually the current one).
  * FSTDocumentKeyReference binds a key to a specific databaseID.
  *
- * TODO(b/64160088): Make FSTDocumentKey aware of the specific databaseID it is tied to.
+ * TODO(b/64160088): Make DocumentKey aware of the specific databaseID it is tied to.
  */
 @interface FSTDocumentKeyReference : NSObject
 
 - (instancetype)init NS_UNAVAILABLE;
 
-- (instancetype)initWithKey:(FSTDocumentKey *)key
+- (instancetype)initWithKey:(firebase::firestore::model::DocumentKey)key
                  databaseID:(const firebase::firestore::model::DatabaseId *)databaseID
     NS_DESIGNATED_INITIALIZER;
 
-@property(nonatomic, strong, readonly) FSTDocumentKey *key;
+- (const firebase::firestore::model::DocumentKey &)key;
+
 // Does not own the DatabaseId instance.
 @property(nonatomic, assign, readonly) const firebase::firestore::model::DatabaseId *databaseID;
 
diff --git a/Firestore/Source/API/FSTUserDataConverter.mm b/Firestore/Source/API/FSTUserDataConverter.mm
index e418996..7ee16de 100644
--- a/Firestore/Source/API/FSTUserDataConverter.mm
+++ b/Firestore/Source/API/FSTUserDataConverter.mm
@@ -28,19 +28,20 @@
 #import "Firestore/Source/API/FIRFieldValue+Internal.h"
 #import "Firestore/Source/API/FIRFirestore+Internal.h"
 #import "Firestore/Source/API/FIRSetOptions+Internal.h"
-#import "Firestore/Source/Model/FSTDocumentKey.h"
 #import "Firestore/Source/Model/FSTFieldValue.h"
 #import "Firestore/Source/Model/FSTMutation.h"
 #import "Firestore/Source/Util/FSTAssert.h"
 #import "Firestore/Source/Util/FSTUsageValidation.h"
 
 #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/field_path.h"
 #include "Firestore/core/src/firebase/firestore/util/string_apple.h"
 #include "absl/memory/memory.h"
 
 namespace util = firebase::firestore::util;
 using firebase::firestore::model::DatabaseId;
+using firebase::firestore::model::DocumentKey;
 using firebase::firestore::model::FieldPath;
 
 NS_ASSUME_NONNULL_BEGIN
@@ -62,7 +63,7 @@ static NSString *const RESERVED_FIELD_DESIGNATOR = @"__";
   return self;
 }
 
-- (NSArray<FSTMutation *> *)mutationsWithKey:(FSTDocumentKey *)key
+- (NSArray<FSTMutation *> *)mutationsWithKey:(const DocumentKey &)key
                                 precondition:(FSTPrecondition *)precondition {
   NSMutableArray<FSTMutation *> *mutations = [NSMutableArray array];
   if (self.fieldMask) {
@@ -99,7 +100,7 @@ static NSString *const RESERVED_FIELD_DESIGNATOR = @"__";
   return self;
 }
 
-- (NSArray<FSTMutation *> *)mutationsWithKey:(FSTDocumentKey *)key
+- (NSArray<FSTMutation *> *)mutationsWithKey:(const DocumentKey &)key
                                 precondition:(FSTPrecondition *)precondition {
   NSMutableArray<FSTMutation *> *mutations = [NSMutableArray array];
   [mutations addObject:[[FSTPatchMutation alloc] initWithKey:key
@@ -312,17 +313,23 @@ typedef NS_ENUM(NSInteger, FSTUserDataSource) {
 
 #pragma mark - FSTDocumentKeyReference
 
-@implementation FSTDocumentKeyReference
+@implementation FSTDocumentKeyReference {
+  DocumentKey _key;
+}
 
-- (instancetype)initWithKey:(FSTDocumentKey *)key databaseID:(const DatabaseId *)databaseID {
+- (instancetype)initWithKey:(DocumentKey)key databaseID:(const DatabaseId *)databaseID {
   self = [super init];
   if (self) {
-    _key = key;
+    _key = std::move(key);
     _databaseID = databaseID;
   }
   return self;
 }
 
+- (const firebase::firestore::model::DocumentKey &)key {
+  return _key;
+}
+
 @end
 
 #pragma mark - FSTUserDataConverter
diff --git a/Firestore/Source/Core/FSTQuery.h b/Firestore/Source/Core/FSTQuery.h
index 8da0878..3a67e7f 100644
--- a/Firestore/Source/Core/FSTQuery.h
+++ b/Firestore/Source/Core/FSTQuery.h
@@ -20,7 +20,6 @@
 #include "Firestore/core/src/firebase/firestore/model/resource_path.h"
 
 @class FSTDocument;
-@class FSTDocumentKey;
 @class FSTFieldValue;
 
 NS_ASSUME_NONNULL_BEGIN
diff --git a/Firestore/Source/Core/FSTQuery.mm b/Firestore/Source/Core/FSTQuery.mm
index 626bbb6..811ad03 100644
--- a/Firestore/Source/Core/FSTQuery.mm
+++ b/Firestore/Source/Core/FSTQuery.mm
@@ -21,15 +21,16 @@
 
 #import "Firestore/Source/API/FIRFirestore+Internal.h"
 #import "Firestore/Source/Model/FSTDocument.h"
-#import "Firestore/Source/Model/FSTDocumentKey.h"
 #import "Firestore/Source/Model/FSTFieldValue.h"
 #import "Firestore/Source/Util/FSTAssert.h"
 
+#include "Firestore/core/src/firebase/firestore/model/document_key.h"
 #include "Firestore/core/src/firebase/firestore/model/field_path.h"
 #include "Firestore/core/src/firebase/firestore/model/resource_path.h"
 #include "Firestore/core/src/firebase/firestore/util/string_apple.h"
 
 namespace util = firebase::firestore::util;
+using firebase::firestore::model::DocumentKey;
 using firebase::firestore::model::FieldPath;
 using firebase::firestore::model::ResourcePath;
 
@@ -613,7 +614,7 @@ NSString *FSTStringFromQueryRelationOperator(FSTRelationFilterOperator filterOpe
 }
 
 - (instancetype)queryByAddingFilter:(id<FSTFilter>)filter {
-  FSTAssert(![FSTDocumentKey isDocumentKey:_path], @"No filtering allowed for document query");
+  FSTAssert(!DocumentKey::IsDocumentKey(_path), @"No filtering allowed for document query");
 
   const FieldPath *newInequalityField = nullptr;
   if ([filter isKindOfClass:[FSTRelationFilter class]] &&
@@ -634,7 +635,7 @@ NSString *FSTStringFromQueryRelationOperator(FSTRelationFilterOperator filterOpe
 }
 
 - (instancetype)queryByAddingSortOrder:(FSTSortOrder *)sortOrder {
-  FSTAssert(![FSTDocumentKey isDocumentKey:_path], @"No ordering is allowed for a document query.");
+  FSTAssert(!DocumentKey::IsDocumentKey(_path), @"No ordering is allowed for a document query.");
 
   // TODO(klimt): Validate that the same key isn't added twice.
   return [[FSTQuery alloc] initWithPath:self.path
@@ -673,7 +674,7 @@ NSString *FSTStringFromQueryRelationOperator(FSTRelationFilterOperator filterOpe
 }
 
 - (BOOL)isDocumentQuery {
-  return [FSTDocumentKey isDocumentKey:_path] && self.filters.count == 0;
+  return DocumentKey::IsDocumentKey(_path) && self.filters.count == 0;
 }
 
 - (BOOL)matchesDocument:(FSTDocument *)document {
@@ -768,7 +769,7 @@ NSString *FSTStringFromQueryRelationOperator(FSTRelationFilterOperator filterOpe
 /* Returns YES if the document matches the path for the receiver. */
 - (BOOL)pathMatchesDocument:(FSTDocument *)document {
   const ResourcePath &documentPath = document.key.path();
-  if ([FSTDocumentKey isDocumentKey:_path]) {
+  if (DocumentKey::IsDocumentKey(_path)) {
     // Exact match for document queries.
     return self.path == documentPath;
   } else {
diff --git a/Firestore/Source/Core/FSTSyncEngine.mm b/Firestore/Source/Core/FSTSyncEngine.mm
index b1f138a..d834cc2 100644
--- a/Firestore/Source/Core/FSTSyncEngine.mm
+++ b/Firestore/Source/Core/FSTSyncEngine.mm
@@ -16,6 +16,7 @@
 
 #import "Firestore/Source/Core/FSTSyncEngine.h"
 
+#include <map>
 #include <unordered_map>
 
 #import <GRPCClient/GRPCCall.h>
@@ -33,7 +34,6 @@
 #import "Firestore/Source/Local/FSTQueryData.h"
 #import "Firestore/Source/Local/FSTReferenceSet.h"
 #import "Firestore/Source/Model/FSTDocument.h"
-#import "Firestore/Source/Model/FSTDocumentKey.h"
 #import "Firestore/Source/Model/FSTDocumentSet.h"
 #import "Firestore/Source/Model/FSTMutationBatch.h"
 #import "Firestore/Source/Remote/FSTRemoteEvent.h"
@@ -49,6 +49,7 @@ using firebase::firestore::auth::HashUser;
 using firebase::firestore::auth::User;
 using firebase::firestore::core::TargetIdGenerator;
 using firebase::firestore::model::DocumentKey;
+using firebase::firestore::model::TargetId;
 
 NS_ASSUME_NONNULL_BEGIN
 
@@ -128,17 +129,6 @@ static const FSTListenSequenceNumber kIrrelevantSequenceNumber = -1;
 @property(nonatomic, strong, readonly)
     NSMutableDictionary<NSNumber *, FSTQueryView *> *queryViewsByTarget;
 
-/**
- * When a document is in limbo, we create a special listen to resolve it. This maps the
- * FSTDocumentKey of each limbo document to the FSTTargetID of the listen resolving it.
- */
-@property(nonatomic, strong, readonly)
-    NSMutableDictionary<FSTDocumentKey *, FSTBoxedTargetID *> *limboTargetsByKey;
-
-/** The inverse of limboTargetsByKey, a map of FSTTargetID to the key of the limbo doc. */
-@property(nonatomic, strong, readonly)
-    NSMutableDictionary<FSTBoxedTargetID *, FSTDocumentKey *> *limboKeysByTarget;
-
 /** Used to track any documents that are currently in limbo. */
 @property(nonatomic, strong, readonly) FSTReferenceSet *limboDocumentRefs;
 
@@ -155,6 +145,15 @@ static const FSTListenSequenceNumber kIrrelevantSequenceNumber = -1;
   std::unordered_map<User, NSMutableDictionary<NSNumber *, FSTVoidErrorBlock> *, HashUser>
       _mutationCompletionBlocks;
 
+  /**
+   * When a document is in limbo, we create a special listen to resolve it. This maps the
+   * DocumentKey of each limbo document to the TargetId of the listen resolving it.
+   */
+  std::map<DocumentKey, TargetId> _limboTargetsByKey;
+
+  /** The inverse of _limboTargetsByKey, a map of TargetId to the key of the limbo doc. */
+  std::map<TargetId, DocumentKey> _limboKeysByTarget;
+
   User _currentUser;
 }
 
@@ -168,8 +167,6 @@ static const FSTListenSequenceNumber kIrrelevantSequenceNumber = -1;
     _queryViewsByQuery = [NSMutableDictionary dictionary];
     _queryViewsByTarget = [NSMutableDictionary dictionary];
 
-    _limboTargetsByKey = [NSMutableDictionary dictionary];
-    _limboKeysByTarget = [NSMutableDictionary dictionary];
     _limboCollector = [[FSTEagerGarbageCollector alloc] init];
     _limboDocumentRefs = [[FSTReferenceSet alloc] init];
     [_limboCollector addGarbageSource:_limboDocumentRefs];
@@ -298,8 +295,12 @@ static const FSTListenSequenceNumber kIrrelevantSequenceNumber = -1;
   [remoteEvent.targetChanges enumerateKeysAndObjectsUsingBlock:^(
                                  FSTBoxedTargetID *_Nonnull targetID,
                                  FSTTargetChange *_Nonnull targetChange, BOOL *_Nonnull stop) {
-    FSTDocumentKey *limboKey = self.limboKeysByTarget[targetID];
-    if (limboKey && targetChange.currentStatusUpdate == FSTCurrentStatusUpdateMarkCurrent &&
+    const auto iter = _limboKeysByTarget.find([targetID intValue]);
+    if (iter == _limboKeysByTarget.end()) {
+      return;
+    }
+    const DocumentKey &limboKey = iter->second;
+    if (targetChange.currentStatusUpdate == FSTCurrentStatusUpdateMarkCurrent &&
         remoteEvent.documentUpdates.find(limboKey) == remoteEvent.documentUpdates.end()) {
       // When listening to a query the server responds with a snapshot containing documents
       // matching the query and a current marker telling us we're now in sync. It's possible for
@@ -344,15 +345,16 @@ static const FSTListenSequenceNumber kIrrelevantSequenceNumber = -1;
   [self.delegate handleViewSnapshots:newViewSnapshots];
 }
 
-- (void)rejectListenWithTargetID:(FSTBoxedTargetID *)targetID error:(NSError *)error {
+- (void)rejectListenWithTargetID:(const TargetId)targetID error:(NSError *)error {
   [self assertDelegateExistsForSelector:_cmd];
 
-  FSTDocumentKey *limboKey = self.limboKeysByTarget[targetID];
-  if (limboKey) {
+  const auto iter = _limboKeysByTarget.find(targetID);
+  if (iter != _limboKeysByTarget.end()) {
+    const DocumentKey limboKey = iter->second;
     // Since this query failed, we won't want to manually unlisten to it.
     // So go ahead and remove it from bookkeeping.
-    [self.limboTargetsByKey removeObjectForKey:limboKey];
-    [self.limboKeysByTarget removeObjectForKey:targetID];
+    _limboTargetsByKey.erase(limboKey);
+    _limboKeysByTarget.erase(targetID);
 
     // TODO(dimond): Retry on transient errors?
 
@@ -368,8 +370,8 @@ static const FSTListenSequenceNumber kIrrelevantSequenceNumber = -1;
                                                      documentUpdates:{{limboKey, doc}}];
     [self applyRemoteEvent:event];
   } else {
-    FSTQueryView *queryView = self.queryViewsByTarget[targetID];
-    FSTAssert(queryView, @"Unknown targetId: %@", targetID);
+    FSTQueryView *queryView = self.queryViewsByTarget[@(targetID)];
+    FSTAssert(queryView, @"Unknown targetId: %d", targetID);
     [self.localStore releaseQuery:queryView.query];
     [self removeAndCleanupQuery:queryView];
     [self.delegate handleError:error forQuery:queryView.query];
@@ -479,7 +481,7 @@ static const FSTListenSequenceNumber kIrrelevantSequenceNumber = -1;
         break;
 
       case FSTLimboDocumentChangeTypeRemoved:
-        FSTLog(@"Document no longer in limbo: %@", limboChange.key);
+        FSTLog(@"Document no longer in limbo: %s", limboChange.key.ToString().c_str());
         [self.limboDocumentRefs removeReferenceToKey:limboChange.key forID:targetID];
         break;
 
@@ -491,19 +493,19 @@ static const FSTListenSequenceNumber kIrrelevantSequenceNumber = -1;
 }
 
 - (void)trackLimboChange:(FSTLimboDocumentChange *)limboChange {
-  FSTDocumentKey *key = limboChange.key;
+  DocumentKey key{limboChange.key};
 
-  if (!self.limboTargetsByKey[key]) {
-    FSTLog(@"New document in limbo: %@", key);
-    FSTTargetID limboTargetID = _targetIdGenerator.NextId();
-    FSTQuery *query = [FSTQuery queryWithPath:key.path];
+  if (_limboTargetsByKey.find(key) == _limboTargetsByKey.end()) {
+    FSTLog(@"New document in limbo: %s", key.ToString().c_str());
+    TargetId limboTargetID = _targetIdGenerator.NextId();
+    FSTQuery *query = [FSTQuery queryWithPath:key.path()];
     FSTQueryData *queryData = [[FSTQueryData alloc] initWithQuery:query
                                                          targetID:limboTargetID
                                              listenSequenceNumber:kIrrelevantSequenceNumber
                                                           purpose:FSTQueryPurposeLimboResolution];
-    self.limboKeysByTarget[@(limboTargetID)] = key;
+    _limboKeysByTarget[limboTargetID] = key;
     [self.remoteStore listenToTargetWithQueryData:queryData];
-    self.limboTargetsByKey[key] = @(limboTargetID);
+    _limboTargetsByKey[key] = limboTargetID;
   }
 }
 
@@ -511,22 +513,22 @@ static const FSTListenSequenceNumber kIrrelevantSequenceNumber = -1;
 - (void)garbageCollectLimboDocuments {
   const std::set<DocumentKey> garbage = [self.limboCollector collectGarbage];
   for (const DocumentKey &key : garbage) {
-    FSTBoxedTargetID *limboTarget = self.limboTargetsByKey[static_cast<FSTDocumentKey *>(key)];
-    if (!limboTarget) {
+    const auto iter = _limboTargetsByKey.find(key);
+    if (iter == _limboTargetsByKey.end()) {
       // This target already got removed, because the query failed.
       return;
     }
-    FSTTargetID limboTargetID = limboTarget.intValue;
+    TargetId limboTargetID = iter->second;
     [self.remoteStore stopListeningToTargetID:limboTargetID];
-    [self.limboTargetsByKey removeObjectForKey:key];
-    [self.limboKeysByTarget removeObjectForKey:limboTarget];
+    _limboTargetsByKey.erase(key);
+    _limboKeysByTarget.erase(limboTargetID);
   }
 }
 
 // Used for testing
-- (NSDictionary<FSTDocumentKey *, FSTBoxedTargetID *> *)currentLimboDocuments {
+- (std::map<DocumentKey, TargetId>)currentLimboDocuments {
   // Return defensive copy
-  return [self.limboTargetsByKey copy];
+  return _limboTargetsByKey;
 }
 
 - (void)userDidChange:(const User &)user {
diff --git a/Firestore/Source/Core/FSTTransaction.h b/Firestore/Source/Core/FSTTransaction.h
index 4e5441b..676ada9 100644
--- a/Firestore/Source/Core/FSTTransaction.h
+++ b/Firestore/Source/Core/FSTTransaction.h
@@ -24,7 +24,6 @@
 
 @class FIRSetOptions;
 @class FSTDatastore;
-@class FSTDocumentKey;
 @class FSTFieldMask;
 @class FSTFieldTransform;
 @class FSTMaybeDocument;
@@ -53,18 +52,20 @@ NS_ASSUME_NONNULL_BEGIN
  * Stores mutation for the given key and set data, to be committed when commitWithCompletion is
  * called.
  */
-- (void)setData:(FSTParsedSetData *)data forDocument:(FSTDocumentKey *)key;
+- (void)setData:(FSTParsedSetData *)data
+    forDocument:(const firebase::firestore::model::DocumentKey &)key;
 
 /**
  * Stores mutations for the given key and update data, to be committed when commitWithCompletion
  * is called.
  */
-- (void)updateData:(FSTParsedUpdateData *)data forDocument:(FSTDocumentKey *)key;
+- (void)updateData:(FSTParsedUpdateData *)data
+       forDocument:(const firebase::firestore::model::DocumentKey &)key;
 
 /**
  * Stores a delete mutation for the given key, to be committed when commitWithCompletion is called.
  */
-- (void)deleteDocument:(FSTDocumentKey *)key;
+- (void)deleteDocument:(const firebase::firestore::model::DocumentKey &)key;
 
 /**
  * Attempts to commit the mutations set on this transaction. Calls the given completion block when
diff --git a/Firestore/Source/Core/FSTTransaction.mm b/Firestore/Source/Core/FSTTransaction.mm
index 58f2dd7..9e67ed4 100644
--- a/Firestore/Source/Core/FSTTransaction.mm
+++ b/Firestore/Source/Core/FSTTransaction.mm
@@ -16,6 +16,7 @@
 
 #import "Firestore/Source/Core/FSTTransaction.h"
 
+#include <map>
 #include <vector>
 
 #import <GRPCClient/GRPCCall.h>
@@ -25,7 +26,6 @@
 #import "Firestore/Source/API/FSTUserDataConverter.h"
 #import "Firestore/Source/Core/FSTSnapshotVersion.h"
 #import "Firestore/Source/Model/FSTDocument.h"
-#import "Firestore/Source/Model/FSTDocumentKey.h"
 #import "Firestore/Source/Model/FSTDocumentKeySet.h"
 #import "Firestore/Source/Model/FSTMutation.h"
 #import "Firestore/Source/Remote/FSTDatastore.h"
@@ -42,8 +42,6 @@ NS_ASSUME_NONNULL_BEGIN
 
 @interface FSTTransaction ()
 @property(nonatomic, strong, readonly) FSTDatastore *datastore;
-@property(nonatomic, strong, readonly)
-    NSMutableDictionary<FSTDocumentKey *, FSTSnapshotVersion *> *readVersions;
 @property(nonatomic, strong, readonly) NSMutableArray *mutations;
 @property(nonatomic, assign) BOOL commitCalled;
 /**
@@ -53,7 +51,9 @@ NS_ASSUME_NONNULL_BEGIN
 @property(nonatomic, strong, nullable) NSError *lastWriteError;
 @end
 
-@implementation FSTTransaction
+@implementation FSTTransaction {
+  std::map<DocumentKey, FSTSnapshotVersion *> _readVersions;
+}
 
 + (instancetype)transactionWithDatastore:(FSTDatastore *)datastore {
   return [[FSTTransaction alloc] initWithDatastore:datastore];
@@ -63,7 +63,6 @@ NS_ASSUME_NONNULL_BEGIN
   self = [super init];
   if (self) {
     _datastore = datastore;
-    _readVersions = [NSMutableDictionary dictionary];
     _mutations = [NSMutableArray array];
     _commitCalled = NO;
   }
@@ -85,8 +84,10 @@ NS_ASSUME_NONNULL_BEGIN
     // when writing.
     docVersion = [FSTSnapshotVersion noVersion];
   }
-  FSTSnapshotVersion *existingVersion = self.readVersions[(FSTDocumentKey *)doc.key];
-  if (existingVersion) {
+  if (_readVersions.find(doc.key) == _readVersions.end()) {
+    _readVersions[doc.key] = docVersion;
+    return YES;
+  } else {
     if (error) {
       *error =
           [NSError errorWithDomain:FIRFirestoreErrorDomain
@@ -97,9 +98,6 @@ NS_ASSUME_NONNULL_BEGIN
                           }];
     }
     return NO;
-  } else {
-    self.readVersions[(FSTDocumentKey *)doc.key] = docVersion;
-    return YES;
   }
 }
 
@@ -138,12 +136,12 @@ NS_ASSUME_NONNULL_BEGIN
  * Returns version of this doc when it was read in this transaction as a precondition, or no
  * precondition if it was not read.
  */
-- (FSTPrecondition *)preconditionForDocumentKey:(FSTDocumentKey *)key {
-  FSTSnapshotVersion *_Nullable snapshotVersion = self.readVersions[key];
-  if (snapshotVersion) {
-    return [FSTPrecondition preconditionWithUpdateTime:snapshotVersion];
-  } else {
+- (FSTPrecondition *)preconditionForDocumentKey:(const DocumentKey &)key {
+  const auto iter = _readVersions.find(key);
+  if (iter == _readVersions.end()) {
     return [FSTPrecondition none];
+  } else {
+    return [FSTPrecondition preconditionWithUpdateTime:iter->second];
   }
 }
 
@@ -151,10 +149,16 @@ NS_ASSUME_NONNULL_BEGIN
  * Returns the precondition for a document if the operation is an update, based on the provided
  * UpdateOptions. Will return nil if an error occurred, in which case it sets the error parameter.
  */
-- (nullable FSTPrecondition *)preconditionForUpdateWithDocumentKey:(FSTDocumentKey *)key
+- (nullable FSTPrecondition *)preconditionForUpdateWithDocumentKey:(const DocumentKey &)key
                                                              error:(NSError **)error {
-  FSTSnapshotVersion *_Nullable version = self.readVersions[key];
-  if (version && [version isEqual:[FSTSnapshotVersion noVersion]]) {
+  const auto iter = _readVersions.find(key);
+  if (iter == _readVersions.end()) {
+    // Document was not read, so we just use the preconditions for an update.
+    return [FSTPrecondition preconditionWithExists:YES];
+  }
+
+  FSTSnapshotVersion *version = iter->second;
+  if ([version isEqual:[FSTSnapshotVersion noVersion]]) {
     // The document was read, but doesn't exist.
     // Return an error because the precondition is impossible
     if (error) {
@@ -166,21 +170,18 @@ NS_ASSUME_NONNULL_BEGIN
                  }];
     }
     return nil;
-  } else if (version) {
+  } else {
     // Document exists, just base precondition on document update time.
     return [FSTPrecondition preconditionWithUpdateTime:version];
-  } else {
-    // Document was not read, so we just use the preconditions for an update.
-    return [FSTPrecondition preconditionWithExists:YES];
   }
 }
 
-- (void)setData:(FSTParsedSetData *)data forDocument:(FSTDocumentKey *)key {
+- (void)setData:(FSTParsedSetData *)data forDocument:(const DocumentKey &)key {
   [self writeMutations:[data mutationsWithKey:key
                                  precondition:[self preconditionForDocumentKey:key]]];
 }
 
-- (void)updateData:(FSTParsedUpdateData *)data forDocument:(FSTDocumentKey *)key {
+- (void)updateData:(FSTParsedUpdateData *)data forDocument:(const DocumentKey &)key {
   NSError *error = nil;
   FSTPrecondition *_Nullable precondition =
       [self preconditionForUpdateWithDocumentKey:key error:&error];
@@ -192,13 +193,13 @@ NS_ASSUME_NONNULL_BEGIN
   }
 }
 
-- (void)deleteDocument:(FSTDocumentKey *)key {
+- (void)deleteDocument:(const DocumentKey &)key {
   [self writeMutations:@[ [[FSTDeleteMutation alloc]
                             initWithKey:key
                            precondition:[self preconditionForDocumentKey:key]] ]];
   // Since the delete will be applied before all following writes, we need to ensure that the
   // precondition for the next write will be exists: false.
-  self.readVersions[key] = [FSTSnapshotVersion noVersion];
+  _readVersions[key] = [FSTSnapshotVersion noVersion];
 }
 
 - (void)commitWithCompletion:(FSTVoidErrorBlock)completion {
@@ -213,11 +214,10 @@ NS_ASSUME_NONNULL_BEGIN
   }
 
   // Make a list of read documents that haven't been written.
-  __block FSTDocumentKeySet *unwritten = [FSTDocumentKeySet keySet];
-  [self.readVersions enumerateKeysAndObjectsUsingBlock:^(FSTDocumentKey *key,
-                                                         FSTSnapshotVersion *version, BOOL *stop) {
-    unwritten = [unwritten setByAddingObject:key];
-  }];
+  FSTDocumentKeySet *unwritten = [FSTDocumentKeySet keySet];
+  for (const auto &kv : _readVersions) {
+    unwritten = [unwritten setByAddingObject:kv.first];
+  };
   // For each mutation, note that the doc was written.
   for (FSTMutation *mutation in self.mutations) {
     unwritten = [unwritten setByRemovingObject:mutation.key];
diff --git a/Firestore/Source/Core/FSTView.h b/Firestore/Source/Core/FSTView.h
index 6ff77cd..38f57fc 100644
--- a/Firestore/Source/Core/FSTView.h
+++ b/Firestore/Source/Core/FSTView.h
@@ -20,7 +20,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 FSTDocumentSet;
 @class FSTDocumentViewChangeSet;
 @class FSTMaybeDocument;
@@ -63,12 +64,14 @@ typedef NS_ENUM(NSInteger, FSTLimboDocumentChangeType) {
 // A change to a particular document wrt to whether it is in "limbo".
 @interface FSTLimboDocumentChange : NSObject
 
-+ (instancetype)changeWithType:(FSTLimboDocumentChangeType)type key:(FSTDocumentKey *)key;
++ (instancetype)changeWithType:(FSTLimboDocumentChangeType)type
+                           key:(firebase::firestore::model::DocumentKey)key;
 
 - (id)init __attribute__((unavailable("Use a static constructor method.")));
 
+- (const firebase::firestore::model::DocumentKey &)key;
+
 @property(nonatomic, assign, readonly) FSTLimboDocumentChangeType type;
-@property(nonatomic, strong, readonly) FSTDocumentKey *key;
 @end
 
 #pragma mark - FSTViewChange
diff --git a/Firestore/Source/Core/FSTView.mm b/Firestore/Source/Core/FSTView.mm
index b3cf189..d87951a 100644
--- a/Firestore/Source/Core/FSTView.mm
+++ b/Firestore/Source/Core/FSTView.mm
@@ -16,15 +16,20 @@
 
 #import "Firestore/Source/Core/FSTView.h"
 
+#include <utility>
+
 #import "Firestore/Source/Core/FSTQuery.h"
 #import "Firestore/Source/Core/FSTViewSnapshot.h"
 #import "Firestore/Source/Model/FSTDocument.h"
-#import "Firestore/Source/Model/FSTDocumentKey.h"
 #import "Firestore/Source/Model/FSTDocumentSet.h"
 #import "Firestore/Source/Model/FSTFieldValue.h"
 #import "Firestore/Source/Remote/FSTRemoteEvent.h"
 #import "Firestore/Source/Util/FSTAssert.h"
 
+#include "Firestore/core/src/firebase/firestore/model/document_key.h"
+
+using firebase::firestore::model::DocumentKey;
+
 NS_ASSUME_NONNULL_BEGIN
 
 #pragma mark - FSTViewDocumentChanges
@@ -61,28 +66,34 @@ NS_ASSUME_NONNULL_BEGIN
 
 @interface FSTLimboDocumentChange ()
 
-+ (instancetype)changeWithType:(FSTLimboDocumentChangeType)type key:(FSTDocumentKey *)key;
++ (instancetype)changeWithType:(FSTLimboDocumentChangeType)type key:(DocumentKey)key;
 
 - (instancetype)initWithType:(FSTLimboDocumentChangeType)type
-                         key:(FSTDocumentKey *)key NS_DESIGNATED_INITIALIZER;
+                         key:(DocumentKey)key NS_DESIGNATED_INITIALIZER;
 
 @end
 
-@implementation FSTLimboDocumentChange
+@implementation FSTLimboDocumentChange {
+  DocumentKey _key;
+}
 
-+ (instancetype)changeWithType:(FSTLimboDocumentChangeType)type key:(FSTDocumentKey *)key {
-  return [[FSTLimboDocumentChange alloc] initWithType:type key:key];
++ (instancetype)changeWithType:(FSTLimboDocumentChangeType)type key:(DocumentKey)key {
+  return [[FSTLimboDocumentChange alloc] initWithType:type key:std::move(key)];
 }
 
-- (instancetype)initWithType:(FSTLimboDocumentChangeType)type key:(FSTDocumentKey *)key {
+- (instancetype)initWithType:(FSTLimboDocumentChangeType)type key:(DocumentKey)key {
   self = [super init];
   if (self) {
     _type = type;
-    _key = key;
+    _key = std::move(key);
   }
   return self;
 }
 
+- (const DocumentKey &)key {
+  return _key;
+}
+
 - (BOOL)isEqual:(id)other {
   if (self == other) {
     return YES;
@@ -357,7 +368,7 @@ static NSComparisonResult FSTCompareDocumentViewChangeTypes(FSTDocumentViewChang
 #pragma mark - Private methods
 
 /** Returns whether the doc for the given key should be in limbo. */
-- (BOOL)shouldBeLimboDocumentKey:(FSTDocumentKey *)key {
+- (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]) {
     return NO;
diff --git a/Firestore/Source/Core/FSTViewSnapshot.mm b/Firestore/Source/Core/FSTViewSnapshot.mm
index e60b785..ae366fb 100644
--- a/Firestore/Source/Core/FSTViewSnapshot.mm
+++ b/Firestore/Source/Core/FSTViewSnapshot.mm
@@ -18,11 +18,14 @@
 
 #import "Firestore/Source/Core/FSTQuery.h"
 #import "Firestore/Source/Model/FSTDocument.h"
-#import "Firestore/Source/Model/FSTDocumentKey.h"
 #import "Firestore/Source/Model/FSTDocumentSet.h"
 #import "Firestore/Source/Util/FSTAssert.h"
 #import "Firestore/third_party/Immutable/FSTImmutableSortedDictionary.h"
 
+#include "Firestore/core/src/firebase/firestore/model/document_key.h"
+
+using firebase::firestore::model::DocumentKey;
+
 NS_ASSUME_NONNULL_BEGIN
 
 #pragma mark - FSTDocumentViewChange
@@ -98,7 +101,7 @@ NS_ASSUME_NONNULL_BEGIN
 }
 
 - (void)addChange:(FSTDocumentViewChange *)change {
-  FSTDocumentKey *key = change.document.key;
+  const DocumentKey &key = change.document.key;
   FSTDocumentViewChange *oldChange = [self.changeMap objectForKey:key];
   if (!oldChange) {
     self.changeMap = [self.changeMap dictionaryBySettingObject:change forKey:key];
diff --git a/Firestore/Source/Model/FSTMutationBatch.mm b/Firestore/Source/Model/FSTMutationBatch.mm
index 6598d27..e62a72c 100644
--- a/Firestore/Source/Model/FSTMutationBatch.mm
+++ b/Firestore/Source/Model/FSTMutationBatch.mm
@@ -20,7 +20,6 @@
 
 #import "Firestore/Source/Core/FSTSnapshotVersion.h"
 #import "Firestore/Source/Model/FSTDocument.h"
-#import "Firestore/Source/Model/FSTDocumentKey.h"
 #import "Firestore/Source/Model/FSTMutation.h"
 #import "Firestore/Source/Util/FSTAssert.h"
 
diff --git a/Firestore/Source/Remote/FSTRemoteStore.h b/Firestore/Source/Remote/FSTRemoteStore.h
index 6f4d565..09e1d32 100644
--- a/Firestore/Source/Remote/FSTRemoteStore.h
+++ b/Firestore/Source/Remote/FSTRemoteStore.h
@@ -22,7 +22,6 @@
 #include "Firestore/core/src/firebase/firestore/auth/user.h"
 
 @class FSTDatastore;
-@class FSTDocumentKey;
 @class FSTLocalStore;
 @class FSTMutationBatch;
 @class FSTMutationBatchResult;
@@ -59,7 +58,8 @@ NS_ASSUME_NONNULL_BEGIN
  *     will be an indication that the user is no longer authorized to see the data matching the
  *     target.
  */
-- (void)rejectListenWithTargetID:(FSTBoxedTargetID *)targetID error:(NSError *)error;
+- (void)rejectListenWithTargetID:(const firebase::firestore::model::TargetId)targetID
+                           error:(NSError *)error;
 
 /**
  * Applies the result of a successful write of a mutation batch to the sync engine, emitting
diff --git a/Firestore/Source/Remote/FSTRemoteStore.mm b/Firestore/Source/Remote/FSTRemoteStore.mm
index 28ee594..8892ffb 100644
--- a/Firestore/Source/Remote/FSTRemoteStore.mm
+++ b/Firestore/Source/Remote/FSTRemoteStore.mm
@@ -465,7 +465,7 @@ static const int kMaxPendingWrites = 10;
   for (FSTBoxedTargetID *targetID in change.targetIDs) {
     if (self.listenTargets[targetID]) {
       [self.listenTargets removeObjectForKey:targetID];
-      [self.syncEngine rejectListenWithTargetID:targetID error:change.cause];
+      [self.syncEngine rejectListenWithTargetID:[targetID intValue] error:change.cause];
     }
   }
 }
diff --git a/Firestore/Source/Remote/FSTStream.h b/Firestore/Source/Remote/FSTStream.h
index e48f1da..fba79d2 100644
--- a/Firestore/Source/Remote/FSTStream.h
+++ b/Firestore/Source/Remote/FSTStream.h
@@ -22,7 +22,6 @@
 #include "Firestore/core/src/firebase/firestore/auth/credentials_provider.h"
 #include "Firestore/core/src/firebase/firestore/core/database_info.h"
 
-@class FSTDocumentKey;
 @class FSTDispatchQueue;
 @class FSTMutation;
 @class FSTMutationResult;
diff --git a/Firestore/Source/Remote/FSTWatchChange.mm b/Firestore/Source/Remote/FSTWatchChange.mm
index 2c1a916..284e980 100644
--- a/Firestore/Source/Remote/FSTWatchChange.mm
+++ b/Firestore/Source/Remote/FSTWatchChange.mm
@@ -70,7 +70,7 @@ NS_ASSUME_NONNULL_BEGIN
 - (NSUInteger)hash {
   NSUInteger hash = self.updatedTargetIDs.hash;
   hash = hash * 31 + self.removedTargetIDs.hash;
-  hash = hash * 31 + std::hash<std::string>{}(self.documentKey.ToString());
+  hash = hash * 31 + self.documentKey.Hash();
   hash = hash * 31 + self.document.hash;
   return hash;
 }
diff --git a/Firestore/core/src/firebase/firestore/model/base_path.h b/Firestore/core/src/firebase/firestore/model/base_path.h
index 73ac611..bc1f89d 100644
--- a/Firestore/core/src/firebase/firestore/model/base_path.h
+++ b/Firestore/core/src/firebase/firestore/model/base_path.h
@@ -162,9 +162,9 @@ class BasePath {
 #if defined(__OBJC__)
   // For Objective-C++ hash; to be removed after migration.
   // Do NOT use in C++ code.
-  uint64_t Hash() const {
+  NSUInteger Hash() const {
     std::hash<std::string> hash_fn;
-    uint64_t hash_result = 0;
+    NSUInteger hash_result = 0;
     for (const std::string& segment : segments_) {
       hash_result = hash_result * 31u + hash_fn(segment);
     }
diff --git a/Firestore/core/src/firebase/firestore/model/database_id.h b/Firestore/core/src/firebase/firestore/model/database_id.h
index 01ccf76..2ad1332 100644
--- a/Firestore/core/src/firebase/firestore/model/database_id.h
+++ b/Firestore/core/src/firebase/firestore/model/database_id.h
@@ -64,7 +64,7 @@ class DatabaseId {
 #if defined(__OBJC__)
   // For objective-c++ hash; to be removed after migration.
   // Do NOT use in C++ code.
-  uint64_t Hash() const {
+  NSUInteger Hash() const {
     std::hash<std::string> hash_fn;
     return hash_fn(project_id_) * 31u + hash_fn(database_id_);
   }
diff --git a/Firestore/core/src/firebase/firestore/model/document_key.h b/Firestore/core/src/firebase/firestore/model/document_key.h
index 134bda5..4bdc04b 100644
--- a/Firestore/core/src/firebase/firestore/model/document_key.h
+++ b/Firestore/core/src/firebase/firestore/model/document_key.h
@@ -59,6 +59,10 @@ class DocumentKey {
   std::string ToString() const {
     return path().CanonicalString();
   }
+
+  NSUInteger Hash() const {
+    return std::hash<std::string>{}(ToString());
+  }
 #endif
 
   /**