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

8 files changed, 113 insertions, 93 deletions

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