Migrate FSTDocumentVersionDictionary to C++ DocumentVersionMap (#1231)

7 files changed, 35 insertions, 95 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,
diff --git a/Firestore/Source/Model/FSTDocumentVersionDictionary.h b/Firestore/Source/Model/FSTDocumentVersionDictionary.h
deleted file mode 100644
index 674614e..0000000
--- a/Firestore/Source/Model/FSTDocumentVersionDictionary.h
+++ /dev/null
@@ -1,40 +0,0 @@
-/*
- * Copyright 2017 Google
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-#import <Foundation/Foundation.h>
-
-#import "Firestore/third_party/Immutable/FSTImmutableSortedDictionary.h"
-
-@class FSTDocumentKey;
-@class FSTSnapshotVersion;
-
-NS_ASSUME_NONNULL_BEGIN
-
-/** A map of key to version number. */
-typedef FSTImmutableSortedDictionary<FSTDocumentKey *, FSTSnapshotVersion *>
-    FSTDocumentVersionDictionary;
-
-/**
- * Extension to FSTImmutableSortedDictionary that allows natural construction of
- * FSTDocumentVersionDictionary.
- */
-@interface FSTImmutableSortedDictionary (FSTDocumentVersionDictionary)
-
-+ (instancetype)documentVersionDictionary;
-
-@end
-
-NS_ASSUME_NONNULL_END
diff --git a/Firestore/Source/Model/FSTDocumentVersionDictionary.mm b/Firestore/Source/Model/FSTDocumentVersionDictionary.mm
deleted file mode 100644
index 870e082..0000000
--- a/Firestore/Source/Model/FSTDocumentVersionDictionary.mm
+++ /dev/null
@@ -1,37 +0,0 @@
-/*
- * Copyright 2017 Google
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-#import "Firestore/Source/Model/FSTDocumentVersionDictionary.h"
-
-#import "Firestore/Source/Core/FSTSnapshotVersion.h"
-#import "Firestore/Source/Model/FSTDocumentKey.h"
-
-NS_ASSUME_NONNULL_BEGIN
-
-@implementation FSTImmutableSortedDictionary (FSTDocumentVersionDictionary)
-
-+ (instancetype)documentVersionDictionary {
-  static FSTDocumentVersionDictionary *singleton;
-  static dispatch_once_t onceToken;
-  dispatch_once(&onceToken, ^{
-    singleton = [FSTDocumentVersionDictionary dictionaryWithComparator:FSTDocumentKeyComparator];
-  });
-  return singleton;
-}
-
-@end
-
-NS_ASSUME_NONNULL_END
diff --git a/Firestore/Source/Model/FSTMutationBatch.h b/Firestore/Source/Model/FSTMutationBatch.h
index f95b145..761a885 100644
--- a/Firestore/Source/Model/FSTMutationBatch.h
+++ b/Firestore/Source/Model/FSTMutationBatch.h
@@ -16,8 +16,9 @@
 
 #import <Foundation/Foundation.h>
 
+#include <unordered_map>
+
 #import "Firestore/Source/Core/FSTTypes.h"
-#import "Firestore/Source/Model/FSTDocumentVersionDictionary.h"
 
 #include "Firestore/core/src/firebase/firestore/model/document_key.h"
 #include "Firestore/core/src/firebase/firestore/model/document_key_set.h"
@@ -28,6 +29,17 @@
 @class FSTMutationResult;
 @class FSTMutationBatchResult;
 
+namespace firebase {
+namespace firestore {
+namespace model {
+
+// TODO(wilhuff): make this type a member of MutationBatchResult once that's a C++ class.
+using DocumentVersionMap = std::unordered_map<DocumentKey, SnapshotVersion, DocumentKeyHash>;
+
+}  // namespace model
+}  // namespace firestore
+}  // namespace firebase
+
 NS_ASSUME_NONNULL_BEGIN
 
 /**
@@ -115,7 +127,8 @@ extern const FSTBatchID kFSTBatchIDUnknown;
 @property(nonatomic, strong, readonly) FSTMutationBatch *batch;
 @property(nonatomic, strong, readonly) NSArray<FSTMutationResult *> *mutationResults;
 @property(nonatomic, strong, readonly, nullable) NSData *streamToken;
-@property(nonatomic, strong, readonly) FSTDocumentVersionDictionary *docVersions;
+
+- (const firebase::firestore::model::DocumentVersionMap &)docVersions;
 
 @end
 
diff --git a/Firestore/Source/Model/FSTMutationBatch.mm b/Firestore/Source/Model/FSTMutationBatch.mm
index 9926bd4..1e9189c 100644
--- a/Firestore/Source/Model/FSTMutationBatch.mm
+++ b/Firestore/Source/Model/FSTMutationBatch.mm
@@ -24,11 +24,11 @@
 #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
 
@@ -133,24 +133,25 @@ const FSTBatchID kFSTBatchIDUnknown = -1;
                 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 {
   SnapshotVersion _commitVersion;
+  DocumentVersionMap _docVersions;
 }
 
 - (instancetype)initWithBatch:(FSTMutationBatch *)batch
                 commitVersion:(SnapshotVersion)commitVersion
               mutationResults:(NSArray<FSTMutationResult *> *)mutationResults
                   streamToken:(nullable NSData *)streamToken
-                  docVersions:(FSTDocumentVersionDictionary *)docVersions {
+                  docVersions:(DocumentVersionMap)docVersions {
   if (self = [super init]) {
     _batch = batch;
     _commitVersion = std::move(commitVersion);
     _mutationResults = mutationResults;
     _streamToken = streamToken;
-    _docVersions = docVersions;
+    _docVersions = std::move(docVersions);
   }
   return self;
 }
@@ -159,6 +160,10 @@ const FSTBatchID kFSTBatchIDUnknown = -1;
   return _commitVersion;
 }
 
+- (const DocumentVersionMap &)docVersions {
+  return _docVersions;
+}
+
 + (instancetype)resultWithBatch:(FSTMutationBatch *)batch
                   commitVersion:(SnapshotVersion)commitVersion
                 mutationResults:(NSArray<FSTMutationResult *> *)mutationResults
@@ -167,8 +172,7 @@ const FSTBatchID kFSTBatchIDUnknown = -1;
             @"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++) {
     absl::optional<SnapshotVersion> version = mutationResults[i].version;
@@ -178,14 +182,14 @@ const FSTBatchID kFSTBatchIDUnknown = -1;
       version = commitVersion;
     }
 
-    docVersions = [docVersions dictionaryBySettingObject:version.value() forKey:mutations[i].key];
+    docVersions[mutations[i].key] = version.value();
   }
 
   return [[FSTMutationBatchResult alloc] initWithBatch:batch
                                          commitVersion:std::move(commitVersion)
                                        mutationResults:mutationResults
                                            streamToken:streamToken
-                                           docVersions:docVersions];
+                                           docVersions:std::move(docVersions)];
 }
 
 @end
diff --git a/Firestore/Source/Remote/FSTRemoteStore.h b/Firestore/Source/Remote/FSTRemoteStore.h
index 09e1d32..9b01ce4 100644
--- a/Firestore/Source/Remote/FSTRemoteStore.h
+++ b/Firestore/Source/Remote/FSTRemoteStore.h
@@ -17,7 +17,6 @@
 #import <Foundation/Foundation.h>
 
 #import "Firestore/Source/Core/FSTTypes.h"
-#import "Firestore/Source/Model/FSTDocumentVersionDictionary.h"
 
 #include "Firestore/core/src/firebase/firestore/auth/user.h"