Prevent deallocation of asyncDataSource and asyncDelegate in ASCollectionView and ASTableView

Grab a strong reference for asyncDataSource and asyncDelegate in ASCollectionView and ASTableView before executing the range update to be sure they are not going away while executing the range update. This can happen in range updates while going back in the view controller hierarchy

AsyncDisplayKit/ASCollectionView.mm

@@ -1139,7 +1139,14 @@ static NSString * const kCellReuseIdentifier = @"_ASCollectionViewCell";
   // Updating the visible node index paths only for not range managed nodes. Range managed nodes will get their
   // their update in the layout pass
   if (![node supportsRangeManagedInterfaceState]) {
-    [_rangeController visibleNodeIndexPathsDidChangeWithScrollDirection:self.scrollDirection];
+    // Grab a strong reference for data source and delegate to be sure they are not going away while executing
+    // the range update. This can happen in range updates while going back in the view controller hierarchy
+    __block id<ASCollectionDataSource> asyncDataSource = _asyncDataSource;
+    __block id<ASCollectionDelegate> asyncDelegate = _asyncDelegate;
+    [_rangeController scheduleRangeUpdateCompletion:^{
+      asyncDataSource = nil;
+      asyncDelegate = nil;
+    }];
   }
 }
 

AsyncDisplayKit/ASTableView.mm

@@ -1072,7 +1072,14 @@ static NSString * const kCellReuseIdentifier = @"_ASTableViewCell";
   // Updating the visible node index paths only for not range managed nodes. Range managed nodes will get their
   // their update in the layout pass
   if (![node supportsRangeManagedInterfaceState]) {
-    [_rangeController visibleNodeIndexPathsDidChangeWithScrollDirection:self.scrollDirection];
+    // Grab a strong reference for data source and delegate to be sure they are not going away while executing
+    // the range update. This can happen in range updates while going back in the view controller hierarchy
+    __block id<ASTableDataSource> asyncDataSource = _asyncDataSource;
+    __block id<ASTableDelegate> asyncDelegate = _asyncDelegate;
+    [_rangeController scheduleRangeUpdateCompletion:^{
+      asyncDataSource = nil;
+      asyncDelegate = nil;
+    }];
   }
 }
 

AsyncDisplayKit/Details/ASRangeController.mm

@@ -27,6 +27,7 @@
   BOOL _didUpdateCurrentRange;
   BOOL _didRegisterForNotifications;
   CFAbsoluteTime _pendingDisplayNodesTimestamp;
+  NSMutableArray *_scheduledRangeUpdateCompletionBlocks;
 }
 
 @end
@@ -46,6 +47,7 @@ static UIApplicationState __ApplicationState = UIApplicationStateActive;
   _rangeIsValid = YES;
   _currentRangeMode = ASLayoutRangeModeInvalid;
   _didUpdateCurrentRange = NO;
+  _scheduledRangeUpdateCompletionBlocks = [NSMutableArray array];
   
   [[[self class] allRangeControllersWeakSet] addObject:self];
   
@@ -111,6 +113,15 @@ static UIApplicationState __ApplicationState = UIApplicationStateActive;
 
 - (void)scheduleRangeUpdate
 {
+  [self scheduleRangeUpdateCompletion:nil];
+}
+
+- (void)scheduleRangeUpdateCompletion:(void (^)(void))completion
+{
+  if (completion) {
+    [_scheduledRangeUpdateCompletionBlocks addObject:completion];
+  }
+ 
   if (_queuedRangeUpdate) {
     return;
   }
@@ -118,8 +129,13 @@ static UIApplicationState __ApplicationState = UIApplicationStateActive;
   // coalesce these events -- handling them multiple times per runloop is noisy and expensive
   _queuedRangeUpdate = YES;
   
+  __block id<ASRangeControllerDataSource> dataSource = _dataSource;
+  __block id<ASRangeControllerDelegate> delegate = _delegate;
   dispatch_async(dispatch_get_main_queue(), ^{
-    [self performRangeUpdate];
+    [self _updateVisibleNodeIndexPaths];
+    
+    dataSource = nil;
+    delegate = nil;
   });
 }
 
@@ -320,6 +336,11 @@ static UIApplicationState __ApplicationState = UIApplicationStateActive;
   _rangeIsValid = YES;
   _queuedRangeUpdate = NO;
   
+  for (void (^completionBlock)(void) in _scheduledRangeUpdateCompletionBlocks) {
+    completionBlock();
+  }
+  [_scheduledRangeUpdateCompletionBlocks removeAllObjects];
+  
 #if ASRangeControllerLoggingEnabled
 //  NSSet *visibleNodePathsSet = [NSSet setWithArray:visibleNodePaths];
 //  BOOL setsAreEqual = [visibleIndexPaths isEqualToSet:visibleNodePathsSet];

AsyncDisplayKit/Details/ASRangeControllerUpdateRangeProtocol+Beta.h

@@ -49,6 +49,12 @@
  */
 - (void)updateCurrentRangeWithMode:(ASLayoutRangeMode)rangeMode;
 
+/**
+ * Schedule a range update and call the completion block if finished. This drives updating the working
+ * ranges, and triggering their actions.
+ */
+- (void)scheduleRangeUpdateCompletion:(void (^)(void))completion;
+
 @end
 
 
@@ -64,7 +70,10 @@
 
 @interface ASViewController (ASRangeControllerUpdateRangeProtocol)
 
-/// Automatically adjust range mode based on view events if the containing node confirms to the ASRangeControllerUpdateRangeProtocol
+/**
+ * Automatically adjust range mode based on view events if the containing node confirms to the
+ * ASRangeControllerUpdateRangeProtocol
+ */
 @property (nonatomic, assign) BOOL automaticallyAdjustRangeModeBasedOnViewEvents;
 
 @end