[ASDataController] Call All Completion Handlers after Nested Batch Updates (#2274)

* Add tests for batch update completion handler calling

* Ensure we call all completion handlers after collection view updates

* Tweak it

* Fix the doc

* Minor improvements

* Document addCompletionHandler better

AsyncDisplayKit/ASCollectionNode.h

@@ -91,7 +91,7 @@ NS_ASSUME_NONNULL_BEGIN
  * the main thread.
  * @warning This method is substantially more expensive than UICollectionView's version.
  */
-- (void)reloadDataWithCompletion:(void (^)())completion;
+- (void)reloadDataWithCompletion:(nullable void (^)())completion;
 
 /**
  * Reload everything from scratch, destroying the working range and all cached nodes.

AsyncDisplayKit/ASCollectionView.h

@@ -128,7 +128,7 @@ NS_ASSUME_NONNULL_BEGIN
  *                    Boolean parameter that contains the value YES if all of the related animations completed successfully or 
  *                    NO if they were interrupted. This parameter may be nil. If supplied, the block is run on the main thread.
  */
-- (void)performBatchAnimated:(BOOL)animated updates:(void (^ _Nullable)())updates completion:(void (^ _Nullable)(BOOL))completion;
+- (void)performBatchAnimated:(BOOL)animated updates:(nullable __attribute((noescape)) void (^)())updates completion:(nullable void (^)(BOOL finished))completion;
 
 /**
  *  Perform a batch of updates asynchronously.  This method must be called from the main thread.
@@ -139,7 +139,7 @@ NS_ASSUME_NONNULL_BEGIN
  *                    Boolean parameter that contains the value YES if all of the related animations completed successfully or
  *                    NO if they were interrupted. This parameter may be nil. If supplied, the block is run on the main thread.
  */
-- (void)performBatchUpdates:(void (^ _Nullable)())updates completion:(void (^ _Nullable)(BOOL))completion;
+- (void)performBatchUpdates:(nullable __attribute((noescape)) void (^)())updates completion:(nullable void (^)(BOOL finished))completion;
 
 /**
  * Reload everything from scratch, destroying the working range and all cached nodes.
@@ -148,7 +148,7 @@ NS_ASSUME_NONNULL_BEGIN
  * the main thread.
  * @warning This method is substantially more expensive than UICollectionView's version.
  */
-- (void)reloadDataWithCompletion:(void (^ _Nullable)())completion;
+- (void)reloadDataWithCompletion:(nullable void (^)())completion;
 
 /**
  * Reload everything from scratch, destroying the working range and all cached nodes.

AsyncDisplayKit/Details/ASChangeSetDataController.mm

@@ -20,6 +20,11 @@
   _ASHierarchyChangeSet *_changeSet;
 }
 
+- (void)dealloc
+{
+  ASDisplayNodeCAssert(_changeSetBatchUpdateCounter == 0, @"ASChangeSetDataController deallocated in the middle of a batch update.");
+}
+
 #pragma mark - Batching (External API)
 
 - (void)beginUpdates
@@ -40,18 +45,20 @@
   // Prevent calling endUpdatesAnimated:completion: in an unbalanced way
   NSAssert(_changeSetBatchUpdateCounter >= 0, @"endUpdatesAnimated:completion: called without having a balanced beginUpdates call");
   
+  [_changeSet addCompletionHandler:completion];
   if (_changeSetBatchUpdateCounter == 0) {
+    [self invalidateDataSourceItemCounts];
+    [_changeSet markCompletedWithNewItemCounts:[self itemCountsFromDataSource]];
+    void (^batchCompletion)(BOOL finished) = _changeSet.completionHandler;
+
     if (!self.initialReloadDataHasBeenCalled) {
-      if (completion) {
+      if (batchCompletion != nil) {
-        completion(YES);
+        batchCompletion(YES);
       }
       _changeSet = nil;
       return;
     }
     
-    [self invalidateDataSourceItemCounts];
-    [_changeSet markCompletedWithNewItemCounts:[self itemCountsFromDataSource]];
-    
     [super beginUpdates];
     
     for (_ASHierarchyItemChange *change in [_changeSet itemChangesOfType:_ASHierarchyChangeTypeDelete]) {
@@ -70,7 +77,7 @@
       [super insertRowsAtIndexPaths:change.indexPaths withAnimationOptions:change.animationOptions];
     }
 
-    [super endUpdatesAnimated:animated completion:completion];
+    [super endUpdatesAnimated:animated completion:batchCompletion];
     
     _changeSet = nil;
   }

AsyncDisplayKit/Private/_ASHierarchyChangeSet.h

@@ -94,6 +94,26 @@ NSString *NSStringFromASHierarchyChangeType(_ASHierarchyChangeType changeType);
 
 - (instancetype)initWithOldData:(std::vector<NSInteger>)oldItemCounts NS_DESIGNATED_INITIALIZER;
 
+
+/**
+ * The combined completion handler.
+ *
+ * @precondition The change set must be completed.
+ * @warning The completion block is discarded after reading because it may have captured
+ *   significant resources that we would like to reclaim as soon as possible.
+ */
+@property (nonatomic, copy, readonly, nullable) void(^completionHandler)(BOOL finished);
+
+/**
+ * Append the given completion handler to the combined @c completionHandler.
+ *
+ * @discussion Since batch updates can be nested, we have to support multiple
+ * completion handlers per update.
+ *
+ * @precondition The change set must not be completed.
+ */
+- (void)addCompletionHandler:(nullable void(^)(BOOL finished))completion;
+
 /// @precondition The change set must be completed.
 @property (nonatomic, strong, readonly) NSIndexSet *deletedSections;
 /// @precondition The change set must be completed.

AsyncDisplayKit/Private/_ASHierarchyChangeSet.mm

@@ -101,6 +101,7 @@ NSString *NSStringFromASHierarchyChangeType(_ASHierarchyChangeType changeType)
 @implementation _ASHierarchyChangeSet {
   std::vector<NSInteger> _oldItemCounts;
   std::vector<NSInteger> _newItemCounts;
+  void (^_completionHandler)(BOOL finished);
 }
 
 - (instancetype)init
@@ -132,6 +133,31 @@ NSString *NSStringFromASHierarchyChangeType(_ASHierarchyChangeType changeType)
 
 #pragma mark External API
 
+- (void (^)(BOOL finished))completionHandler
+{
+  [self _ensureCompleted];
+
+  void (^completionHandler)(BOOL) = _completionHandler;
+  _completionHandler = nil;
+  return completionHandler;
+}
+
+- (void)addCompletionHandler:(void (^)(BOOL))completion
+{
+  [self _ensureNotCompleted];
+  if (completion == nil) {
+    return;
+  }
+
+  void (^oldCompletionHandler)(BOOL finished) = _completionHandler;
+  _completionHandler = ^(BOOL finished) {
+    if (oldCompletionHandler != nil) {
+    	oldCompletionHandler(finished);
+    }
+    completion(finished);
+  };
+}
+
 - (void)markCompletedWithNewItemCounts:(std::vector<NSInteger>)newItemCounts
 {
   NSAssert(!_completed, @"Attempt to mark already-completed changeset as completed.");

AsyncDisplayKitTests/ASCollectionViewTests.mm

@@ -441,4 +441,49 @@
   [[UIDevice currentDevice] setValue:@(oldDeviceOrientation) forKey:@"orientation"];
 }
 
+/**
+ * See corresponding test in ASUICollectionViewTests
+ *
+ * @discussion Currently, we do not replicate UICollectionView's call order (outer, inner0, inner1, ...)
+ *   and instead call (inner0, inner1, outer, ...). This is because we primarily provide a
+ *   beginUpdates/endUpdatesWithCompletion: interface (like UITableView). With UICollectionView's
+ *   performBatchUpdates:completion:, the completion block is enqueued at -beginUpdates time.
+ *   With our tableView-like scheme, the completion block is provided at -endUpdates time
+ *   and it is naturally enqueued at this time. It is assumed that this is an acceptable deviation,
+ *   and that developers do not expect a particular completion order guarantee.
+ */
+- (void)testThatNestedBatchCompletionsAreCalledInOrder
+{
+  ASCollectionViewTestController *testController = [[ASCollectionViewTestController alloc] initWithNibName:nil bundle:nil];
+
+  ASCollectionView *cv = testController.collectionView;
+
+  XCTestExpectation *inner0 = [self expectationWithDescription:@"Inner completion 0 is called"];
+  XCTestExpectation *inner1 = [self expectationWithDescription:@"Inner completion 1 is called"];
+  XCTestExpectation *outer = [self expectationWithDescription:@"Outer completion is called"];
+
+  NSMutableArray<XCTestExpectation *> *completions = [NSMutableArray array];
+
+  [cv performBatchUpdates:^{
+    [cv performBatchUpdates:^{
+
+    } completion:^(BOOL finished) {
+      [completions addObject:inner0];
+      [inner0 fulfill];
+    }];
+    [cv performBatchUpdates:^{
+
+    } completion:^(BOOL finished) {
+      [completions addObject:inner1];
+      [inner1 fulfill];
+    }];
+  } completion:^(BOOL finished) {
+    [completions addObject:outer];
+    [outer fulfill];
+  }];
+
+  [self waitForExpectationsWithTimeout:5 handler:nil];
+  XCTAssertEqualObjects(completions, (@[ inner0, inner1, outer ]), @"Expected completion order to be correct");
+}
+
 @end

AsyncDisplayKitTests/ASUICollectionViewTests.m

@@ -34,6 +34,45 @@
   [self _testSupplementaryNodeAtIndexPath:[NSIndexPath indexPathWithIndex:3] sectionCount:2 expectException:NO];
 }
 
