port C++ `DocumentKey` to `Remote/*` (#965)

* port C++ DocumentKey to Remote's

* port C++ DocumentKey to Remote's

* address changes

* address changes

* address changes

14 files changed, 144 insertions, 109 deletions

diff --git a/Firestore/Source/API/FIRTransaction.mm b/Firestore/Source/API/FIRTransaction.mm
index 5edff19..2d1dd68 100644
--- a/Firestore/Source/API/FIRTransaction.mm
+++ b/Firestore/Source/API/FIRTransaction.mm
@@ -94,7 +94,7 @@ NS_ASSUME_NONNULL_BEGIN
                               NSError *_Nullable error))completion {
   [self validateReference:document];
   [self.internalTransaction
-      lookupDocumentsForKeys:@[ document.key ]
+      lookupDocumentsForKeys:{document.key}
                   completion:^(NSArray<FSTMaybeDocument *> *_Nullable documents,
                                NSError *_Nullable error) {
                     if (error) {
diff --git a/Firestore/Source/Core/FSTSyncEngine.mm b/Firestore/Source/Core/FSTSyncEngine.mm
index d567878..b1f138a 100644
--- a/Firestore/Source/Core/FSTSyncEngine.mm
+++ b/Firestore/Source/Core/FSTSyncEngine.mm
@@ -300,7 +300,7 @@ static const FSTListenSequenceNumber kIrrelevantSequenceNumber = -1;
                                  FSTTargetChange *_Nonnull targetChange, BOOL *_Nonnull stop) {
     FSTDocumentKey *limboKey = self.limboKeysByTarget[targetID];
     if (limboKey && targetChange.currentStatusUpdate == FSTCurrentStatusUpdateMarkCurrent &&
-        remoteEvent.documentUpdates[limboKey] == nil) {
+        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
       // these to arrive as separate remote events or as a single remote event. For a document
@@ -363,11 +363,9 @@ static const FSTListenSequenceNumber kIrrelevantSequenceNumber = -1;
         [NSMutableDictionary dictionary];
     FSTDeletedDocument *doc =
         [FSTDeletedDocument documentWithKey:limboKey version:[FSTSnapshotVersion noVersion]];
-    NSMutableDictionary<FSTDocumentKey *, FSTMaybeDocument *> *docUpdate =
-        [NSMutableDictionary dictionaryWithObject:doc forKey:limboKey];
     FSTRemoteEvent *event = [FSTRemoteEvent eventWithSnapshotVersion:[FSTSnapshotVersion noVersion]
                                                        targetChanges:targetChanges
-                                                     documentUpdates:docUpdate];
+                                                     documentUpdates:{{limboKey, doc}}];
     [self applyRemoteEvent:event];
   } else {
     FSTQueryView *queryView = self.queryViewsByTarget[targetID];
diff --git a/Firestore/Source/Core/FSTTransaction.h b/Firestore/Source/Core/FSTTransaction.h
index d4f3d8d..4e5441b 100644
--- a/Firestore/Source/Core/FSTTransaction.h
+++ b/Firestore/Source/Core/FSTTransaction.h
@@ -16,8 +16,12 @@
 
 #import <Foundation/Foundation.h>
 
+#include <vector>
+
 #import "Firestore/Source/Core/FSTTypes.h"
 
+#include "Firestore/core/src/firebase/firestore/model/document_key.h"
+
 @class FIRSetOptions;
 @class FSTDatastore;
 @class FSTDocumentKey;
@@ -42,7 +46,7 @@ NS_ASSUME_NONNULL_BEGIN
  * Takes a set of keys and asynchronously attempts to fetch all the documents from the backend,
  * ignoring any local changes.
  */
-- (void)lookupDocumentsForKeys:(NSArray<FSTDocumentKey *> *)keys
+- (void)lookupDocumentsForKeys:(const std::vector<firebase::firestore::model::DocumentKey> &)keys
                     completion:(FSTVoidMaybeDocumentArrayErrorBlock)completion;
 
 /**
diff --git a/Firestore/Source/Core/FSTTransaction.mm b/Firestore/Source/Core/FSTTransaction.mm
index b980a4e..58f2dd7 100644
--- a/Firestore/Source/Core/FSTTransaction.mm
+++ b/Firestore/Source/Core/FSTTransaction.mm
@@ -16,6 +16,8 @@
 
 #import "Firestore/Source/Core/FSTTransaction.h"
 
+#include <vector>
+
 #import <GRPCClient/GRPCCall.h>
 
 #import "FIRFirestoreErrors.h"
@@ -32,6 +34,8 @@
 
 #include "Firestore/core/src/firebase/firestore/model/document_key.h"
 
+using firebase::firestore::model::DocumentKey;
+
 NS_ASSUME_NONNULL_BEGIN
 
 #pragma mark - FSTTransaction
@@ -99,7 +103,7 @@ NS_ASSUME_NONNULL_BEGIN
   }
 }
 
-- (void)lookupDocumentsForKeys:(NSArray<FSTDocumentKey *> *)keys
+- (void)lookupDocumentsForKeys:(const std::vector<DocumentKey> &)keys
                     completion:(FSTVoidMaybeDocumentArrayErrorBlock)completion {
   [self ensureCommitNotCalled];
   if (self.mutations.count) {
diff --git a/Firestore/Source/Local/FSTLocalStore.mm b/Firestore/Source/Local/FSTLocalStore.mm
index 922f281..1bcc6d0 100644
--- a/Firestore/Source/Local/FSTLocalStore.mm
+++ b/Firestore/Source/Local/FSTLocalStore.mm
@@ -47,6 +47,7 @@
 using firebase::firestore::auth::User;
 using firebase::firestore::model::DocumentKey;
 using firebase::firestore::core::TargetIdGenerator;
+using firebase::firestore::model::DocumentKey;
 
 NS_ASSUME_NONNULL_BEGIN
 
@@ -320,28 +321,29 @@ NS_ASSUME_NONNULL_BEGIN
   }];
 
   // TODO(klimt): This could probably be an NSMutableDictionary.
-  __block FSTDocumentKeySet *changedDocKeys = [FSTDocumentKeySet keySet];
-  [remoteEvent.documentUpdates
-      enumerateKeysAndObjectsUsingBlock:^(FSTDocumentKey *key, FSTMaybeDocument *doc, BOOL *stop) {
-        changedDocKeys = [changedDocKeys setByAddingObject:key];
-        FSTMaybeDocument *existingDoc = [remoteDocuments entryForKey:key];
-        // Make sure we don't apply an old document version to the remote cache, though we
-        // make an exception for [SnapshotVersion noVersion] which can happen for manufactured
-        // events (e.g. in the case of a limbo document resolution failing).
-        if (!existingDoc || [doc.version isEqual:[FSTSnapshotVersion noVersion]] ||
-            [doc.version compare:existingDoc.version] != NSOrderedAscending) {
-          [remoteDocuments addEntry:doc];
-        } else {
-          FSTLog(
-              @"FSTLocalStore Ignoring outdated watch update for %@. "
-               "Current version: %@  Watch version: %@",
-              key, existingDoc.version, doc.version);
-        }
-
-        // The document might be garbage because it was unreferenced by everything.
-        // Make sure to mark it as garbage if it is...
-        [self.garbageCollector addPotentialGarbageKey:key];
-      }];
+  FSTDocumentKeySet *changedDocKeys = [FSTDocumentKeySet keySet];
+  for (const auto &kv : remoteEvent.documentUpdates) {
+    const DocumentKey &key = kv.first;
+    FSTMaybeDocument *doc = kv.second;
+    changedDocKeys = [changedDocKeys setByAddingObject:key];
+    FSTMaybeDocument *existingDoc = [remoteDocuments entryForKey:key];
+    // Make sure we don't apply an old document version to the remote cache, though we
+    // make an exception for [SnapshotVersion noVersion] which can happen for manufactured
+    // events (e.g. in the case of a limbo document resolution failing).
+    if (!existingDoc || [doc.version isEqual:[FSTSnapshotVersion noVersion]] ||
+        [doc.version compare:existingDoc.version] != NSOrderedAscending) {
+      [remoteDocuments addEntry:doc];
+    } else {
+      FSTLog(
+          @"FSTLocalStore Ignoring outdated watch update for %s. "
+           "Current version: %@  Watch version: %@",
+          key.ToString().c_str(), existingDoc.version, doc.version);
+    }
+
+    // The document might be garbage because it was unreferenced by everything.
+    // Make sure to mark it as garbage if it is...
+    [self.garbageCollector addPotentialGarbageKey:key];
+  }
 
   // HACK: The only reason we allow omitting snapshot version is so we can synthesize remote events
   // when we get permission denied errors while trying to resolve the state of a locally cached
diff --git a/Firestore/Source/Remote/FSTDatastore.h b/Firestore/Source/Remote/FSTDatastore.h
index 7b8274c..b3ba46c 100644
--- a/Firestore/Source/Remote/FSTDatastore.h
+++ b/Firestore/Source/Remote/FSTDatastore.h
@@ -16,14 +16,16 @@
 
 #import <Foundation/Foundation.h>
 
+#include <vector>
+
 #import "Firestore/Source/Core/FSTTypes.h"
 
 #include "Firestore/core/src/firebase/firestore/auth/credentials_provider.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"
 #include "absl/strings/string_view.h"
 
-@class FSTDocumentKey;
 @class FSTDispatchQueue;
 @class FSTMutation;
 @class FSTMutationResult;
@@ -88,7 +90,7 @@ NS_ASSUME_NONNULL_BEGIN
                        token:(const absl::string_view)token;
 
 /** Looks up a list of documents in datastore. */
-- (void)lookupDocuments:(NSArray<FSTDocumentKey *> *)keys
+- (void)lookupDocuments:(const std::vector<firebase::firestore::model::DocumentKey> &)keys
              completion:(FSTVoidMaybeDocumentArrayErrorBlock)completion;
 
 /** Commits data to datastore. */
diff --git a/Firestore/Source/Remote/FSTDatastore.mm b/Firestore/Source/Remote/FSTDatastore.mm
index 63499df..501237f 100644
--- a/Firestore/Source/Remote/FSTDatastore.mm
+++ b/Firestore/Source/Remote/FSTDatastore.mm
@@ -16,6 +16,9 @@
 
 #import "Firestore/Source/Remote/FSTDatastore.h"
 
+#include <memory>
+#include <vector>
+
 #import <GRPCClient/GRPCCall+OAuth2.h>
 #import <ProtoRPC/ProtoRPC.h>
 
@@ -24,7 +27,6 @@
 #import "Firestore/Source/API/FIRFirestoreVersion.h"
 #import "Firestore/Source/Local/FSTLocalStore.h"
 #import "Firestore/Source/Model/FSTDocument.h"
-#import "Firestore/Source/Model/FSTDocumentKey.h"
 #import "Firestore/Source/Model/FSTMutation.h"
 #import "Firestore/Source/Remote/FSTSerializerBeta.h"
 #import "Firestore/Source/Remote/FSTStream.h"
@@ -38,6 +40,7 @@
 #include "Firestore/core/src/firebase/firestore/auth/token.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"
 #include "Firestore/core/src/firebase/firestore/util/error_apple.h"
 #include "Firestore/core/src/firebase/firestore/util/string_apple.h"
 
@@ -46,6 +49,7 @@ using firebase::firestore::auth::CredentialsProvider;
 using firebase::firestore::auth::Token;
 using firebase::firestore::core::DatabaseInfo;
 using firebase::firestore::model::DatabaseId;
+using firebase::firestore::model::DocumentKey;
 
 NS_ASSUME_NONNULL_BEGIN
 
@@ -245,17 +249,21 @@ typedef GRPCProtoCall * (^RPCFactory)(void);
   [self invokeRPCWithFactory:rpcFactory errorHandler:completion];
 }
 
-- (void)lookupDocuments:(NSArray<FSTDocumentKey *> *)keys
+- (void)lookupDocuments:(const std::vector<DocumentKey> &)keys
              completion:(FSTVoidMaybeDocumentArrayErrorBlock)completion {
   GCFSBatchGetDocumentsRequest *request = [GCFSBatchGetDocumentsRequest message];
   request.database = [self.serializer encodedDatabaseID];
-  for (FSTDocumentKey *key in keys) {
+  for (const DocumentKey &key : keys) {
     [request.documentsArray addObject:[self.serializer encodedDocumentKey:key]];
   }
 
-  __block FSTMaybeDocumentDictionary *results =
-      [FSTMaybeDocumentDictionary maybeDocumentDictionary];
+  struct Closure {
+    std::vector<DocumentKey> keys;
+    FSTMaybeDocumentDictionary *results;
+  };
 
+  __block std::shared_ptr<Closure> closure = std::make_shared<Closure>(
+      Closure{keys, [FSTMaybeDocumentDictionary maybeDocumentDictionary]});
   RPCFactory rpcFactory = ^GRPCProtoCall * {
     __block GRPCProtoCall *rpc = [self.service
         RPCToBatchGetDocumentsWithRequest:request
@@ -275,16 +283,19 @@ typedef GRPCProtoCall * (^RPCFactory)(void);
                                    // Streaming response, accumulate result
                                    FSTMaybeDocument *doc =
                                        [self.serializer decodedMaybeDocumentFromBatch:response];
-                                   results = [results dictionaryBySettingObject:doc forKey:doc.key];
+                                   closure->results =
+                                       [closure->results dictionaryBySettingObject:doc
+                                                                            forKey:doc.key];
                                  } else {
                                    // Streaming response is done, call completion
                                    FSTLog(@"RPC BatchGetDocuments completed successfully.");
                                    [FSTDatastore logHeadersForRPC:rpc RPCName:@"BatchGetDocuments"];
                                    FSTAssert(!response, @"Got response after done.");
                                    NSMutableArray<FSTMaybeDocument *> *docs =
-                                       [NSMutableArray arrayWithCapacity:keys.count];
-                                   for (FSTDocumentKey *key in keys) {
-                                     [docs addObject:results[key]];
+                                       [NSMutableArray arrayWithCapacity:closure->keys.size()];
+                                   for (const DocumentKey &key : closure->keys) {
+                                     [docs addObject:closure->results[static_cast<FSTDocumentKey *>(
+                                                         key)]];
                                    }
                                    completion(docs, nil);
                                  }
diff --git a/Firestore/Source/Remote/FSTRemoteEvent.h b/Firestore/Source/Remote/FSTRemoteEvent.h
index 72f71a5..dd45117 100644
--- a/Firestore/Source/Remote/FSTRemoteEvent.h
+++ b/Firestore/Source/Remote/FSTRemoteEvent.h
@@ -16,12 +16,15 @@
 
 #import <Foundation/Foundation.h>
 
+#include <map>
+
 #import "Firestore/Source/Core/FSTTypes.h"
 #import "Firestore/Source/Model/FSTDocumentDictionary.h"
 #import "Firestore/Source/Model/FSTDocumentKeySet.h"
 
+#include "Firestore/core/src/firebase/firestore/model/document_key.h"
+
 @class FSTDocument;
-@class FSTDocumentKey;
 @class FSTExistenceFilter;
 @class FSTMaybeDocument;
 @class FSTSnapshotVersion;
@@ -148,7 +151,7 @@ typedef NS_ENUM(NSUInteger, FSTCurrentStatusUpdate) {
 eventWithSnapshotVersion:(FSTSnapshotVersion *)snapshotVersion
            targetChanges:(NSMutableDictionary<NSNumber *, FSTTargetChange *> *)targetChanges
          documentUpdates:
-             (NSMutableDictionary<FSTDocumentKey *, FSTMaybeDocument *> *)documentUpdates;
+             (std::map<firebase::firestore::model::DocumentKey, FSTMaybeDocument *>)documentUpdates;
 
 /** The snapshot version this event brings us up to. */
 @property(nonatomic, strong, readonly) FSTSnapshotVersion *snapshotVersion;
@@ -161,8 +164,7 @@ eventWithSnapshotVersion:(FSTSnapshotVersion *)snapshotVersion
  * A set of which documents have changed or been deleted, along with the doc's new values
  * (if not deleted).
  */
-@property(nonatomic, strong, readonly)
-    NSDictionary<FSTDocumentKey *, FSTMaybeDocument *> *documentUpdates;
+- (const std::map<firebase::firestore::model::DocumentKey, FSTMaybeDocument *> &)documentUpdates;
 
 /** Adds a document update to this remote event */
 - (void)addDocumentUpdate:(FSTMaybeDocument *)document;
diff --git a/Firestore/Source/Remote/FSTRemoteEvent.mm b/Firestore/Source/Remote/FSTRemoteEvent.mm
index cd073b3..a3c85df 100644
--- a/Firestore/Source/Remote/FSTRemoteEvent.mm
+++ b/Firestore/Source/Remote/FSTRemoteEvent.mm
@@ -16,9 +16,11 @@
 
 #import "Firestore/Source/Remote/FSTRemoteEvent.h"
 
+#include <map>
+#include <utility>
+
 #import "Firestore/Source/Core/FSTSnapshotVersion.h"
 #import "Firestore/Source/Model/FSTDocument.h"
-#import "Firestore/Source/Model/FSTDocumentKey.h"
 #import "Firestore/Source/Remote/FSTWatchChange.h"
 #import "Firestore/Source/Util/FSTAssert.h"
 #import "Firestore/Source/Util/FSTClasses.h"
@@ -26,6 +28,8 @@
 
 #include "Firestore/core/src/firebase/firestore/model/document_key.h"
 
+using firebase::firestore::model::DocumentKey;
+
 NS_ASSUME_NONNULL_BEGIN
 
 #pragma mark - FSTTargetMapping
@@ -33,20 +37,20 @@ NS_ASSUME_NONNULL_BEGIN
 @interface FSTTargetMapping ()
 
 /** Private mutator method to add a document key to the mapping */
-- (void)addDocumentKey:(FSTDocumentKey *)documentKey;
+- (void)addDocumentKey:(const DocumentKey &)documentKey;
 
 /** Private mutator method to remove a document key from the mapping */
-- (void)removeDocumentKey:(FSTDocumentKey *)documentKey;
+- (void)removeDocumentKey:(const DocumentKey &)documentKey;
 
 @end
 
 @implementation FSTTargetMapping
 
-- (void)addDocumentKey:(FSTDocumentKey *)documentKey {
+- (void)addDocumentKey:(const DocumentKey &)documentKey {
   @throw FSTAbstractMethodException();  // NOLINT
 }
 
-- (void)removeDocumentKey:(FSTDocumentKey *)documentKey {
+- (void)removeDocumentKey:(const DocumentKey &)documentKey {
   @throw FSTAbstractMethodException();  // NOLINT
 }
 
@@ -92,11 +96,11 @@ NS_ASSUME_NONNULL_BEGIN
   return self.documents.hash;
 }
 
-- (void)addDocumentKey:(FSTDocumentKey *)documentKey {
+- (void)addDocumentKey:(const DocumentKey &)documentKey {
   self.documents = [self.documents setByAddingObject:documentKey];
 }
 
-- (void)removeDocumentKey:(FSTDocumentKey *)documentKey {
+- (void)removeDocumentKey:(const DocumentKey &)documentKey {
   self.documents = [self.documents setByRemovingObject:documentKey];
 }
 
@@ -160,12 +164,12 @@ NS_ASSUME_NONNULL_BEGIN
   return result;
 }
 
-- (void)addDocumentKey:(FSTDocumentKey *)documentKey {
+- (void)addDocumentKey:(const DocumentKey &)documentKey {
   self.addedDocuments = [self.addedDocuments setByAddingObject:documentKey];
   self.removedDocuments = [self.removedDocuments setByRemovingObject:documentKey];
 }
 
-- (void)removeDocumentKey:(FSTDocumentKey *)documentKey {
+- (void)removeDocumentKey:(const DocumentKey &)documentKey {
   self.addedDocuments = [self.addedDocuments setByRemovingObject:documentKey];
   self.removedDocuments = [self.removedDocuments setByAddingObject:documentKey];
 }
@@ -240,42 +244,39 @@ NS_ASSUME_NONNULL_BEGIN
 #pragma mark - FSTRemoteEvent
 
 @interface FSTRemoteEvent () {
-  NSMutableDictionary<FSTDocumentKey *, FSTMaybeDocument *> *_documentUpdates;
   NSMutableDictionary<FSTBoxedTargetID *, FSTTargetChange *> *_targetChanges;
 }
 
 - (instancetype)
 initWithSnapshotVersion:(FSTSnapshotVersion *)snapshotVersion
           targetChanges:(NSMutableDictionary<FSTBoxedTargetID *, FSTTargetChange *> *)targetChanges
-        documentUpdates:
-            (NSMutableDictionary<FSTDocumentKey *, FSTMaybeDocument *> *)documentUpdates;
+        documentUpdates:(std::map<DocumentKey, FSTMaybeDocument *>)documentUpdates;
 
 @property(nonatomic, strong) FSTSnapshotVersion *snapshotVersion;
 
 @end
 
-@implementation FSTRemoteEvent
-
+@implementation FSTRemoteEvent {
+  std::map<DocumentKey, FSTMaybeDocument *> _documentUpdates;
+}
 + (instancetype)
 eventWithSnapshotVersion:(FSTSnapshotVersion *)snapshotVersion
            targetChanges:(NSMutableDictionary<NSNumber *, FSTTargetChange *> *)targetChanges
-         documentUpdates:
-             (NSMutableDictionary<FSTDocumentKey *, FSTMaybeDocument *> *)documentUpdates {
+         documentUpdates:(std::map<DocumentKey, FSTMaybeDocument *>)documentUpdates {
   return [[FSTRemoteEvent alloc] initWithSnapshotVersion:snapshotVersion
                                            targetChanges:targetChanges
-                                         documentUpdates:documentUpdates];
+                                         documentUpdates:std::move(documentUpdates)];
 }
 
-- (instancetype)
-initWithSnapshotVersion:(FSTSnapshotVersion *)snapshotVersion
-          targetChanges:(NSMutableDictionary<NSNumber *, FSTTargetChange *> *)targetChanges
-        documentUpdates:
-            (NSMutableDictionary<FSTDocumentKey *, FSTMaybeDocument *> *)documentUpdates {
+- (instancetype)initWithSnapshotVersion:(FSTSnapshotVersion *)snapshotVersion
+                          targetChanges:
+                              (NSMutableDictionary<NSNumber *, FSTTargetChange *> *)targetChanges
+                        documentUpdates:(std::map<DocumentKey, FSTMaybeDocument *>)documentUpdates {
   self = [super init];
   if (self) {
     _snapshotVersion = snapshotVersion;
     _targetChanges = targetChanges;
-    _documentUpdates = documentUpdates;
+    _documentUpdates = std::move(documentUpdates);
   }
   return self;
 }
@@ -284,13 +285,13 @@ initWithSnapshotVersion:(FSTSnapshotVersion *)snapshotVersion
   return static_cast<NSDictionary<FSTBoxedTargetID *, FSTTargetChange *> *>(_targetChanges);
 }
 
-- (NSDictionary<FSTDocumentKey *, FSTMaybeDocument *> *)documentUpdates {
-  return static_cast<NSDictionary<FSTDocumentKey *, FSTMaybeDocument *> *>(_documentUpdates);
+- (const std::map<DocumentKey, FSTMaybeDocument *> &)documentUpdates {
+  return _documentUpdates;
 }
 
 /** Adds a document update to this remote event */
 - (void)addDocumentUpdate:(FSTMaybeDocument *)document {
-  _documentUpdates[(FSTDocumentKey *)document.key] = document;
+  _documentUpdates[document.key] = document;
 }
 
 /** Handles an existence filter mismatch */
@@ -326,10 +327,6 @@ initWithSnapshotVersion:(FSTSnapshotVersion *)snapshotVersion
 @property(nonatomic, strong, readonly)
     NSMutableDictionary<FSTBoxedTargetID *, FSTTargetChange *> *targetChanges;
 
-/** Keeps track of document to update */
-@property(nonatomic, strong, readonly)
-    NSMutableDictionary<FSTDocumentKey *, FSTMaybeDocument *> *documentUpdates;
-
 /** The set of open listens on the client */
 @property(nonatomic, strong, readonly)
     NSDictionary<FSTBoxedTargetID *, FSTQueryData *> *listenTargets;
@@ -341,6 +338,8 @@ initWithSnapshotVersion:(FSTSnapshotVersion *)snapshotVersion
 
 @implementation FSTWatchChangeAggregator {
   NSMutableDictionary<FSTBoxedTargetID *, FSTExistenceFilter *> *_existenceFilters;
+  /** Keeps track of document to update */
+  std::map<DocumentKey, FSTMaybeDocument *> _documentUpdates;
 }
 
 - (instancetype)
@@ -357,7 +356,6 @@ initWithSnapshotVersion:(FSTSnapshotVersion *)snapshotVersion
     _pendingTargetResponses = [NSMutableDictionary dictionaryWithDictionary:pendingTargetResponses];
 
     _existenceFilters = [NSMutableDictionary dictionary];
-    _documentUpdates = [NSMutableDictionary dictionary];
   }
   return self;
 }
@@ -418,7 +416,7 @@ initWithSnapshotVersion:(FSTSnapshotVersion *)snapshotVersion
   // Only update the document if there is a new document to replace, this might be just a target
   // update instead.
   if (docChange.document && relevant) {
-    self.documentUpdates[docChange.documentKey] = docChange.document;
+    _documentUpdates[docChange.documentKey] = docChange.document;
   }
 }
 
@@ -522,7 +520,7 @@ initWithSnapshotVersion:(FSTSnapshotVersion *)snapshotVersion
   self.frozen = YES;
   return [FSTRemoteEvent eventWithSnapshotVersion:self.snapshotVersion
                                     targetChanges:targetChanges
-                                  documentUpdates:self.documentUpdates];
+                                  documentUpdates:_documentUpdates];
 }
 
 @end
diff --git a/Firestore/Source/Remote/FSTRemoteStore.mm b/Firestore/Source/Remote/FSTRemoteStore.mm
index f484278..28ee594 100644
--- a/Firestore/Source/Remote/FSTRemoteStore.mm
+++ b/Firestore/Source/Remote/FSTRemoteStore.mm
@@ -24,7 +24,6 @@
 #import "Firestore/Source/Local/FSTLocalStore.h"
 #import "Firestore/Source/Local/FSTQueryData.h"
 #import "Firestore/Source/Model/FSTDocument.h"
-#import "Firestore/Source/Model/FSTDocumentKey.h"
 #import "Firestore/Source/Model/FSTMutation.h"
 #import "Firestore/Source/Model/FSTMutationBatch.h"
 #import "Firestore/Source/Remote/FSTDatastore.h"
@@ -37,10 +36,12 @@
 #import "Firestore/Source/Util/FSTLogger.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;
 
 NS_ASSUME_NONNULL_BEGIN
 
@@ -382,7 +383,7 @@ static const int kMaxPendingWrites = 10;
         // updates. Without applying a deleted document there might be another query that will
         // raise this document as part of a snapshot until it is resolved, essentially exposing
         // inconsistency between queries
-        FSTDocumentKey *key = [FSTDocumentKey keyWithPath:query.path];
+        const DocumentKey key{query.path};
         FSTDeletedDocument *deletedDoc =
             [FSTDeletedDocument documentWithKey:key version:snapshotVersion];
         [remoteEvent addDocumentUpdate:deletedDoc];
diff --git a/Firestore/Source/Remote/FSTSerializerBeta.h b/Firestore/Source/Remote/FSTSerializerBeta.h
index 0f1c3ae..d96dbeb 100644
--- a/Firestore/Source/Remote/FSTSerializerBeta.h
+++ b/Firestore/Source/Remote/FSTSerializerBeta.h
@@ -17,8 +17,8 @@
 #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 FSTDocumentKey;
 @class FSTFieldValue;
 @class FSTMaybeDocument;
 @class FSTMutation;
@@ -70,8 +70,8 @@ NS_ASSUME_NONNULL_BEGIN
 /** Returns the database ID, such as `projects/{project id}/databases/{database_id}`. */
 - (NSString *)encodedDatabaseID;
 
-- (NSString *)encodedDocumentKey:(FSTDocumentKey *)key;
-- (FSTDocumentKey *)decodedDocumentKey:(NSString *)key;
+- (NSString *)encodedDocumentKey:(const firebase::firestore::model::DocumentKey &)key;
+- (firebase::firestore::model::DocumentKey)decodedDocumentKey:(NSString *)key;
 
 - (GCFSValue *)encodedFieldValue:(FSTFieldValue *)fieldValue;
 - (FSTFieldValue *)decodedFieldValue:(GCFSValue *)valueProto;
@@ -95,7 +95,8 @@ NS_ASSUME_NONNULL_BEGIN
 - (FSTWatchChange *)decodedWatchChange:(GCFSListenResponse *)watchChange;
 - (FSTSnapshotVersion *)versionFromListenResponse:(GCFSListenResponse *)watchChange;
 
-- (GCFSDocument *)encodedDocumentWithFields:(FSTObjectValue *)objectValue key:(FSTDocumentKey *)key;
+- (GCFSDocument *)encodedDocumentWithFields:(FSTObjectValue *)objectValue
+                                        key:(const firebase::firestore::model::DocumentKey &)key;
 
 /**
  * Encodes an FSTObjectValue into a dictionary.
diff --git a/Firestore/Source/Remote/FSTSerializerBeta.mm b/Firestore/Source/Remote/FSTSerializerBeta.mm
index 77d2ec7..9531ddc 100644
--- a/Firestore/Source/Remote/FSTSerializerBeta.mm
+++ b/Firestore/Source/Remote/FSTSerializerBeta.mm
@@ -35,7 +35,6 @@
 #import "Firestore/Source/Core/FSTSnapshotVersion.h"
 #import "Firestore/Source/Local/FSTQueryData.h"
 #import "Firestore/Source/Model/FSTDocument.h"
-#import "Firestore/Source/Model/FSTDocumentKey.h"
 #import "Firestore/Source/Model/FSTFieldValue.h"
 #import "Firestore/Source/Model/FSTMutation.h"
 #import "Firestore/Source/Model/FSTMutationBatch.h"
@@ -44,12 +43,14 @@
 #import "Firestore/Source/Util/FSTAssert.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/model/resource_path.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;
 using firebase::firestore::model::FieldPath;
 using firebase::firestore::model::ResourcePath;
 
@@ -104,19 +105,19 @@ NS_ASSUME_NONNULL_BEGIN
   return [[FIRGeoPoint alloc] initWithLatitude:latLng.latitude longitude:latLng.longitude];
 }
 
-#pragma mark - FSTDocumentKey <=> Key proto
+#pragma mark - DocumentKey <=> Key proto
 
-- (NSString *)encodedDocumentKey:(FSTDocumentKey *)key {
-  return [self encodedResourcePathForDatabaseID:self.databaseID path:key.path];
+- (NSString *)encodedDocumentKey:(const DocumentKey &)key {
+  return [self encodedResourcePathForDatabaseID:self.databaseID path:key.path()];
 }
 
-- (FSTDocumentKey *)decodedDocumentKey:(NSString *)name {
+- (DocumentKey)decodedDocumentKey:(NSString *)name {
   const ResourcePath path = [self decodedResourcePathWithDatabaseID:name];
   FSTAssert(path[1] == self.databaseID->project_id(),
             @"Tried to deserialize key from different project.");
   FSTAssert(path[3] == self.databaseID->database_id(),
             @"Tried to deserialize key from different datbase.");
-  return [FSTDocumentKey keyWithPath:[self localResourcePathForQualifiedResourcePath:path]];
+  return DocumentKey{[self localResourcePathForQualifiedResourcePath:path]};
 }
 
 - (NSString *)encodedResourcePathForDatabaseID:(const DatabaseId *)databaseID
@@ -306,12 +307,12 @@ NS_ASSUME_NONNULL_BEGIN
 }
 
 - (GCFSValue *)encodedReferenceValueForDatabaseID:(const DatabaseId *)databaseID
-                                              key:(FSTDocumentKey *)key {
+                                              key:(const DocumentKey &)key {
   FSTAssert(*databaseID == *self.databaseID, @"Database %s:%s cannot encode reference from %s:%s",
             self.databaseID->project_id().c_str(), self.databaseID->database_id().c_str(),
             databaseID->project_id().c_str(), databaseID->database_id().c_str());
   GCFSValue *result = [GCFSValue message];
-  result.referenceValue = [self encodedResourcePathForDatabaseID:databaseID path:key.path];
+  result.referenceValue = [self encodedResourcePathForDatabaseID:databaseID path:key.path()];
   return result;
 }
 
@@ -319,8 +320,7 @@ NS_ASSUME_NONNULL_BEGIN
   const ResourcePath path = [self decodedResourcePathWithDatabaseID:resourceName];
   const std::string &project = path[1];
   const std::string &database = path[3];
-  FSTDocumentKey *key =
-      [FSTDocumentKey keyWithPath:[self localResourcePathForQualifiedResourcePath:path]];
+  const DocumentKey key{[self localResourcePathForQualifiedResourcePath:path]};
 
   const DatabaseId database_id(project, database);
   FSTAssert(database_id == *self.databaseID, @"Database %s:%s cannot encode reference from %s:%s",
@@ -390,7 +390,7 @@ NS_ASSUME_NONNULL_BEGIN
 #pragma mark - FSTObjectValue <=> Document proto
 
 - (GCFSDocument *)encodedDocumentWithFields:(FSTObjectValue *)objectValue
-                                        key:(FSTDocumentKey *)key {
+                                        key:(const DocumentKey &)key {
   GCFSDocument *proto = [GCFSDocument message];
   proto.name = [self encodedDocumentKey:key];
   proto.fields = [self encodedFields:objectValue];
@@ -412,7 +412,7 @@ NS_ASSUME_NONNULL_BEGIN
 
 - (FSTDocument *)decodedFoundDocument:(GCFSBatchGetDocumentsResponse *)response {
   FSTAssert(!!response.found, @"Tried to deserialize a found document from a deleted document.");
-  FSTDocumentKey *key = [self decodedDocumentKey:response.found.name];
+  const DocumentKey key = [self decodedDocumentKey:response.found.name];
   FSTObjectValue *value = [self decodedFields:response.found.fields];
   FSTSnapshotVersion *version = [self decodedVersion:response.found.updateTime];
   FSTAssert(![version isEqual:[FSTSnapshotVersion noVersion]],
@@ -423,7 +423,7 @@ NS_ASSUME_NONNULL_BEGIN
 
 - (FSTDeletedDocument *)decodedDeletedDocument:(GCFSBatchGetDocumentsResponse *)response {
   FSTAssert(!!response.missing, @"Tried to deserialize a deleted document from a found document.");
-  FSTDocumentKey *key = [self decodedDocumentKey:response.missing];
+  const DocumentKey key = [self decodedDocumentKey:response.missing];
   FSTSnapshotVersion *version = [self decodedVersion:response.readTime];
   FSTAssert(![version isEqual:[FSTSnapshotVersion noVersion]],
             @"Got a no document response with no snapshot version");
@@ -1052,7 +1052,7 @@ NS_ASSUME_NONNULL_BEGIN
 
 - (FSTDocumentWatchChange *)decodedDocumentChange:(GCFSDocumentChange *)change {
   FSTObjectValue *value = [self decodedFields:change.document.fields];
-  FSTDocumentKey *key = [self decodedDocumentKey:change.document.name];
+  const DocumentKey key = [self decodedDocumentKey:change.document.name];
   FSTSnapshotVersion *version = [self decodedVersion:change.document.updateTime];
   FSTAssert(![version isEqual:[FSTSnapshotVersion noVersion]],
             @"Got a document change with no snapshot version");
@@ -1069,7 +1069,7 @@ NS_ASSUME_NONNULL_BEGIN
 }
 
 - (FSTDocumentWatchChange *)decodedDocumentDelete:(GCFSDocumentDelete *)change {
-  FSTDocumentKey *key = [self decodedDocumentKey:change.document];
+  const DocumentKey key = [self decodedDocumentKey:change.document];
   // Note that version might be unset in which case we use [FSTSnapshotVersion noVersion]
   FSTSnapshotVersion *version = [self decodedVersion:change.readTime];
   FSTMaybeDocument *document = [FSTDeletedDocument documentWithKey:key version:version];
@@ -1083,7 +1083,7 @@ NS_ASSUME_NONNULL_BEGIN
 }
 
 - (FSTDocumentWatchChange *)decodedDocumentRemove:(GCFSDocumentRemove *)change {
-  FSTDocumentKey *key = [self decodedDocumentKey:change.document];
+  const DocumentKey key = [self decodedDocumentKey:change.document];
   NSArray<NSNumber *> *removedTargetIds = [self decodedIntegerArray:change.removedTargetIdsArray];
 
   return [[FSTDocumentWatchChange alloc] initWithUpdatedTargetIDs:@[]
diff --git a/Firestore/Source/Remote/FSTWatchChange.h b/Firestore/Source/Remote/FSTWatchChange.h
index 8ce24fa..8f730de 100644
--- a/Firestore/Source/Remote/FSTWatchChange.h
+++ b/Firestore/Source/Remote/FSTWatchChange.h
@@ -18,7 +18,8 @@
 
 #import "Firestore/Source/Core/FSTTypes.h"
 
-@class FSTDocumentKey;
+#include "Firestore/core/src/firebase/firestore/model/document_key.h"
+
 @class FSTExistenceFilter;
 @class FSTMaybeDocument;
 @class FSTSnapshotVersion;
@@ -41,21 +42,21 @@ NS_ASSUME_NONNULL_BEGIN
 
 - (instancetype)initWithUpdatedTargetIDs:(NSArray<NSNumber *> *)updatedTargetIDs
                         removedTargetIDs:(NSArray<NSNumber *> *)removedTargetIDs
-                             documentKey:(FSTDocumentKey *)documentKey
+                             documentKey:(firebase::firestore::model::DocumentKey)documentKey
                                 document:(nullable FSTMaybeDocument *)document
     NS_DESIGNATED_INITIALIZER;
 
 - (instancetype)init NS_UNAVAILABLE;
 
+/** The key of the document for this change. */
+- (const firebase::firestore::model::DocumentKey &)documentKey;
+
 /** The new document applies to all of these targets. */
 @property(nonatomic, strong, readonly) NSArray<NSNumber *> *updatedTargetIDs;
 
 /** The new document is removed from all of these targets. */
 @property(nonatomic, strong, readonly) NSArray<NSNumber *> *removedTargetIDs;
 
-/** The key of the document for this change. */
-@property(nonatomic, strong, readonly) FSTDocumentKey *documentKey;
-
 /**
  * The new document or DeletedDocument if it was deleted. Is null if the document went out of
  * view without the server sending a new document.
diff --git a/Firestore/Source/Remote/FSTWatchChange.mm b/Firestore/Source/Remote/FSTWatchChange.mm
index 926d027..2c1a916 100644
--- a/Firestore/Source/Remote/FSTWatchChange.mm
+++ b/Firestore/Source/Remote/FSTWatchChange.mm
@@ -16,31 +16,42 @@
 
 #import "Firestore/Source/Remote/FSTWatchChange.h"
 
+#include <utility>
+
 #import "Firestore/Source/Model/FSTDocument.h"
-#import "Firestore/Source/Model/FSTDocumentKey.h"
 #import "Firestore/Source/Remote/FSTExistenceFilter.h"
 
+#include "Firestore/core/src/firebase/firestore/model/document_key.h"
+
+using firebase::firestore::model::DocumentKey;
+
 NS_ASSUME_NONNULL_BEGIN
 
 @implementation FSTWatchChange
 @end
 
-@implementation FSTDocumentWatchChange
+@implementation FSTDocumentWatchChange {
+  DocumentKey _documentKey;
+}
 
 - (instancetype)initWithUpdatedTargetIDs:(NSArray<NSNumber *> *)updatedTargetIDs
                         removedTargetIDs:(NSArray<NSNumber *> *)removedTargetIDs
-                             documentKey:(FSTDocumentKey *)documentKey
+                             documentKey:(DocumentKey)documentKey
                                 document:(nullable FSTMaybeDocument *)document {
   self = [super init];
   if (self) {
     _updatedTargetIDs = updatedTargetIDs;
     _removedTargetIDs = removedTargetIDs;
-    _documentKey = documentKey;
+    _documentKey = std::move(documentKey);
     _document = document;
   }
   return self;
 }
 
+- (const firebase::firestore::model::DocumentKey &)documentKey {
+  return _documentKey;
+}
+
 - (BOOL)isEqual:(id)other {
   if (other == self) {
     return YES;
@@ -59,7 +70,7 @@ NS_ASSUME_NONNULL_BEGIN
 - (NSUInteger)hash {
   NSUInteger hash = self.updatedTargetIDs.hash;
   hash = hash * 31 + self.removedTargetIDs.hash;
-  hash = hash * 31 + self.documentKey.hash;
+  hash = hash * 31 + std::hash<std::string>{}(self.documentKey.ToString());
   hash = hash * 31 + self.document.hash;
   return hash;
 }