Migrate FSTDocumentVersionDictionary to C++ DocumentVersionMap (#1231)

2 files changed, 7 insertions, 6 deletions

diff --git a/Firestore/Source/Local/FSTLocalStore.h b/Firestore/Source/Local/FSTLocalStore.h
index 1312c28..1f4146a 100644
--- a/Firestore/Source/Local/FSTLocalStore.h
+++ b/Firestore/Source/Local/FSTLocalStore.h
@@ -18,7 +18,6 @@
 
 #import "Firestore/Source/Core/FSTTypes.h"
 #import "Firestore/Source/Model/FSTDocumentDictionary.h"
-#import "Firestore/Source/Model/FSTDocumentVersionDictionary.h"
 
 #include "Firestore/core/src/firebase/firestore/auth/user.h"
 #include "Firestore/core/src/firebase/firestore/model/document_key.h"
diff --git a/Firestore/Source/Local/FSTLocalStore.mm b/Firestore/Source/Local/FSTLocalStore.mm
index e721579..29d0395 100644
--- a/Firestore/Source/Local/FSTLocalStore.mm
+++ b/Firestore/Source/Local/FSTLocalStore.mm
@@ -49,6 +49,7 @@ using firebase::firestore::core::TargetIdGenerator;
 using firebase::firestore::model::DocumentKey;
 using firebase::firestore::model::SnapshotVersion;
 using firebase::firestore::model::DocumentKeySet;
+using firebase::firestore::model::DocumentVersionMap;
 
 NS_ASSUME_NONNULL_BEGIN
 
@@ -513,15 +514,16 @@ NS_ASSUME_NONNULL_BEGIN
 - (void)applyBatchResult:(FSTMutationBatchResult *)batchResult {
   FSTMutationBatch *batch = batchResult.batch;
   DocumentKeySet docKeys = batch.keys;
+  const DocumentVersionMap &versions = batchResult.docVersions;
   for (const DocumentKey &docKey : docKeys) {
     FSTMaybeDocument *_Nullable remoteDoc = [self.remoteDocumentCache entryForKey:docKey];
     FSTMaybeDocument *_Nullable doc = remoteDoc;
-    // TODO(zxu): Once ported to use C++ version of FSTMutationBatchResult, we should be able to
-    // check ackVersion instead, which will be an absl::optional type.
-    FSTAssert(batchResult.docVersions[static_cast<FSTDocumentKey *>(docKey)],
+
+    auto ackVersionIter = versions.find(docKey);
+    FSTAssert(ackVersionIter != versions.end(),
               @"docVersions should contain every doc in the write.");
-    SnapshotVersion ackVersion = batchResult.docVersions[static_cast<FSTDocumentKey *>(docKey)];
-    if (!doc || SnapshotVersion{doc.version} < ackVersion) {
+    const SnapshotVersion &ackVersion = ackVersionIter->second;
+    if (!doc || doc.version < ackVersion) {
       doc = [batch applyTo:doc documentKey:docKey mutationBatchResult:batchResult];
       if (!doc) {
         FSTAssert(!remoteDoc, @"Mutation batch %@ applied to document %@ resulted in nil.", batch,