Batch render supplementary views on reload data

AsyncDisplayKit/ASCollectionView.mm

@@ -153,6 +153,8 @@ static BOOL _isInterceptedSelector(SEL sel)
   CGSize _maxSizeForNodesConstrainedSize;
   BOOL _ignoreMaxSizeChange;
   
+  NSMutableArray *_registeredSupplementaryKinds;
+  
   /**
    * If YES, the `UICollectionView` will reload its data on next layout pass so we should not forward any updates to it.
    
@@ -233,6 +235,8 @@ static BOOL _isInterceptedSelector(SEL sel)
     _layoutDelegate = _flowLayoutInspector;
   }
 
+  _registeredSupplementaryKinds = [NSMutableArray array];
+  
   self.backgroundColor = [UIColor whiteColor];
   
   [self registerClass:[_ASCollectionViewCell class] forCellWithReuseIdentifier:@"_ASCollectionViewCell"];
@@ -390,6 +394,7 @@ static BOOL _isInterceptedSelector(SEL sel)
 
 - (void)registerSupplementaryNodeOfKind:(NSString *)elementKind
 {
+  [_registeredSupplementaryKinds addObject:elementKind];
   [self registerClass:[UICollectionReusableView class] forSupplementaryViewOfKind:elementKind
                                             withReuseIdentifier:[self __reuseIdentifierForKind:elementKind]];
 }
@@ -637,11 +642,6 @@ static BOOL _isInterceptedSelector(SEL sel)
   return node;
 }
 
-- (ASDisplayNode *)dataController:(ASDataController *)dataController supplementaryNodeOfKind:(NSString *)kind atIndexPath:(NSIndexPath *)indexPath
-{
-  return [_asyncDataSource collectionView:self nodeForSupplementaryElementOfKind:kind atIndexPath:indexPath];
-}
-
 - (ASSizeRange)dataController:(ASDataController *)dataController constrainedSizeForNodeAtIndexPath:(NSIndexPath *)indexPath
 {
   ASSizeRange constrainedSize;
@@ -718,6 +718,16 @@ static BOOL _isInterceptedSelector(SEL sel)
 
 #pragma mark - ASCollectionViewDataControllerSource Supplementary view support
 
+- (ASDisplayNode *)dataController:(ASCollectionDataController *)dataController supplementaryNodeOfKind:(NSString *)kind atIndexPath:(NSIndexPath *)indexPath
+{
+  return [_asyncDataSource collectionView:self nodeForSupplementaryElementOfKind:kind atIndexPath:indexPath];
+}
+
+- (NSArray *)supplementaryNodeKindsInDataController:(ASCollectionDataController *)dataController
+{
+  return _registeredSupplementaryKinds;
+}
+
 - (ASSizeRange)dataController:(ASCollectionDataController *)dataController constrainedSizeForSupplementaryNodeOfKind:(NSString *)kind atIndexPath:(NSIndexPath *)indexPath
 {
   return [_layoutDelegate collectionView:self constrainedSizeForSupplementaryNodeOfKind:kind atIndexPath:indexPath];

AsyncDisplayKit/Details/ASCollectionDataController.mm

@@ -23,29 +23,32 @@
 @implementation ASCollectionDataController {
   NSMutableDictionary *_completedSupplementaryNodes;
   NSMutableDictionary *_editingSupplementaryNodes;
+  
+  NSMutableDictionary *_pendingNodes;
+  NSMutableDictionary *_pendingIndexPaths;
 }
 
-- (void)initialSupplementaryLoading
+- (void)prepareForReloadData
 {
-  [self performEditCommandWithBlock:^{
+  NSArray *elementKinds = [self.collectionDataSource supplementaryNodeKindsInDataController:self];
-    ASDisplayNodeAssertMainThread();
+  [elementKinds enumerateObjectsUsingBlock:^(NSString *kind, NSUInteger idx, BOOL * _Nonnull stop) {
-    [self accessDataSourceWithBlock:^{
+    NSMutableArray *indexPaths = [NSMutableArray array];
-      NSArray *elementKinds = [self.collectionDataSource supplementaryNodeKindsInDataController:self];
+    NSMutableArray *nodes = [NSMutableArray array];
-      [elementKinds enumerateObjectsUsingBlock:^(NSString *kind, NSUInteger idx, BOOL * _Nonnull stop) {
+    [self _populateSupplementaryNodesOfKind:kind withMutableNodes:nodes mutableIndexPaths:indexPaths];
-        _completedSupplementaryNodes[kind] = [NSMutableArray array];
+    _pendingNodes[kind] = nodes;
-        _editingSupplementaryNodes[kind] = [NSMutableArray array];
+    _pendingIndexPaths[kind] = indexPaths;
-        
-        NSMutableArray *indexPaths = [NSMutableArray array];
-        NSMutableArray *nodes = [NSMutableArray array];
-        [self _populateSupplementaryNodesOfKind:kind withMutableNodes:nodes mutableIndexPaths:indexPaths];
-        [self batchLayoutNodes:nodes atIndexPaths:indexPaths constrainedSize:^ASSizeRange(NSIndexPath *indexPath) {
-          return [self.collectionDataSource dataController:self constrainedSizeForSupplementaryNodeOfKind:kind atIndexPath:indexPath];
-        } completion:nil];
-      }];
-    }];
   }];
 }
-       
+
+- (void)willReloadData
+{
+  [_pendingNodes enumerateKeysAndObjectsUsingBlock:^(NSString *kind, NSMutableArray *nodes, BOOL *stop) {
+    [self batchLayoutNodes:nodes atIndexPaths:_pendingIndexPaths[kind] constrainedSize:^ASSizeRange(NSIndexPath *indexPath) {
+      return [self.collectionDataSource dataController:self constrainedSizeForSupplementaryNodeOfKind:kind atIndexPath:indexPath];
+    } completion:nil];
+  }];
+}
+
 - (void)_populateSupplementaryNodesOfKind:(NSString *)kind withMutableNodes:(NSMutableArray *)nodes mutableIndexPaths:(NSMutableArray *)indexPaths
 {
   NSUInteger sectionCount = [self.collectionDataSource dataController:self numberOfSectionsForSupplementaryKind:kind];

AsyncDisplayKit/Details/ASDataController+Subclasses.h

@@ -30,4 +30,14 @@
  */
 - (void)batchLayoutNodes:(NSArray *)nodes atIndexPaths:(NSArray *)indexPaths constrainedSize:(ASSizeRange (^)(NSIndexPath *indexPath))constraintedSizeBlock completion:(void (^)(NSArray *nodes, NSArray *indexPaths))completionBlock;
 
