1 files changed, 98 insertions, 85 deletions

diff --git a/Firestore/Source/Remote/FSTRemoteEvent.mm b/Firestore/Source/Remote/FSTRemoteEvent.mm
index 718fd11..438072e 100644
--- a/Firestore/Source/Remote/FSTRemoteEvent.mm
+++ b/Firestore/Source/Remote/FSTRemoteEvent.mm
@@ -26,9 +26,12 @@
 #import "Firestore/Source/Util/FSTClasses.h"
 #import "Firestore/Source/Util/FSTLogger.h"
 #include "Firestore/core/src/firebase/firestore/model/document_key.h"
+#include "Firestore/core/src/firebase/firestore/util/hashing.h"
 
 using firebase::firestore::model::DocumentKey;
 using firebase::firestore::model::SnapshotVersion;
+using firebase::firestore::util::Hash;
+using firebase::firestore::model::DocumentKeySet;
 
 NS_ASSUME_NONNULL_BEGIN
 
@@ -54,7 +57,7 @@ NS_ASSUME_NONNULL_BEGIN
   @throw FSTAbstractMethodException();  // NOLINT
 }
 
-- (void)filterUpdatesUsingExistingKeys:(FSTDocumentKeySet *)existingKeys {
+- (void)filterUpdatesUsingExistingKeys:(const DocumentKeySet &)existingKeys {
   @throw FSTAbstractMethodException();  // NOLINT
 }
 
@@ -62,28 +65,30 @@ NS_ASSUME_NONNULL_BEGIN
 
 #pragma mark - FSTResetMapping
 
-@interface FSTResetMapping ()
-@property(nonatomic, strong) FSTDocumentKeySet *documents;
-@end
-
-@implementation FSTResetMapping
+@implementation FSTResetMapping {
+  DocumentKeySet _documents;
+}
 
 + (instancetype)mappingWithDocuments:(NSArray<FSTDocument *> *)documents {
-  FSTResetMapping *mapping = [[FSTResetMapping alloc] init];
+  DocumentKeySet keys;
   for (FSTDocument *doc in documents) {
-    mapping.documents = [mapping.documents setByAddingObject:doc.key];
+    keys = keys.insert(doc.key);
   }
-  return mapping;
+  return [[FSTResetMapping alloc] initWithDocuments:std::move(keys)];
 }
 
-- (instancetype)init {
+- (instancetype)initWithDocuments:(DocumentKeySet)documents {
   self = [super init];
   if (self) {
-    _documents = [FSTDocumentKeySet keySet];
+    _documents = std::move(documents);
   }
   return self;
 }
 
+- (const DocumentKeySet &)documents {
+  return _documents;
+}
+
 - (BOOL)isEqual:(id)other {
   if (other == self) {
     return YES;
@@ -93,22 +98,22 @@ NS_ASSUME_NONNULL_BEGIN
   }
 
   FSTResetMapping *otherMapping = (FSTResetMapping *)other;
-  return [self.documents isEqual:otherMapping.documents];
+  return _documents == otherMapping.documents;
 }
 
 - (NSUInteger)hash {
-  return self.documents.hash;
+  return Hash(_documents);
 }
 
 - (void)addDocumentKey:(const DocumentKey &)documentKey {
-  self.documents = [self.documents setByAddingObject:documentKey];
+  _documents = _documents.insert(documentKey);
 }
 
 - (void)removeDocumentKey:(const DocumentKey &)documentKey {
-  self.documents = [self.documents setByRemovingObject:documentKey];
+  _documents = _documents.erase(documentKey);
 }
 
-- (void)filterUpdatesUsingExistingKeys:(FSTDocumentKeySet *)existingKeys {
+- (void)filterUpdatesUsingExistingKeys:(const DocumentKeySet &)existingKeys {
   // No-op. Resets are not filtered.
 }
 
@@ -116,34 +121,43 @@ NS_ASSUME_NONNULL_BEGIN
 
 #pragma mark - FSTUpdateMapping
 
-@interface FSTUpdateMapping ()
-@property(nonatomic, strong) FSTDocumentKeySet *addedDocuments;
-@property(nonatomic, strong) FSTDocumentKeySet *removedDocuments;
-@end
-
-@implementation FSTUpdateMapping
+@implementation FSTUpdateMapping {
+  DocumentKeySet _addedDocuments;
+  DocumentKeySet _removedDocuments;
+}
 
 + (FSTUpdateMapping *)mappingWithAddedDocuments:(NSArray<FSTDocument *> *)added
                                removedDocuments:(NSArray<FSTDocument *> *)removed {
-  FSTUpdateMapping *mapping = [[FSTUpdateMapping alloc] init];
+  DocumentKeySet addedDocuments;
+  DocumentKeySet removedDocuments;
   for (FSTDocument *doc in added) {
-    mapping.addedDocuments = [mapping.addedDocuments setByAddingObject:doc.key];
+    addedDocuments = addedDocuments.insert(doc.key);
   }
   for (FSTDocument *doc in removed) {
-    mapping.removedDocuments = [mapping.removedDocuments setByAddingObject:doc.key];
+    removedDocuments = removedDocuments.insert(doc.key);
   }
-  return mapping;
+  return [[FSTUpdateMapping alloc] initWithAddedDocuments:std::move(addedDocuments)
+                                         removedDocuments:std::move(removedDocuments)];
 }
 
-- (instancetype)init {
+- (instancetype)initWithAddedDocuments:(DocumentKeySet)addedDocuments
+                      removedDocuments:(DocumentKeySet)removedDocuments {
   self = [super init];
   if (self) {
-    _addedDocuments = [FSTDocumentKeySet keySet];
-    _removedDocuments = [FSTDocumentKeySet keySet];
+    _addedDocuments = std::move(addedDocuments);
+    _removedDocuments = std::move(removedDocuments);
   }
   return self;
 }
 
+- (const DocumentKeySet &)addedDocuments {
+  return _addedDocuments;
+}
+
+- (const DocumentKeySet &)removedDocuments {
+  return _removedDocuments;
+}
+
 - (BOOL)isEqual:(id)other {
   if (other == self) {
     return YES;
@@ -153,42 +167,42 @@ NS_ASSUME_NONNULL_BEGIN
   }
 
   FSTUpdateMapping *otherMapping = (FSTUpdateMapping *)other;
-  return [self.addedDocuments isEqual:otherMapping.addedDocuments] &&
-         [self.removedDocuments isEqual:otherMapping.removedDocuments];
+  return _addedDocuments == otherMapping.addedDocuments &&
+         _removedDocuments == otherMapping.removedDocuments;
 }
 
 - (NSUInteger)hash {
-  return self.addedDocuments.hash * 31 + self.removedDocuments.hash;
+  return Hash(_addedDocuments, _removedDocuments);
 }
 
-- (FSTDocumentKeySet *)applyTo:(FSTDocumentKeySet *)keys {
-  __block FSTDocumentKeySet *result = keys;
-  [self.addedDocuments enumerateObjectsUsingBlock:^(FSTDocumentKey *key, BOOL *stop) {
-    result = [result setByAddingObject:key];
-  }];
-  [self.removedDocuments enumerateObjectsUsingBlock:^(FSTDocumentKey *key, BOOL *stop) {
-    result = [result setByRemovingObject:key];
-  }];
+- (DocumentKeySet)applyTo:(const DocumentKeySet &)keys {
+  DocumentKeySet result = keys;
+  for (const DocumentKey &key : _addedDocuments) {
+    result = result.insert(key);
+  }
+  for (const DocumentKey &key : _removedDocuments) {
+    result = result.erase(key);
+  }
   return result;
 }
 
 - (void)addDocumentKey:(const DocumentKey &)documentKey {
-  self.addedDocuments = [self.addedDocuments setByAddingObject:documentKey];
-  self.removedDocuments = [self.removedDocuments setByRemovingObject:documentKey];
+  _addedDocuments = _addedDocuments.insert(documentKey);
+  _removedDocuments = _removedDocuments.erase(documentKey);
 }
 
 - (void)removeDocumentKey:(const DocumentKey &)documentKey {
-  self.addedDocuments = [self.addedDocuments setByRemovingObject:documentKey];
-  self.removedDocuments = [self.removedDocuments setByAddingObject:documentKey];
+  _addedDocuments = _addedDocuments.erase(documentKey);
+  _removedDocuments = _removedDocuments.insert(documentKey);
 }
 
-- (void)filterUpdatesUsingExistingKeys:(FSTDocumentKeySet *)existingKeys {
-  __block FSTDocumentKeySet *result = _addedDocuments;
-  [_addedDocuments enumerateObjectsUsingBlock:^(FSTDocumentKey *docKey, BOOL *stop) {
-    if ([existingKeys containsObject:docKey]) {
-      result = [result setByRemovingObject:docKey];
+- (void)filterUpdatesUsingExistingKeys:(const DocumentKeySet &)existingKeys {
+  DocumentKeySet result = _addedDocuments;
+  for (const DocumentKey &key : _addedDocuments) {
+    if (existingKeys.contains(key)) {
+      result = result.erase(key);
     }
-  }];
+  }
   _addedDocuments = result;
 }
 
@@ -227,14 +241,19 @@ NS_ASSUME_NONNULL_BEGIN
 
 + (instancetype)changeWithDocuments:(NSArray<FSTMaybeDocument *> *)docs
                 currentStatusUpdate:(FSTCurrentStatusUpdate)currentStatusUpdate {
-  FSTUpdateMapping *mapping = [[FSTUpdateMapping alloc] init];
+  DocumentKeySet addedDocuments;
+  DocumentKeySet removedDocuments;
   for (FSTMaybeDocument *doc in docs) {
     if ([doc isKindOfClass:[FSTDeletedDocument class]]) {
-      mapping.removedDocuments = [mapping.removedDocuments setByAddingObject:doc.key];
+      removedDocuments = removedDocuments.insert(doc.key);
     } else {
-      mapping.addedDocuments = [mapping.addedDocuments setByAddingObject:doc.key];
+      addedDocuments = addedDocuments.insert(doc.key);
     }
   }
+  FSTUpdateMapping *mapping =
+      [[FSTUpdateMapping alloc] initWithAddedDocuments:std::move(addedDocuments)
+                                      removedDocuments:std::move(removedDocuments)];
+
   FSTTargetChange *change = [[FSTTargetChange alloc] init];
   change.mapping = mapping;
   change.currentStatusUpdate = currentStatusUpdate;
@@ -273,38 +292,44 @@ NS_ASSUME_NONNULL_BEGIN
 
 #pragma mark - FSTRemoteEvent
 
-@interface FSTRemoteEvent () {
-  NSMutableDictionary<FSTBoxedTargetID *, FSTTargetChange *> *_targetChanges;
-}
-
-- (instancetype)
-initWithSnapshotVersion:(SnapshotVersion)snapshotVersion
-          targetChanges:(NSMutableDictionary<FSTBoxedTargetID *, FSTTargetChange *> *)targetChanges
-        documentUpdates:(std::map<DocumentKey, FSTMaybeDocument *>)documentUpdates
-         limboDocuments:(FSTDocumentKeySet *)limboDocuments;
-
-@end
-
 @implementation FSTRemoteEvent {
-  std::map<DocumentKey, FSTMaybeDocument *> _documentUpdates;
   SnapshotVersion _snapshotVersion;
+  NSMutableDictionary<FSTBoxedTargetID *, FSTTargetChange *> *_targetChanges;
+  std::map<DocumentKey, FSTMaybeDocument *> _documentUpdates;
+  DocumentKeySet _limboDocumentChanges;
 }
 
 - (instancetype)initWithSnapshotVersion:(SnapshotVersion)snapshotVersion
                           targetChanges:
                               (NSMutableDictionary<NSNumber *, FSTTargetChange *> *)targetChanges
                         documentUpdates:(std::map<DocumentKey, FSTMaybeDocument *>)documentUpdates
-                         limboDocuments:(FSTDocumentKeySet *)limboDocuments {
+                         limboDocuments:(DocumentKeySet)limboDocuments {
   self = [super init];
   if (self) {
     _snapshotVersion = std::move(snapshotVersion);
     _targetChanges = targetChanges;
     _documentUpdates = std::move(documentUpdates);
-    _limboDocumentChanges = limboDocuments;
+    _limboDocumentChanges = std::move(limboDocuments);
   }
   return self;
 }
 
+- (NSDictionary<FSTBoxedTargetID *, FSTTargetChange *> *)targetChanges {
+  return _targetChanges;
+}
+
+- (const DocumentKeySet &)limboDocumentChanges {
+  return _limboDocumentChanges;
+}
+
+- (const std::map<DocumentKey, FSTMaybeDocument *> &)documentUpdates {
+  return _documentUpdates;
+}
+
+- (const SnapshotVersion &)snapshotVersion {
+  return _snapshotVersion;
+}
+
 - (void)synthesizeDeleteForLimboTargetChange:(FSTTargetChange *)targetChange
                                          key:(const DocumentKey &)key {
   if (targetChange.currentStatusUpdate == FSTCurrentStatusUpdateMarkCurrent &&
@@ -328,22 +353,10 @@ initWithSnapshotVersion:(SnapshotVersion)snapshotVersion
     // TODO(dimond): Ideally we would have an explicit lookup query instead resulting in an
     // explicit delete message and we could remove this special logic.
     _documentUpdates[key] = [FSTDeletedDocument documentWithKey:key version:_snapshotVersion];
-    _limboDocumentChanges = [_limboDocumentChanges setByAddingObject:key];
+    _limboDocumentChanges = _limboDocumentChanges.insert(key);
   }
 }
 
-- (NSDictionary<FSTBoxedTargetID *, FSTTargetChange *> *)targetChanges {
-  return _targetChanges;
-}
-
-- (const std::map<DocumentKey, FSTMaybeDocument *> &)documentUpdates {
-  return _documentUpdates;
-}
-
-- (const SnapshotVersion &)snapshotVersion {
-  return _snapshotVersion;
-}
-
 /** Adds a document update to this remote event */
 - (void)addDocumentUpdate:(FSTMaybeDocument *)document {
   _documentUpdates[document.key] = document;
@@ -393,7 +406,7 @@ initWithSnapshotVersion:(SnapshotVersion)snapshotVersion
   /** Keeps track of document to update */
   std::map<DocumentKey, FSTMaybeDocument *> _documentUpdates;
 
-  FSTDocumentKeySet *_limboDocuments;
+  DocumentKeySet _limboDocuments;
   /** The snapshot version for every target change this creates. */
   SnapshotVersion _snapshotVersion;
 }
@@ -410,7 +423,7 @@ initWithSnapshotVersion:(SnapshotVersion)snapshotVersion
     _targetChanges = [NSMutableDictionary dictionary];
     _listenTargets = listenTargets;
     _pendingTargetResponses = [NSMutableDictionary dictionaryWithDictionary:pendingTargetResponses];
-    _limboDocuments = [FSTDocumentKeySet keySet];
+    _limboDocuments = DocumentKeySet{};
     _existenceFilters = [NSMutableDictionary dictionary];
   }
   return self;
@@ -467,7 +480,7 @@ initWithSnapshotVersion:(SnapshotVersion)snapshotVersion
     // We haven't seen a document update for this key yet.
     if (queryData.purpose == FSTQueryPurposeLimboResolution) {
       // We haven't seen this document before, and this target is a limbo target.
-      _limboDocuments = [_limboDocuments setByAddingObject:documentKey];
+      _limboDocuments = _limboDocuments.insert(documentKey);
       return YES;
     } else {
       // We haven't seen the document before, but this is a non-limbo target.
@@ -481,7 +494,7 @@ initWithSnapshotVersion:(SnapshotVersion)snapshotVersion
   } else {
     // We haven't marked this as non-limbo yet, but this target is not a limbo target.
     // Mark the key as non-limbo and make sure it isn't in our set.
-    _limboDocuments = [_limboDocuments setByRemovingObject:documentKey];
+    _limboDocuments = _limboDocuments.erase(documentKey);
     return NO;
   }
 }