Eliminate unnecessary DocumentKey->FSTDocumentKey conversions (#1507)

10 files changed, 50 insertions, 36 deletions

diff --git a/Firestore/Example/Tests/Integration/API/FIRWriteBatchTests.mm b/Firestore/Example/Tests/Integration/API/FIRWriteBatchTests.mm
index 63616aa..1815444 100644
--- a/Firestore/Example/Tests/Integration/API/FIRWriteBatchTests.mm
+++ b/Firestore/Example/Tests/Integration/API/FIRWriteBatchTests.mm
@@ -336,7 +336,7 @@ int64_t GetCurrentMemoryUsedInMb() {
   const auto errorCode =
       task_info(mach_task_self(), MACH_TASK_BASIC_INFO, (task_info_t)&taskInfo, &taskInfoSize);
   if (errorCode == KERN_SUCCESS) {
-    const int bytesInMegabyte = 1000 * 1000;
+    const int bytesInMegabyte = 1024 * 1024;
     return taskInfo.resident_size / bytesInMegabyte;
   }
   return -1;
@@ -349,8 +349,9 @@ int64_t GetCurrentMemoryUsedInMb() {
   FIRDocumentReference *mainDoc = [self documentRef];
   FIRWriteBatch *batch = [mainDoc.firestore batch];
 
-  // >= 500 mutations will be rejected, so use 500-1 mutations
-  for (int i = 0; i != 500 - 1; ++i) {
+  // > 500 mutations will be rejected.
+  const int maxMutations = 500;
+  for (int i = 0; i != maxMutations; ++i) {
     FIRDocumentReference *nestedDoc = [[mainDoc collectionWithPath:@"nested"] documentWithAutoID];
     // The exact data doesn't matter; what is important is the large number of mutations.
     [batch setData:@{
@@ -367,9 +368,9 @@ int64_t GetCurrentMemoryUsedInMb() {
     const int64_t memoryUsedAfterCommitMb = GetCurrentMemoryUsedInMb();
     XCTAssertNotEqual(memoryUsedAfterCommitMb, -1);
     const int64_t memoryDeltaMb = memoryUsedAfterCommitMb - memoryUsedBeforeCommitMb;
-    // This by its nature cannot be a precise value. In debug mode, the increase on simulator
-    // seems to be around 90 MB. A regression would be on the scale of 500Mb.
-    XCTAssertLessThan(memoryDeltaMb, 150);
+    // This by its nature cannot be a precise value. Runs on simulator seem to give an increase of
+    // 10MB in debug mode pretty consistently. A regression would be on the scale of 500Mb.
+    XCTAssertLessThan(memoryDeltaMb, 20);
 
     [expectation fulfill];
   }];
diff --git a/Firestore/Source/Core/FSTView.mm b/Firestore/Source/Core/FSTView.mm
index f954731..40631ed 100644
--- a/Firestore/Source/Core/FSTView.mm
+++ b/Firestore/Source/Core/FSTView.mm
@@ -109,7 +109,7 @@ NS_ASSUME_NONNULL_BEGIN
     return NO;
   }
   FSTLimboDocumentChange *otherChange = (FSTLimboDocumentChange *)other;
-  return self.type == otherChange.type && [self.key isEqual:otherChange.key];
+  return self.type == otherChange.type && self.key == otherChange.key;
 }
 
 - (NSUInteger)hash {
diff --git a/Firestore/Source/Local/FSTDocumentReference.mm b/Firestore/Source/Local/FSTDocumentReference.mm
index 3e755bc..215801d 100644
--- a/Firestore/Source/Local/FSTDocumentReference.mm
+++ b/Firestore/Source/Local/FSTDocumentReference.mm
@@ -45,7 +45,7 @@ NS_ASSUME_NONNULL_BEGIN
 
   FSTDocumentReference *reference = (FSTDocumentReference *)other;
 
-  return [self.key isEqualToKey:reference.key] && self.ID == reference.ID;
+  return self.key == reference.key && self.ID == reference.ID;
 }
 
 - (NSUInteger)hash {
diff --git a/Firestore/Source/Local/FSTLevelDBRemoteDocumentCache.mm b/Firestore/Source/Local/FSTLevelDBRemoteDocumentCache.mm
index 534d2a4..e3f917e 100644
--- a/Firestore/Source/Local/FSTLevelDBRemoteDocumentCache.mm
+++ b/Firestore/Source/Local/FSTLevelDBRemoteDocumentCache.mm
@@ -122,7 +122,7 @@ using leveldb::Status;
   }
 
   FSTMaybeDocument *maybeDocument = [self.serializer decodedMaybeDocument:proto];
-  HARD_ASSERT([maybeDocument.key isEqualToKey:documentKey],
+  HARD_ASSERT(maybeDocument.key == documentKey,
               "Read document has key (%s) instead of expected key (%s).",
               maybeDocument.key.ToString(), documentKey.ToString());
   return maybeDocument;
diff --git a/Firestore/Source/Model/FSTDocument.mm b/Firestore/Source/Model/FSTDocument.mm
index 376075e..f96552c 100644
--- a/Firestore/Source/Model/FSTDocument.mm
+++ b/Firestore/Source/Model/FSTDocument.mm
@@ -102,7 +102,7 @@ NS_ASSUME_NONNULL_BEGIN
   }
 
   FSTDocument *otherDoc = other;
-  return [self.key isEqual:otherDoc.key] && self.version == otherDoc.version &&
+  return self.key == otherDoc.key && self.version == otherDoc.version &&
          [self.data isEqual:otherDoc.data] && self.hasLocalMutations == otherDoc.hasLocalMutations;
 }
 
@@ -142,7 +142,7 @@ NS_ASSUME_NONNULL_BEGIN
   }
 
   FSTDocument *otherDoc = other;
-  return [self.key isEqual:otherDoc.key] && self.version == otherDoc.version;
+  return self.key == otherDoc.key && self.version == otherDoc.version;
 }
 
 - (NSUInteger)hash {
diff --git a/Firestore/Source/Model/FSTDocumentKey.h b/Firestore/Source/Model/FSTDocumentKey.h
index a403117..ebc88c9 100644
--- a/Firestore/Source/Model/FSTDocumentKey.h
+++ b/Firestore/Source/Model/FSTDocumentKey.h
@@ -21,6 +21,15 @@
 
 #include "Firestore/core/src/firebase/firestore/model/resource_path.h"
 
+// Using forward declaration to avoid circular dependency (`document_key.h` includes this header).`
+namespace firebase {
+namespace firestore {
+namespace model {
+class DocumentKey;
+}
+}
+}
+
 NS_ASSUME_NONNULL_BEGIN
 
 /** FSTDocumentKey represents the location of a document in the Firestore database. */
@@ -33,6 +42,8 @@ NS_ASSUME_NONNULL_BEGIN
  * @return A new instance of FSTDocumentKey.
  */
 + (instancetype)keyWithPath:(firebase::firestore::model::ResourcePath)path;
+
++ (instancetype)keyWithDocumentKey:(const firebase::firestore::model::DocumentKey &)documentKey;
 /**
  * Creates and returns a new document key with a path with the given segments.
  *
diff --git a/Firestore/Source/Model/FSTDocumentKey.mm b/Firestore/Source/Model/FSTDocumentKey.mm
index ad3968e..0d77136 100644
--- a/Firestore/Source/Model/FSTDocumentKey.mm
+++ b/Firestore/Source/Model/FSTDocumentKey.mm
@@ -21,26 +21,32 @@
 
 #import "Firestore/Source/Core/FSTFirestoreClient.h"
 
-#include "Firestore/core/src/firebase/firestore/model/resource_path.h"
+#include "Firestore/core/src/firebase/firestore/model/document_key.h"
 #include "Firestore/core/src/firebase/firestore/util/hard_assert.h"
 #include "Firestore/core/src/firebase/firestore/util/hashing.h"
 #include "Firestore/core/src/firebase/firestore/util/string_apple.h"
 
 namespace util = firebase::firestore::util;
 using firebase::firestore::model::ResourcePath;
+using firebase::firestore::model::DocumentKey;
 
 NS_ASSUME_NONNULL_BEGIN
 
 @interface FSTDocumentKey () {
-  /** The path to the document. */
-  ResourcePath _path;
+  // Forward most of the logic to the C++ implementation until FSTDocumentKey usages are completely
+  // migrated.
+  DocumentKey _delegate;
 }
 @end
 
 @implementation FSTDocumentKey
 
 + (instancetype)keyWithPath:(ResourcePath)path {
-  return [[FSTDocumentKey alloc] initWithPath:std::move(path)];
+  return [[FSTDocumentKey alloc] initWithDocumentKey:DocumentKey{path}];
+}
+
++ (instancetype)keyWithDocumentKey:(const firebase::firestore::model::DocumentKey &)documentKey {
+  return [[FSTDocumentKey alloc] initWithDocumentKey:documentKey];
 }
 
 + (instancetype)keyWithSegments:(std::initializer_list<std::string>)segments {
@@ -52,12 +58,9 @@ NS_ASSUME_NONNULL_BEGIN
 }
 
 /** Designated initializer. */
-- (instancetype)initWithPath:(ResourcePath)path {
-  HARD_ASSERT([FSTDocumentKey isDocumentKey:path], "invalid document key path: %s",
-              path.CanonicalString());
-
+- (instancetype)initWithDocumentKey:(const DocumentKey &)key {
   if (self = [super init]) {
-    _path = path;
+    _delegate = key;
   }
   return self;
 }
@@ -73,11 +76,11 @@ NS_ASSUME_NONNULL_BEGIN
 }
 
 - (NSUInteger)hash {
-  return util::Hash(_path);
+  return _delegate.Hash();
 }
 
 - (NSString *)description {
-  return [NSString stringWithFormat:@"<FSTDocumentKey: %s>", _path.CanonicalString().c_str()];
+  return [NSString stringWithFormat:@"<FSTDocumentKey: %s>", _delegate.ToString().c_str()];
 }
 
 /** Implements NSCopying without actually copying because FSTDocumentKeys are immutable. */
@@ -100,11 +103,11 @@ NS_ASSUME_NONNULL_BEGIN
 }
 
 + (BOOL)isDocumentKey:(const ResourcePath &)path {
-  return path.size() % 2 == 0;
+  return DocumentKey::IsDocumentKey(path);
 }
 
 - (const ResourcePath &)path {
-  return _path;
+  return _delegate.path();
 }
 
 @end