+/**
+ * An opportunity for a subclass to access the data source before entering into the editing queue
+ */
+- (void)prepareForReloadData;
+
+/**
+ * Subclasses can override this to reload data after the abstract data controller deletes its old data and before it reloads the new.
+ */
+- (void)willReloadData;
+
 @end

AsyncDisplayKit/Details/ASDataController.mm

@@ -308,6 +308,9 @@ static void *kASSizingQueueContext = &kASSizingQueueContext;
       
       // Measure nodes whose views are loaded before we leave the main thread
       [self _layoutNodesWithMainThreadAffinity:updatedNodes atIndexPaths:updatedIndexPaths];
+
+      // Allow subclasses to perform setup before going into the edit transaction
+      [self prepareForReloadData];
       
       [_editingTransactionQueue addOperationWithBlock:^{
         ASLOG(@"Edit Transaction - reloadData");
@@ -319,6 +322,8 @@ static void *kASSizingQueueContext = &kASSizingQueueContext;
         NSMutableIndexSet *indexSet = [[NSMutableIndexSet alloc] initWithIndexesInRange:NSMakeRange(0, _editingNodes.count)];
         [self _deleteSectionsAtIndexSet:indexSet withAnimationOptions:animationOptions];
         
+        [self willReloadData];
+        
         // Insert each section
         NSMutableArray *sections = [NSMutableArray arrayWithCapacity:sectionCount];
         for (int i = 0; i < sectionCount; i++) {
@@ -337,6 +342,16 @@ static void *kASSizingQueueContext = &kASSizingQueueContext;
   }];
 }
 
+- (void)prepareForReloadData
+{
+  // Implemented by subclasses
+}
+
+- (void)willReloadData
+{
+  // Implemented by subclasses
+}
+
 #pragma mark - Data Source Access (Calling _dataSource)
 
 - (void)accessDataSourceWithBlock:(dispatch_block_t)block