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

9 files changed, 100 insertions, 73 deletions

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