diff --git a/Firestore/Source/Model/FSTMutation.mm b/Firestore/Source/Model/FSTMutation.mm
index d124c78..0337fbb 100644
--- a/Firestore/Source/Model/FSTMutation.mm
+++ b/Firestore/Source/Model/FSTMutation.mm
@@ -137,7 +137,7 @@ NS_ASSUME_NONNULL_BEGIN
   }
 
   FSTSetMutation *otherMutation = (FSTSetMutation *)other;
-  return [self.key isEqual:otherMutation.key] && [self.value isEqual:otherMutation.value] &&
+  return self.key == otherMutation.key && [self.value isEqual:otherMutation.value] &&
          self.precondition == otherMutation.precondition;
 }
 
@@ -173,7 +173,7 @@ NS_ASSUME_NONNULL_BEGIN
               [maybeDoc class]);
   FSTDocument *doc = (FSTDocument *)maybeDoc;
 
-  HARD_ASSERT([doc.key isEqual:self.key], "Can only set a document with the same key");
+  HARD_ASSERT(doc.key == self.key, "Can only set a document with the same key");
   return [FSTDocument documentWithData:self.value
                                    key:doc.key
                                version:doc.version
@@ -212,7 +212,7 @@ NS_ASSUME_NONNULL_BEGIN
   }
 
   FSTPatchMutation *otherMutation = (FSTPatchMutation *)other;
