Relayout all nodes in backing store on relayoutAllNodes

AsyncDisplayKit/ASCollectionView.mm

@@ -596,7 +596,7 @@ static BOOL _isInterceptedSelector(SEL sel)
       _ignoreMaxSizeChange = NO;
     } else {
       [self performBatchAnimated:NO updates:^{
-        [_dataController relayoutAllRows];
+        [_dataController relayoutAllNodes];
       } completion:nil];
     }
   }

AsyncDisplayKit/ASTableView.mm

@@ -405,7 +405,7 @@ void ASPerformBlockWithoutAnimation(BOOL withoutAnimation, void (^block)()) {
       _ignoreMaxWidthChange = NO;
     } else {
       [self beginUpdates];
-      [_dataController relayoutAllRows];
+      [_dataController relayoutAllNodes];
       [self endUpdates];
     }
   }

AsyncDisplayKit/Details/ASDataController.h

@@ -157,10 +157,12 @@ typedef NSUInteger ASDataControllerAnimationOptions;
 - (void)reloadRowsAtIndexPaths:(NSArray *)indexPaths withAnimationOptions:(ASDataControllerAnimationOptions)animationOptions;
 
 /**
- * Re-measures all loaded nodes. Used to respond to a change in size of the containing view 
+ * Re-measures all loaded nodes in the backing store.
+ * 
+ * @discussion Used to respond to a change in size of the containing view
  * (e.g. ASTableView or ASCollectionView after an orientation change).
  */
-- (void)relayoutAllRows;
+- (void)relayoutAllNodes;
 
 - (void)moveRowAtIndexPath:(NSIndexPath *)indexPath toIndexPath:(NSIndexPath *)newIndexPath withAnimationOptions:(ASDataControllerAnimationOptions)animationOptions;
 

AsyncDisplayKit/Details/ASDataController.mm

@@ -729,41 +729,46 @@ static void *kASSizingQueueContext = &kASSizingQueueContext;
   }];
 }
 
-- (void)relayoutAllRows
+- (void)relayoutAllNodes
 {
   [self performEditCommandWithBlock:^{
     ASDisplayNodeAssertMainThread();
     LOG(@"Edit Command - relayoutRows");
     [_editingTransactionQueue waitUntilAllOperationsAreFinished];
-    
-    void (^relayoutNodesBlock)(NSMutableArray *) = ^void(NSMutableArray *nodes) {
-      if (!nodes.count) {
-        return;
-      }
-      
-      [self accessDataSourceWithBlock:^{
-        [nodes enumerateObjectsUsingBlock:^(NSMutableArray *section, NSUInteger sectionIndex, BOOL *stop) {
-          [section enumerateObjectsUsingBlock:^(ASCellNode *node, NSUInteger rowIndex, BOOL *stop) {
-            NSIndexPath *indexPath = [NSIndexPath indexPathForRow:rowIndex inSection:sectionIndex];
-            ASSizeRange constrainedSize = [_dataSource dataController:self constrainedSizeForNodeAtIndexPath:indexPath];
-            [node measureWithSizeRange:constrainedSize];
-            node.frame = CGRectMake(0.0f, 0.0f, node.calculatedSize.width, node.calculatedSize.height);
-          }];
-        }];
-      }];
-    };
 
     // Can't relayout right away because _completedNodes may not be up-to-date,
     // i.e there might be some nodes that were measured using the old constrained size but haven't been added to _completedNodes
     // (see _layoutNodes:atIndexPaths:withAnimationOptions:).
     [_editingTransactionQueue addOperationWithBlock:^{
       ASDisplayNodePerformBlockOnMainThread(^{
-        relayoutNodesBlock(_completedNodes[ASDataControllerRowNodeKind]);
+        [_completedNodes enumerateKeysAndObjectsUsingBlock:^(NSString *kind, NSMutableArray *nodes, BOOL *stop) {
+          [self _relayoutNodesOfKind:kind];
+        }];
       });
     }];
   }];
 }
 
+- (void)_relayoutNodesOfKind:(NSString *)kind
+{
+  ASDisplayNodeAssertMainThread();
+  NSArray *nodes = [self completedNodesOfKind:kind];
+  if (!nodes.count) {
+    return;
+  }
+  
+  [self accessDataSourceWithBlock:^{
+    [nodes enumerateObjectsUsingBlock:^(NSMutableArray *section, NSUInteger sectionIndex, BOOL *stop) {
+      [section enumerateObjectsUsingBlock:^(ASCellNode *node, NSUInteger rowIndex, BOOL *stop) {
+        NSIndexPath *indexPath = [NSIndexPath indexPathForRow:rowIndex inSection:sectionIndex];
+        ASSizeRange constrainedSize = [self constrainedSizeForNodeOfKind:kind atIndexPath:indexPath];
+        [node measureWithSizeRange:constrainedSize];
+        node.frame = CGRectMake(0.0f, 0.0f, node.calculatedSize.width, node.calculatedSize.height);
+      }];
+    }];
+  }];
+}
+
 - (void)moveRowAtIndexPath:(NSIndexPath *)indexPath toIndexPath:(NSIndexPath *)newIndexPath withAnimationOptions:(ASDataControllerAnimationOptions)animationOptions
 {
   [self performEditCommandWithBlock:^{