Add allMutationsAffectingDocumentKeys

Author: wilhuff

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
 
@@ -315,6 +318,87 @@ - (void)testAllMutationBatchesAffectingDocumentKey {
   });
 }
 
+- (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", [&]() {
+    NSMutableArray<FSTMutationBatch *> *batches = [NSMutableArray array];
+
+    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;
 

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 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;
@@ -396,6 +396,39 @@ - (nullable FSTMutationBatch *)nextMutationBatchAfterBatchID:(FSTBatchID)batchID
   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()];
+    indexIterator->Seek(indexPrefix);
+    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.
+      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 +450,10 @@ - (nullable FSTMutationBatch *)nextMutationBatchAfterBatchID:(FSTBatchID)batchID
   // 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 +462,7 @@ - (nullable FSTMutationBatch *)nextMutationBatchAfterBatchID:(FSTBatchID)batchID
   // 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 +476,25 @@ - (nullable FSTMutationBatch *)nextMutationBatchAfterBatchID:(FSTBatchID)batchID
       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) {

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 @@ - (nullable FSTMutationBatch *)nextMutationBatchAfterBatchID:(FSTBatchID)batchID
   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 @@ - (nullable FSTMutationBatch *)nextMutationBatchAfterBatchID:(FSTBatchID)batchID
       [[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 @@ - (nullable FSTMutationBatch *)nextMutationBatchAfterBatchID:(FSTBatchID)batchID
       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;
 }

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,10 +116,21 @@ 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