-  return [self.key isEqual:otherMutation.key] && self.fieldMask == otherMutation.fieldMask &&
+  return self.key == otherMutation.key && self.fieldMask == otherMutation.fieldMask &&
          [self.value isEqual:otherMutation.value] &&
          self.precondition == otherMutation.precondition;
 }
@@ -259,7 +259,7 @@ NS_ASSUME_NONNULL_BEGIN
               [maybeDoc class]);
   FSTDocument *doc = (FSTDocument *)maybeDoc;
 
-  HARD_ASSERT([doc.key isEqual:self.key], "Can only patch a document with the same key");
+  HARD_ASSERT(doc.key == self.key, "Can only patch a document with the same key");
 
   FSTObjectValue *newData = [self patchObjectValue:doc.data];
   return [FSTDocument documentWithData:newData
@@ -312,8 +312,7 @@ NS_ASSUME_NONNULL_BEGIN
   }
 
   FSTTransformMutation *otherMutation = (FSTTransformMutation *)other;
-  return [self.key isEqual:otherMutation.key] &&
-         self.fieldTransforms == otherMutation.fieldTransforms &&
+  return self.key == otherMutation.key && self.fieldTransforms == otherMutation.fieldTransforms &&
          self.precondition == otherMutation.precondition;
 }
 
