Change removeMutationBatch to remove a single batch

Firestore/Example/Tests/Local/FSTMutationQueueTests.mm

@@ -71,10 +71,10 @@ - (void)testCountBatches {
     FSTMutationBatch *batch2 = [self addMutationBatch];
     XCTAssertEqual(2, [self batchCount]);
 
-    [self.mutationQueue removeMutationBatches:@[ batch2 ]];
+    [self.mutationQueue removeMutationBatch:batch2];
     XCTAssertEqual(1, [self batchCount]);
 
-    [self.mutationQueue removeMutationBatches:@[ batch1 ]];
+    [self.mutationQueue removeMutationBatch:batch1];
     XCTAssertEqual(0, [self batchCount]);
     XCTAssertTrue([self.mutationQueue isEmpty]);
   });
@@ -101,14 +101,14 @@ - (void)testAcknowledgeBatchID {
     [self.mutationQueue acknowledgeBatch:batch2 streamToken:nil];
     XCTAssertEqual([self.mutationQueue highestAcknowledgedBatchID], batch2.batchID);
 
-    [self.mutationQueue removeMutationBatches:@[ batch1 ]];
+    [self.mutationQueue removeMutationBatch:batch1];
     XCTAssertEqual([self.mutationQueue highestAcknowledgedBatchID], batch2.batchID);
 
-    [self.mutationQueue removeMutationBatches:@[ batch2 ]];
+    [self.mutationQueue removeMutationBatch:batch2];
     XCTAssertEqual([self.mutationQueue highestAcknowledgedBatchID], batch2.batchID);
 
     // Batch 3 never acknowledged.
-    [self.mutationQueue removeMutationBatches:@[ batch3 ]];
+    [self.mutationQueue removeMutationBatch:batch3];
     XCTAssertEqual([self.mutationQueue highestAcknowledgedBatchID], batch2.batchID);
   });
 }
@@ -120,7 +120,7 @@ - (void)testAcknowledgeThenRemove {
     FSTMutationBatch *batch1 = [self addMutationBatch];
 
     [self.mutationQueue acknowledgeBatch:batch1 streamToken:nil];
-    [self.mutationQueue removeMutationBatches:@[ batch1 ]];
+    [self.mutationQueue removeMutationBatch:batch1];
 
     XCTAssertEqual([self batchCount], 0);
     XCTAssertEqual([self.mutationQueue highestAcknowledgedBatchID], batch1.batchID);
@@ -142,7 +142,8 @@ - (void)testHighestAcknowledgedBatchIDNeverExceedsNextBatchID {
     [self.mutationQueue acknowledgeBatch:batch2 streamToken:nil];
     XCTAssertEqual([self.mutationQueue highestAcknowledgedBatchID], batch2.batchID);
 
-    [self.mutationQueue removeMutationBatches:@[ batch1, batch2 ]];
+    [self.mutationQueue removeMutationBatch:batch1];
+    [self.mutationQueue removeMutationBatch:batch2];
     XCTAssertEqual([self.mutationQueue highestAcknowledgedBatchID], batch2.batchID);
   });
 