+- (void)testThatNestedBatchCompletionsAreCalledInOrder
+{
+  UICollectionViewLayout *layout = [[UICollectionViewLayout alloc] init];
+  id layoutMock = [OCMockObject partialMockForObject:layout];
+
+  UICollectionView *cv = [[UICollectionView alloc] initWithFrame:CGRectMake(0, 0, 100, 100) collectionViewLayout:layoutMock];
+  id dataSource = [OCMockObject niceMockForProtocol:@protocol(UICollectionViewDataSource)];
+  [[[dataSource stub] andReturnValue:[NSNumber numberWithInteger:1]] collectionView:cv numberOfItemsInSection:0];
+
+  cv.dataSource = dataSource;
+
+  XCTestExpectation *inner0 = [self expectationWithDescription:@"Inner completion 0 is called"];
+  XCTestExpectation *inner1 = [self expectationWithDescription:@"Inner completion 1 is called"];
+  XCTestExpectation *outer = [self expectationWithDescription:@"Outer completion is called"];
+
+  NSMutableArray<XCTestExpectation *> *completions = [NSMutableArray array];
+
+  [cv performBatchUpdates:^{
+    [cv performBatchUpdates:^{
+
+    } completion:^(BOOL finished) {
+      [completions addObject:inner0];
+      [inner0 fulfill];
+    }];
+    [cv performBatchUpdates:^{
+
+    } completion:^(BOOL finished) {
+      [completions addObject:inner1];
+      [inner1 fulfill];
+    }];
+  } completion:^(BOOL finished) {
+    [completions addObject:outer];
+    [outer fulfill];
+  }];
+
+  [self waitForExpectationsWithTimeout:5 handler:nil];
+  XCTAssertEqualObjects(completions, (@[ outer, inner0, inner1 ]), @"Expected completion order to be correct");
+}
+
 - (void)_testSupplementaryNodeAtIndexPath:(NSIndexPath *)indexPath sectionCount:(NSInteger)sectionCount expectException:(BOOL)shouldFail
 {
   UICollectionViewLayoutAttributes *attr = [UICollectionViewLayoutAttributes layoutAttributesForSupplementaryViewOfKind:@"SuppKind" withIndexPath:indexPath];