1 files changed, 33 insertions, 20 deletions
diff --git a/Firestore/Source/Model/FSTMutationBatch.mm b/Firestore/Source/Model/FSTMutationBatch.mm
index e62a72c..1e9189c 100644
--- a/Firestore/Source/Model/FSTMutationBatch.mm
+++ b/Firestore/Source/Model/FSTMutationBatch.mm
@@ -16,16 +16,19 @@
 
 #import "Firestore/Source/Model/FSTMutationBatch.h"
 
+#include <utility>
+
 #import "FIRTimestamp.h"
 
-#import "Firestore/Source/Core/FSTSnapshotVersion.h"
 #import "Firestore/Source/Model/FSTDocument.h"
 #import "Firestore/Source/Model/FSTMutation.h"
 #import "Firestore/Source/Util/FSTAssert.h"
 
-#include "Firestore/core/src/firebase/firestore/model/document_key.h"
-
 using firebase::firestore::model::DocumentKey;
+using firebase::firestore::model::DocumentKeyHash;
+using firebase::firestore::model::SnapshotVersion;
+using firebase::firestore::model::DocumentKeySet;
+using firebase::firestore::model::DocumentVersionMap;
 
 NS_ASSUME_NONNULL_BEGIN
 
@@ -113,10 +116,10 @@ const FSTBatchID kFSTBatchIDUnknown = -1;
 }
 
 // TODO(klimt): This could use NSMutableDictionary instead.
-- (FSTDocumentKeySet *)keys {
-  FSTDocumentKeySet *set = [FSTDocumentKeySet keySet];
+- (DocumentKeySet)keys {
+  DocumentKeySet set;
   for (FSTMutation *mutation in self.mutations) {
-    set = [set setByAddingObject:mutation.key];
+    set = set.insert(mutation.key);
   }
   return set;
 }
@@ -127,56 +130,66 @@ const FSTBatchID kFSTBatchIDUnknown = -1;
 
 @interface FSTMutationBatchResult ()
 - (instancetype)initWithBatch:(FSTMutationBatch *)batch
-                commitVersion:(FSTSnapshotVersion *)commitVersion
+                commitVersion:(SnapshotVersion)commitVersion
               mutationResults:(NSArray<FSTMutationResult *> *)mutationResults
                   streamToken:(nullable NSData *)streamToken
-                  docVersions:(FSTDocumentVersionDictionary *)docVersions NS_DESIGNATED_INITIALIZER;
+                  docVersions:(DocumentVersionMap)docVersions NS_DESIGNATED_INITIALIZER;
 @end
 
-@implementation FSTMutationBatchResult
+@implementation FSTMutationBatchResult {
+  SnapshotVersion _commitVersion;
+  DocumentVersionMap _docVersions;
+}
 
 - (instancetype)initWithBatch:(FSTMutationBatch *)batch
-                commitVersion:(FSTSnapshotVersion *)commitVersion
+                commitVersion:(SnapshotVersion)commitVersion
               mutationResults:(NSArray<FSTMutationResult *> *)mutationResults
                   streamToken:(nullable NSData *)streamToken
-                  docVersions:(FSTDocumentVersionDictionary *)docVersions {
+                  docVersions:(DocumentVersionMap)docVersions {
   if (self = [super init]) {
     _batch = batch;
-    _commitVersion = commitVersion;
+    _commitVersion = std::move(commitVersion);
     _mutationResults = mutationResults;
     _streamToken = streamToken;
-    _docVersions = docVersions;
+    _docVersions = std::move(docVersions);
   }
   return self;
 }
 
+- (const SnapshotVersion &)commitVersion {
+  return _commitVersion;
+}
+
+- (const DocumentVersionMap &)docVersions {
+  return _docVersions;
+}
+
 + (instancetype)resultWithBatch:(FSTMutationBatch *)batch
-                  commitVersion:(FSTSnapshotVersion *)commitVersion
+                  commitVersion:(SnapshotVersion)commitVersion
                 mutationResults:(NSArray<FSTMutationResult *> *)mutationResults
                     streamToken:(nullable NSData *)streamToken {
   FSTAssert(batch.mutations.count == mutationResults.count,
             @"Mutations sent %lu must equal results received %lu",
             (unsigned long)batch.mutations.count, (unsigned long)mutationResults.count);
 
-  FSTDocumentVersionDictionary *docVersions =
-      [FSTDocumentVersionDictionary documentVersionDictionary];
+  DocumentVersionMap docVersions;
   NSArray<FSTMutation *> *mutations = batch.mutations;
   for (NSUInteger i = 0; i < mutations.count; i++) {
-    FSTSnapshotVersion *_Nullable version = mutationResults[i].version;
+    absl::optional<SnapshotVersion> version = mutationResults[i].version;
     if (!version) {
       // deletes don't have a version, so we substitute the commitVersion
       // of the entire batch.
       version = commitVersion;
     }
 
-    docVersions = [docVersions dictionaryBySettingObject:version forKey:mutations[i].key];
+    docVersions[mutations[i].key] = version.value();
   }
 
   return [[FSTMutationBatchResult alloc] initWithBatch:batch
-                                         commitVersion:commitVersion
+                                         commitVersion:std::move(commitVersion)
                                        mutationResults:mutationResults
                                            streamToken:streamToken
-                                           docVersions:docVersions];
+                                           docVersions:std::move(docVersions)];
 }
 
 @end