@@ -439,7 +440,7 @@ - (void)testRemoveMutationBatches {
   self.persistence.run("testRemoveMutationBatches", [&]() {
     NSMutableArray<FSTMutationBatch *> *batches = [self createBatches:10];
 
-    [self.mutationQueue removeMutationBatches:@[ batches[0] ]];
+    [self.mutationQueue removeMutationBatch:batches[0]];
     [batches removeObjectAtIndex:0];
 
     FSTMutationBatch *last = batches[batches.count - 1];
@@ -451,27 +452,29 @@ - (void)testRemoveMutationBatches {
     XCTAssertEqualObjects(found, batches);
     XCTAssertEqual(found.count, 9);
 
-    [self.mutationQueue removeMutationBatches:@[ batches[0], batches[1], batches[2] ]];
+    [self.mutationQueue removeMutationBatch:batches[0]];
+    [self.mutationQueue removeMutationBatch:batches[1]];
+    [self.mutationQueue removeMutationBatch:batches[2]];
     [batches removeObjectsInRange:NSMakeRange(0, 3)];
     XCTAssertEqual([self batchCount], 6);
 
     found = [self.mutationQueue allMutationBatchesThroughBatchID:last.batchID];
     XCTAssertEqualObjects(found, batches);
     XCTAssertEqual(found.count, 6);
 
-    [self.mutationQueue removeMutationBatches:@[ batches[batches.count - 1] ]];
+    [self.mutationQueue removeMutationBatch:batches[batches.count - 1]];
     [batches removeObjectAtIndex:batches.count - 1];
     XCTAssertEqual([self batchCount], 5);
 
     found = [self.mutationQueue allMutationBatchesThroughBatchID:last.batchID];
     XCTAssertEqualObjects(found, batches);
     XCTAssertEqual(found.count, 5);
 
-    [self.mutationQueue removeMutationBatches:@[ batches[3] ]];
+    [self.mutationQueue removeMutationBatch:batches[3]];
     [batches removeObjectAtIndex:3];
     XCTAssertEqual([self batchCount], 4);
 
-    [self.mutationQueue removeMutationBatches:@[ batches[1] ]];
+    [self.mutationQueue removeMutationBatch:batches[1]];
     [batches removeObjectAtIndex:1];
     XCTAssertEqual([self batchCount], 3);
 
@@ -480,7 +483,9 @@ - (void)testRemoveMutationBatches {
     XCTAssertEqual(found.count, 3);
     XCTAssertFalse([self.mutationQueue isEmpty]);
 
-    [self.mutationQueue removeMutationBatches:batches];
+    for (FSTMutationBatch *batch in batches) {
+      [self.mutationQueue removeMutationBatch:batch];
+    }
     found = [self.mutationQueue allMutationBatchesThroughBatchID:last.batchID];
     XCTAssertEqualObjects(found, @[]);
     XCTAssertEqual(found.count, 0);
@@ -562,7 +567,7 @@ - (NSUInteger)batchCount {
   for (NSUInteger i = 0; i < holes.count; i++) {
     NSUInteger index = holes[i].unsignedIntegerValue - i;
     FSTMutationBatch *batch = batches[index];
-    [self.mutationQueue removeMutationBatches:@[ batch ]];
+    [self.mutationQueue removeMutationBatch:batch];
 
     [batches removeObjectAtIndex:index];
     [removed addObject:batch];

Firestore/Source/Local/FSTLevelDBMutationQueue.mm

@@ -521,27 +521,25 @@ - (nullable FSTMutationBatch *)nextMutationBatchAfterBatchID:(BatchId)batchID {
   return result;
 }
 
-- (void)removeMutationBatches:(NSArray<FSTMutationBatch *> *)batches {
+- (void)removeMutationBatch:(FSTMutationBatch *)batch {
   auto checkIterator = _db.currentTransaction->NewIterator();
 
-  for (FSTMutationBatch *batch in batches) {
-    BatchId batchID = batch.batchID;
-    std::string key = LevelDbMutationKey::Key(_userID, batchID);
+  BatchId batchID = batch.batchID;
+  std::string key = LevelDbMutationKey::Key(_userID, batchID);
 
-    // As a sanity check, verify that the mutation batch exists before deleting it.
-    checkIterator->Seek(key);
-    HARD_ASSERT(checkIterator->Valid(), "Mutation batch %s did not exist", DescribeKey(key));
+  // As a sanity check, verify that the mutation batch exists before deleting it.
+  checkIterator->Seek(key);
+  HARD_ASSERT(checkIterator->Valid(), "Mutation batch %s did not exist", DescribeKey(key));
 
-    HARD_ASSERT(key == checkIterator->key(), "Mutation batch %s not found; found %s",
-                DescribeKey(key), DescribeKey(checkIterator));
+  HARD_ASSERT(key == checkIterator->key(), "Mutation batch %s not found; found %s",
+              DescribeKey(key), DescribeKey(checkIterator));
 
-    _db.currentTransaction->Delete(key);
+  _db.currentTransaction->Delete(key);
 
-    for (FSTMutation *mutation in batch.mutations) {
-      key = LevelDbDocumentMutationKey::Key(_userID, mutation.key, batchID);
-      _db.currentTransaction->Delete(key);
-      [_db.referenceDelegate removeMutationReference:mutation.key];
-    }
+  for (FSTMutation *mutation in batch.mutations) {
+    key = LevelDbDocumentMutationKey::Key(_userID, mutation.key, batchID);
+    _db.currentTransaction->Delete(key);
+    [_db.referenceDelegate removeMutationReference:mutation.key];
   }
 }
 

Firestore/Source/Local/FSTLocalStore.mm

@@ -148,7 +148,7 @@ - (void)startMutationQueue {
       if (batches.count > 0) {
         // NOTE: This could be more efficient if we had a removeBatchesThroughBatchID, but this set
         // should be very small and this code should go away eventually.
-        [self.mutationQueue removeMutationBatches:batches];
+        [self removeMutationBatches:batches];
       }
     }
   });
@@ -533,9 +533,9 @@ - (DocumentKeySet)removeMutationBatches:(NSArray<FSTMutationBatch *> *)batches {
       const DocumentKey &key = mutation.key;
       affectedDocs = affectedDocs.insert(key);
     }
+    [self.mutationQueue removeMutationBatch:batch];
   }
 
-  [self.mutationQueue removeMutationBatches:batches];
   return affectedDocs;
 }
 

Firestore/Source/Local/FSTMemoryMutationQueue.mm

@@ -326,67 +326,42 @@ - (nullable FSTMutationBatch *)nextMutationBatchAfterBatchID:(BatchId)batchID {
   return result;
 }
 
-- (void)removeMutationBatches:(NSArray<FSTMutationBatch *> *)batches {
-  NSUInteger batchCount = batches.count;
-  HARD_ASSERT(batchCount > 0, "Should not remove mutations when none exist.");
-
-  BatchId firstBatchID = batches[0].batchID;
-
+- (void)removeMutationBatch:(FSTMutationBatch *)batch {
   NSMutableArray<FSTMutationBatch *> *queue = self.queue;
   NSUInteger queueCount = queue.count;
+  BatchId batchID = batch.batchID;
 
   // Find the position of the first batch for removal. This need not be the first entry in the
   // queue.
-  NSUInteger startIndex = [self indexOfExistingBatchID:firstBatchID action:@"removed"];
-  HARD_ASSERT(queue[startIndex].batchID == firstBatchID, "Removed batches must exist in the queue");
-
-  // Check that removed batches are contiguous (while excluding tombstones).
-  NSUInteger batchIndex = 1;
-  NSUInteger queueIndex = startIndex + 1;
-  while (batchIndex < batchCount && queueIndex < queueCount) {
-    FSTMutationBatch *batch = queue[queueIndex];
-    if ([batch isTombstone]) {
-      queueIndex++;
-      continue;
-    }
-
-    HARD_ASSERT(batch.batchID == batches[batchIndex].batchID,
-                "Removed batches must be contiguous in the queue");
-    batchIndex++;
-    queueIndex++;
-  }
+  NSUInteger batchIndex = [self indexOfExistingBatchID:batchID action:@"removed"];
+  HARD_ASSERT(queue[batchIndex].batchID == batchID, "Removed batches must exist in the queue");
 
   // Only actually remove batches if removing at the front of the queue. Previously rejected batches
   // may have left tombstones in the queue, so expand the removal range to include any tombstones.
-  if (startIndex == 0) {
-    for (; queueIndex < queueCount; queueIndex++) {
-      FSTMutationBatch *batch = queue[queueIndex];
+  if (batchIndex == 0) {
+    int endIndex = 1;
+    for (; endIndex < queueCount; endIndex++) {
+      FSTMutationBatch *batch = queue[endIndex];
       if (![batch isTombstone]) {
         break;
       }
     }
 
-    NSUInteger length = queueIndex - startIndex;
-    [queue removeObjectsInRange:NSMakeRange(startIndex, length)];
+    NSUInteger length = endIndex - batchIndex;
+    [queue removeObjectsInRange:NSMakeRange(batchIndex, length)];
 
   } else {
-    // Mark tombstones
-    for (NSUInteger i = startIndex; i < queueIndex; i++) {
-      queue[i] = [queue[i] toTombstone];
-    }
+    queue[batchIndex] = [queue[batchIndex] toTombstone];
   }
 
   // Remove entries from the index too.
   FSTImmutableSortedSet<FSTDocumentReference *> *references = self.batchesByDocumentKey;
-  for (FSTMutationBatch *batch in batches) {
-    BatchId batchID = batch.batchID;
-    for (FSTMutation *mutation in batch.mutations) {
-      const DocumentKey &key = mutation.key;
-      [_persistence.referenceDelegate removeMutationReference:key];
-
-      FSTDocumentReference *reference = [[FSTDocumentReference alloc] initWithKey:key ID:batchID];
-      references = [references setByRemovingObject:reference];
-    }
+  for (FSTMutation *mutation in batch.mutations) {
+    const DocumentKey &key = mutation.key;
+    [_persistence.referenceDelegate removeMutationReference:key];
+
+    FSTDocumentReference *reference = [[FSTDocumentReference alloc] initWithKey:key ID:batchID];
+    references = [references setByRemovingObject:reference];
   }
   self.batchesByDocumentKey = references;
 }

Firestore/Source/Local/FSTMutationQueue.h

@@ -147,15 +147,12 @@ NS_ASSUME_NONNULL_BEGIN
 - (NSArray<FSTMutationBatch *> *)allMutationBatchesAffectingQuery:(FSTQuery *)query;
 
 /**
- * Removes the given mutation batches from the queue. This is useful in two circumstances:
+ * Removes the given mutation batch from the queue. This is useful in two circumstances:
  *
  * + Removing applied mutations from the head of the queue
  * + Removing rejected mutations from anywhere in the queue
- *
- * In both cases, the array of mutations to remove must be a contiguous range of batchIds. This is
- * most easily accomplished by loading mutations with @a -allMutationBatchesThroughBatchID:.
  */
-- (void)removeMutationBatches:(NSArray<FSTMutationBatch *> *)batches;
+- (void)removeMutationBatch:(FSTMutationBatch *)batch;
 
 /** Performs a consistency check, examining the mutation queue for any leaks, if possible. */
 - (void)performConsistencyCheck;