@@ -355,7 +354,7 @@ NS_ASSUME_NONNULL_BEGIN
               [maybeDoc class]);
   FSTDocument *doc = (FSTDocument *)maybeDoc;
 
-  HARD_ASSERT([doc.key isEqual:self.key], "Can only transform a document with the same key");
+  HARD_ASSERT(doc.key == self.key, "Can only transform a document with the same key");
 
   BOOL hasLocalMutations = (mutationResult == nil);
   NSArray<FSTFieldValue *> *transformResults;
@@ -460,7 +459,7 @@ serverTransformResultsWithBaseDocument:(nullable FSTMaybeDocument *)baseDocument
   }
 
   FSTDeleteMutation *otherMutation = (FSTDeleteMutation *)other;
-  return [self.key isEqual:otherMutation.key] && self.precondition == otherMutation.precondition;
+  return self.key == otherMutation.key && self.precondition == otherMutation.precondition;
 }
 
 - (NSUInteger)hash {
@@ -488,7 +487,7 @@ serverTransformResultsWithBaseDocument:(nullable FSTMaybeDocument *)baseDocument
   }
 
   if (maybeDoc) {
-    HARD_ASSERT([maybeDoc.key isEqual:self.key], "Can only delete a document with the same key");
+    HARD_ASSERT(maybeDoc.key == self.key, "Can only delete a document with the same key");
   }
 
   return [FSTDeletedDocument documentWithKey:self.key version:SnapshotVersion::None()];
diff --git a/Firestore/Source/Model/FSTMutationBatch.mm b/Firestore/Source/Model/FSTMutationBatch.mm
index b3f4dde..11d83e9 100644
--- a/Firestore/Source/Model/FSTMutationBatch.mm
+++ b/Firestore/Source/Model/FSTMutationBatch.mm
@@ -77,7 +77,7 @@ const FSTBatchID kFSTBatchIDUnknown = -1;
 - (FSTMaybeDocument *_Nullable)applyTo:(FSTMaybeDocument *_Nullable)maybeDoc
                            documentKey:(const DocumentKey &)documentKey
                    mutationBatchResult:(FSTMutationBatchResult *_Nullable)mutationBatchResult {
-  HARD_ASSERT(!maybeDoc || [maybeDoc.key isEqualToKey:documentKey],
+  HARD_ASSERT(!maybeDoc || maybeDoc.key == documentKey,
               "applyTo: key %s doesn't match maybeDoc key %s", documentKey.ToString(),
               maybeDoc.key.ToString());
   FSTMaybeDocument *baseDoc = maybeDoc;
@@ -90,7 +90,7 @@ const FSTBatchID kFSTBatchIDUnknown = -1;
   for (NSUInteger i = 0; i < self.mutations.count; i++) {
     FSTMutation *mutation = self.mutations[i];
     FSTMutationResult *_Nullable mutationResult = mutationBatchResult.mutationResults[i];
-    if ([mutation.key isEqualToKey:documentKey]) {
+    if (mutation.key == documentKey) {
       maybeDoc = [mutation applyTo:maybeDoc
                       baseDocument:baseDoc
                     localWriteTime:self.localWriteTime
diff --git a/Firestore/core/src/firebase/firestore/model/document_key.h b/Firestore/core/src/firebase/firestore/model/document_key.h
index 3f5f342..19671f9 100644
--- a/Firestore/core/src/firebase/firestore/model/document_key.h
+++ b/Firestore/core/src/firebase/firestore/model/document_key.h
@@ -56,7 +56,7 @@ class DocumentKey {
   }
 
   operator FSTDocumentKey*() const {
-    return [FSTDocumentKey keyWithPath:path()];
+    return [FSTDocumentKey keyWithDocumentKey:*this];
   }
 
   NSUInteger Hash() const {