Add allMutationsAffectingDocumentKeys to FSTMutationQueue (#1479)

* Pod updates for Cocapods 1.5.3

* Add allMutationsAffectingDocumentKeys

4 files changed, 199 insertions, 22 deletions

diff --git a/Firestore/Example/Tests/Local/FSTMutationQueueTests.mm b/Firestore/Example/Tests/Local/FSTMutationQueueTests.mm
index 5b5bdd1..17e6ccb 100644
--- a/Firestore/Example/Tests/Local/FSTMutationQueueTests.mm
+++ b/Firestore/Example/Tests/Local/FSTMutationQueueTests.mm
@@ -31,11 +31,14 @@
 
 #include "Firestore/core/src/firebase/firestore/auth/user.h"
 #include "Firestore/core/src/firebase/firestore/model/document_key.h"
+#include "Firestore/core/src/firebase/firestore/model/document_key_set.h"
 #include "Firestore/core/test/firebase/firestore/testutil/testutil.h"
 
 namespace testutil = firebase::firestore::testutil;
 using firebase::firestore::auth::User;
 using firebase::firestore::model::DocumentKey;
+using firebase::firestore::model::DocumentKeySet;
+using firebase::firestore::testutil::Key;
 
 NS_ASSUME_NONNULL_BEGIN
 
@@ -284,7 +287,7 @@ NS_ASSUME_NONNULL_BEGIN
 
   self.persistence.run("testAllMutationBatchesAffectingDocumentKey", [&]() {
     NSArray<FSTMutation *> *mutations = @[
-      FSTTestSetMutation(@"fob/bar",
+      FSTTestSetMutation(@"foi/bar",
                          @{ @"a" : @1 }),
       FSTTestSetMutation(@"foo/bar",
                          @{ @"a" : @1 }),
@@ -315,6 +318,85 @@ NS_ASSUME_NONNULL_BEGIN
   });
 }
 
+- (void)testAllMutationBatchesAffectingDocumentKeys {
+  if ([self isTestBaseClass]) return;
+
+  self.persistence.run("testAllMutationBatchesAffectingDocumentKey", [&]() {
+    NSArray<FSTMutation *> *mutations = @[
+      FSTTestSetMutation(@"fob/bar",
+                         @{ @"a" : @1 }),
+      FSTTestSetMutation(@"foo/bar",
+                         @{ @"a" : @1 }),
+      FSTTestPatchMutation("foo/bar",
+                           @{ @"b" : @1 }, {}),
+      FSTTestSetMutation(@"foo/bar/suffix/key",
+                         @{ @"a" : @1 }),
+      FSTTestSetMutation(@"foo/baz",
+                         @{ @"a" : @1 }),
+      FSTTestSetMutation(@"food/bar",
+                         @{ @"a" : @1 })
+    ];
+
+    // Store all the mutations.
+    NSMutableArray<FSTMutationBatch *> *batches = [NSMutableArray array];
+    for (FSTMutation *mutation in mutations) {
+      FSTMutationBatch *batch =
+          [self.mutationQueue addMutationBatchWithWriteTime:[FIRTimestamp timestamp]
+                                                  mutations:@[ mutation ]];
+      [batches addObject:batch];
+    }
+
+    DocumentKeySet keys{
+        Key("foo/bar"),
+        Key("foo/baz"),
+    };
+
+    NSArray<FSTMutationBatch *> *expected = @[ batches[1], batches[2], batches[4] ];
+    NSArray<FSTMutationBatch *> *matches =
+        [self.mutationQueue allMutationBatchesAffectingDocumentKeys:keys];
+
+    XCTAssertEqualObjects(matches, expected);
+  });
+}
+
+- (void)testAllMutationBatchesAffectingDocumentKeys_handlesOverlap {
+  if ([self isTestBaseClass]) return;
+
+  self.persistence.run("testAllMutationBatchesAffectingDocumentKeys_handlesOverlap", [&]() {
+    NSArray<FSTMutation *> *group1 = @[
+      FSTTestSetMutation(@"foo/bar",
+                         @{ @"a" : @1 }),
+      FSTTestSetMutation(@"foo/baz",
+                         @{ @"a" : @1 }),
+    ];
+    FSTMutationBatch *batch1 =
+        [self.mutationQueue addMutationBatchWithWriteTime:[FIRTimestamp timestamp]
+                                                mutations:group1];
+
+    NSArray<FSTMutation *> *group2 = @[ FSTTestSetMutation(@"food/bar", @{ @"a" : @1 }) ];
+    [self.mutationQueue addMutationBatchWithWriteTime:[FIRTimestamp timestamp] mutations:group2];
+
+    NSArray<FSTMutation *> *group3 = @[
+      FSTTestSetMutation(@"foo/bar",
+                         @{ @"b" : @1 }),
+    ];
+    FSTMutationBatch *batch3 =
+        [self.mutationQueue addMutationBatchWithWriteTime:[FIRTimestamp timestamp]
+                                                mutations:group3];
+
+    DocumentKeySet keys{
+        Key("foo/bar"),
+        Key("foo/baz"),
+    };
+
+    NSArray<FSTMutationBatch *> *expected = @[ batch1, batch3 ];
+    NSArray<FSTMutationBatch *> *matches =
+        [self.mutationQueue allMutationBatchesAffectingDocumentKeys:keys];
+
+    XCTAssertEqualObjects(matches, expected);
+  });
+}
+
 - (void)testAllMutationBatchesAffectingQuery {
   if ([self isTestBaseClass]) return;
 
diff --git a/Firestore/Source/Local/FSTLevelDBMutationQueue.mm b/Firestore/Source/Local/FSTLevelDBMutationQueue.mm
index 94ab8a5..685d00f 100644
--- a/Firestore/Source/Local/FSTLevelDBMutationQueue.mm
+++ b/Firestore/Source/Local/FSTLevelDBMutationQueue.mm
@@ -30,7 +30,6 @@
 
 #include "Firestore/core/src/firebase/firestore/auth/user.h"
 #include "Firestore/core/src/firebase/firestore/local/leveldb_transaction.h"
-#include "Firestore/core/src/firebase/firestore/model/document_key.h"
 #include "Firestore/core/src/firebase/firestore/model/resource_path.h"
 #include "Firestore/core/src/firebase/firestore/util/hard_assert.h"
 #include "Firestore/core/src/firebase/firestore/util/string_apple.h"
@@ -46,6 +45,7 @@ using firebase::firestore::local::LevelDbTransaction;
 using Firestore::StringView;
 using firebase::firestore::auth::User;
 using firebase::firestore::model::DocumentKey;
+using firebase::firestore::model::DocumentKeySet;
 using firebase::firestore::model::ResourcePath;
 using leveldb::DB;
 using leveldb::Iterator;
@@ -364,11 +364,16 @@ using leveldb::WriteOptions;
   NSMutableArray *result = [NSMutableArray array];
   FSTLevelDBDocumentMutationKey *rowKey = [[FSTLevelDBDocumentMutationKey alloc] init];
   for (; indexIterator->Valid(); indexIterator->Next()) {
-    // Only consider rows matching exactly the specific key of interest. Note that because we order
-    // by path first, and we order terminators before path separators, we'll encounter all the
-    // index rows for documentKey contiguously. In particular, all the rows for documentKey will
-    // occur before any rows for documents nested in a subcollection beneath documentKey so we can
-    // stop as soon as we hit any such row.
+    // Only consider rows matching exactly the specific key of interest. Index rows have this
+    // form (with markers in brackets):
+    //
+    // <User>user <Path>collection <Path>doc <BatchId>2 <Terminator>
+    // <User>user <Path>collection <Path>doc <BatchId>3 <Terminator>
+    // <User>user <Path>collection <Path>doc <Path>sub <Path>doc <BatchId>3 <Terminator>
+    //
+    // Note that Path markers sort after BatchId markers so this means that when searching for
+    // collection/doc, all the entries for it will be contiguous in the table, allowing a break
+    // after any mismatch.
     if (!absl::StartsWith(indexIterator->key(), indexPrefix) ||
         ![rowKey decodeKey:indexIterator->key()] ||
         DocumentKey{rowKey.documentKey} != documentKey) {
@@ -396,6 +401,43 @@ using leveldb::WriteOptions;
   return result;
 }
 
+- (NSArray<FSTMutationBatch *> *)allMutationBatchesAffectingDocumentKeys:
+    (const DocumentKeySet &)documentKeys {
+  NSString *userID = self.userID;
+
+  // Take a pass through the document keys and collect the set of unique mutation batchIDs that
+  // affect them all. Some batches can affect more than one key.
+  std::set<FSTBatchID> batchIDs;
+
+  auto indexIterator = _db.currentTransaction->NewIterator();
+  FSTLevelDBDocumentMutationKey *rowKey = [[FSTLevelDBDocumentMutationKey alloc] init];
+  for (const DocumentKey &documentKey : documentKeys) {
+    std::string indexPrefix =
+        [FSTLevelDBDocumentMutationKey keyPrefixWithUserID:userID resourcePath:documentKey.path()];
+    for (indexIterator->Seek(indexPrefix); indexIterator->Valid(); indexIterator->Next()) {
+      // Only consider rows matching exactly the specific key of interest. Index rows have this
+      // form (with markers in brackets):
+      //
+      // <User>user <Path>collection <Path>doc <BatchId>2 <Terminator>
+      // <User>user <Path>collection <Path>doc <BatchId>3 <Terminator>
+      // <User>user <Path>collection <Path>doc <Path>sub <Path>doc <BatchId>3 <Terminator>
+      //
+      // Note that Path markers sort after BatchId markers so this means that when searching for
+      // collection/doc, all the entries for it will be contiguous in the table, allowing a break
+      // after any mismatch.
+      if (!absl::StartsWith(indexIterator->key(), indexPrefix) ||
+          ![rowKey decodeKey:indexIterator->key()] ||
+          DocumentKey{rowKey.documentKey} != documentKey) {
+        break;
+      }
+
+      batchIDs.insert(rowKey.batchID);
+    }
+  }
+
+  return [self allMutationBatchesWithBatchIDs:batchIDs];
+}
+
 - (NSArray<FSTMutationBatch *> *)allMutationBatchesAffectingQuery:(FSTQuery *)query {
   HARD_ASSERT(![query isDocumentQuery], "Document queries shouldn't go down this path");
   NSString *userID = self.userID;
@@ -417,11 +459,10 @@ using leveldb::WriteOptions;
   // index for more than a single document so the associated batchIDs will be neither necessarily
   // unique nor in order. This means an efficient simultaneous scan isn't possible.
   std::string indexPrefix =
-      [FSTLevelDBDocumentMutationKey keyPrefixWithUserID:self.userID resourcePath:queryPath];
+      [FSTLevelDBDocumentMutationKey keyPrefixWithUserID:userID resourcePath:queryPath];
   auto indexIterator = _db.currentTransaction->NewIterator();
   indexIterator->Seek(indexPrefix);
 
-  NSMutableArray *result = [NSMutableArray array];
   FSTLevelDBDocumentMutationKey *rowKey = [[FSTLevelDBDocumentMutationKey alloc] init];
 
   // Collect up unique batchIDs encountered during a scan of the index. Use a set<FSTBatchID> to
@@ -430,7 +471,7 @@ using leveldb::WriteOptions;
   // This method is faster than performing lookups of the keys with _db->Get and keeping a hash of
   // batchIDs that have already been looked up. The performance difference is minor for small
   // numbers of keys but > 30% faster for larger numbers of keys.
-  std::set<FSTBatchID> uniqueBatchIds;
+  std::set<FSTBatchID> uniqueBatchIDs;
   for (; indexIterator->Valid(); indexIterator->Next()) {
     if (!absl::StartsWith(indexIterator->key(), indexPrefix) ||
         ![rowKey decodeKey:indexIterator->key()]) {
@@ -444,14 +485,25 @@ using leveldb::WriteOptions;
       continue;
     }
 
-    uniqueBatchIds.insert(rowKey.batchID);
+    uniqueBatchIDs.insert(rowKey.batchID);
   }
 
+  return [self allMutationBatchesWithBatchIDs:uniqueBatchIDs];
+}
+
+/**
+ * Constructs an array of matching batches, sorted by batchID to ensure that multiple mutations
+ * affecting the same document key are applied in order.
+ */
+- (NSArray<FSTMutationBatch *> *)allMutationBatchesWithBatchIDs:
+    (const std::set<FSTBatchID> &)batchIDs {
+  NSMutableArray *result = [NSMutableArray array];
+  NSString *userID = self.userID;
+
   // Given an ordered set of unique batchIDs perform a skipping scan over the main table to find
   // the mutation batches.
   auto mutationIterator = _db.currentTransaction->NewIterator();
-
-  for (FSTBatchID batchID : uniqueBatchIds) {
+  for (FSTBatchID batchID : batchIDs) {
     std::string mutationKey = [FSTLevelDBMutationKey keyWithUserID:userID batchID:batchID];
     mutationIterator->Seek(mutationKey);
     if (!mutationIterator->Valid() || mutationIterator->key() != mutationKey) {
diff --git a/Firestore/Source/Local/FSTMemoryMutationQueue.mm b/Firestore/Source/Local/FSTMemoryMutationQueue.mm
index 8faf893..30caa95 100644
--- a/Firestore/Source/Local/FSTMemoryMutationQueue.mm
+++ b/Firestore/Source/Local/FSTMemoryMutationQueue.mm
@@ -16,6 +16,8 @@
 
 #import "Firestore/Source/Local/FSTMemoryMutationQueue.h"
 
+#include <set>
+
 #import "Firestore/Source/Core/FSTQuery.h"
 #import "Firestore/Source/Local/FSTDocumentReference.h"
 #import "Firestore/Source/Local/FSTMemoryPersistence.h"
@@ -28,6 +30,7 @@
 #include "Firestore/core/src/firebase/firestore/util/hard_assert.h"
 
 using firebase::firestore::model::DocumentKey;
+using firebase::firestore::model::DocumentKeySet;
 using firebase::firestore::model::ResourcePath;
 
 NS_ASSUME_NONNULL_BEGIN
@@ -242,6 +245,28 @@ static const NSComparator NumberComparator = ^NSComparisonResult(NSNumber *left,
   return result;
 }
 
+- (NSArray<FSTMutationBatch *> *)allMutationBatchesAffectingDocumentKeys:
+    (const DocumentKeySet &)documentKeys {
+  // First find the set of affected batch IDs.
+  __block std::set<FSTBatchID> batchIDs;
+  for (const DocumentKey &key : documentKeys) {
+    FSTDocumentReference *start = [[FSTDocumentReference alloc] initWithKey:key ID:0];
+
+    FSTDocumentReferenceBlock block = ^(FSTDocumentReference *reference, BOOL *stop) {
+      if (![key isEqualToKey:reference.key]) {
+        *stop = YES;
+        return;
+      }
+
+      batchIDs.insert(reference.ID);
+    };
+
+    [self.batchesByDocumentKey enumerateObjectsFrom:start to:nil usingBlock:block];
+  }
+
+  return [self allMutationBatchesWithBatchIDs:batchIDs];
+}
+
 - (NSArray<FSTMutationBatch *> *)allMutationBatchesAffectingQuery:(FSTQuery *)query {
   // Use the query path as a prefix for testing if a document matches the query.
   const ResourcePath &prefix = query.path;
@@ -258,8 +283,7 @@ static const NSComparator NumberComparator = ^NSComparisonResult(NSNumber *left,
       [[FSTDocumentReference alloc] initWithKey:DocumentKey{startPath} ID:0];
 
   // Find unique batchIDs referenced by all documents potentially matching the query.
-  __block FSTImmutableSortedSet<NSNumber *> *uniqueBatchIDs =
-      [FSTImmutableSortedSet setWithComparator:NumberComparator];
+  __block std::set<FSTBatchID> uniqueBatchIDs;
   FSTDocumentReferenceBlock block = ^(FSTDocumentReference *reference, BOOL *stop) {
     const ResourcePath &rowKeyPath = reference.key.path();
     if (!prefix.IsPrefixOf(rowKeyPath)) {
@@ -274,19 +298,26 @@ static const NSComparator NumberComparator = ^NSComparisonResult(NSNumber *left,
       return;
     }
 
-    uniqueBatchIDs = [uniqueBatchIDs setByAddingObject:@(reference.ID)];
+    uniqueBatchIDs.insert(reference.ID);
   };
   [self.batchesByDocumentKey enumerateObjectsFrom:start to:nil usingBlock:block];
 
-  // Construct an array of matching batches, sorted by batchID to ensure that multiple mutations
-  // affecting the same document key are applied in order.
+  return [self allMutationBatchesWithBatchIDs:uniqueBatchIDs];
+}
+
+/**
+ * Constructs an array of matching batches, sorted by batchID to ensure that multiple mutations
+ * affecting the same document key are applied in order.
+ */
+- (NSArray<FSTMutationBatch *> *)allMutationBatchesWithBatchIDs:
+    (const std::set<FSTBatchID> &)batchIDs {
   NSMutableArray<FSTMutationBatch *> *result = [NSMutableArray array];
-  [uniqueBatchIDs enumerateObjectsUsingBlock:^(NSNumber *batchID, BOOL *stop) {
-    FSTMutationBatch *batch = [self lookupMutationBatch:[batchID intValue]];
+  for (FSTBatchID batchID : batchIDs) {
+    FSTMutationBatch *batch = [self lookupMutationBatch:batchID];
     if (batch) {
       [result addObject:batch];
     }
-  }];
+  };
 
   return result;
 }
diff --git a/Firestore/Source/Local/FSTMutationQueue.h b/Firestore/Source/Local/FSTMutationQueue.h
index 89b1a52..7d117b5 100644
--- a/Firestore/Source/Local/FSTMutationQueue.h
+++ b/Firestore/Source/Local/FSTMutationQueue.h
@@ -20,6 +20,7 @@
 #import "Firestore/Source/Local/FSTGarbageCollector.h"
 
 #include "Firestore/core/src/firebase/firestore/model/document_key.h"
+#include "Firestore/core/src/firebase/firestore/model/document_key_set.h"
 
 @class FSTMutation;
 @class FSTMutationBatch;
@@ -115,11 +116,22 @@ NS_ASSUME_NONNULL_BEGIN
  * don't contain the document key at all if it's convenient.
  */
 // TODO(mcg): This should really return an NSEnumerator
-// also for b/32992024, all backing stores should really index by document key
 - (NSArray<FSTMutationBatch *> *)allMutationBatchesAffectingDocumentKey:
     (const firebase::firestore::model::DocumentKey &)documentKey;
 
 /**
+ * Finds all mutation batches that could @em possibly affect the given document keys. Not all
+ * mutations in a batch will necessarily affect each key, so when looping through the batches you'll
+ * need to check that the mutation itself matches the key.
+ *
+ * Note that because of this requirement implementations are free to return mutation batches that
+ * don't contain any of the given document keys at all if it's convenient.
+ */
+// TODO(mcg): This should really return an NSEnumerator
+- (NSArray<FSTMutationBatch *> *)allMutationBatchesAffectingDocumentKeys:
+    (const firebase::firestore::model::DocumentKeySet &)documentKeys;
+
+/**
  * Finds all mutation batches that could affect the results for the given query. Not all
  * mutations in a batch will necessarily affect the query, so when looping through the batch
  * you'll need to check that the mutation itself matches the query.