[ASDisplayNode] Add onDidLoad Method to Perform Work When Loaded (#2128)

* [ASDisplayNode] Add `onDidLoad:` method

* Prevent user from rasterizing wrapper nodes – they can't be reloaded in the future

AsyncDisplayKit/ASDisplayNode.h

@@ -40,7 +40,7 @@ typedef CALayer * _Nonnull(^ASDisplayNodeLayerBlock)();
 /**
  * ASDisplayNode loaded callback block. This block is called BEFORE the -didLoad method and is always called on the main thread.
  */
-typedef void (^ASDisplayNodeDidLoadBlock)(ASDisplayNode * _Nonnull node);
+typedef void (^ASDisplayNodeDidLoadBlock)(__kindof ASDisplayNode * _Nonnull node);
 
 /**
  * ASDisplayNode will / did render node content in context.
@@ -50,7 +50,7 @@ typedef void (^ASDisplayNodeContextModifier)(_Nonnull CGContextRef context);
 /**
  * ASDisplayNode layout spec block. This block can be used instead of implementing layoutSpecThatFits: in subclass
  */
-typedef ASLayoutSpec * _Nonnull(^ASLayoutSpecBlock)(ASDisplayNode * _Nonnull node, ASSizeRange constrainedSize);
+typedef ASLayoutSpec * _Nonnull(^ASLayoutSpecBlock)(__kindof ASDisplayNode * _Nonnull node, ASSizeRange constrainedSize);
 
 /**
  Interface state is available on ASDisplayNode and ASViewController, and
@@ -160,6 +160,18 @@ NS_ASSUME_NONNULL_BEGIN
  */
 - (instancetype)initWithLayerBlock:(ASDisplayNodeLayerBlock)layerBlock didLoadBlock:(nullable ASDisplayNodeDidLoadBlock)didLoadBlock;
 
+/**
+ * @abstract Add a block of work to be performed on the main thread when the node's view or layer is loaded. Thread safe.
+ * @warning Be careful not to retain self in `body`. Change the block parameter list to `^(MYCustomNode *self) {}` if you
+ *   want to shadow self (e.g. if calling this during `init`).
+ *
+ * @param body The work to be performed when the node is loaded.
+ *
+ * @precondition The node is not already loaded.
+ * @note This will only be called the next time the node is loaded. If the node is later added to a subtree of a node
+ *    that has `shouldRasterizeDescendants=YES`, and is unloaded, this block will not be called if it is loaded again.
+ */
+- (void)onDidLoad:(ASDisplayNodeDidLoadBlock)body;
 
 /** @name Properties */
 

AsyncDisplayKit/ASDisplayNode.mm

@@ -372,8 +372,10 @@ static ASDisplayNodeMethodOverrides GetASDisplayNodeMethodOverrides(Class c)
   
   [self _initializeInstance];
   _viewBlock = viewBlock;
-  _nodeLoadedBlock = didLoadBlock;
   _flags.synchronous = YES;
+  if (didLoadBlock != nil) {
+    _onDidLoadBlocks = [NSMutableArray arrayWithObject:didLoadBlock];
+  }
   
   return self;
 }
@@ -392,13 +394,30 @@ static ASDisplayNodeMethodOverrides GetASDisplayNodeMethodOverrides(Class c)
   
   [self _initializeInstance];
   _layerBlock = layerBlock;
-  _nodeLoadedBlock = didLoadBlock;
   _flags.synchronous = YES;
   _flags.layerBacked = YES;
+  if (didLoadBlock != nil) {
+    _onDidLoadBlocks = [NSMutableArray arrayWithObject:didLoadBlock];
+  }
   
   return self;
 }
 
+- (void)onDidLoad:(ASDisplayNodeDidLoadBlock)body
+{
+  ASDN::MutexLocker l(__instanceLock__);
+  if ([self _isNodeLoaded]) {
+    ASDisplayNodeFailAssert(@"Attempt to call %@ on node after it was loaded. Node: %@", NSStringFromSelector(_cmd), self);
+    return;
+  }
+  
+  if (_onDidLoadBlocks == nil) {
+    _onDidLoadBlocks = [NSMutableArray arrayWithObject:body];
+  } else {
+    [_onDidLoadBlocks addObject:body];
+  }
+}
+
 - (void)dealloc
 {
   ASDisplayNodeAssertMainThread();
@@ -437,6 +456,7 @@ static ASDisplayNodeMethodOverrides GetASDisplayNodeMethodOverrides(Class c)
 {
   ASDisplayNodeAssertThreadAffinity(self);
   ASDisplayNodeAssert([self isNodeLoaded], @"Implementation shouldn't call __unloadNode if not loaded: %@", self);
+  ASDisplayNodeAssert(_flags.synchronous == NO, @"Node created using -initWithViewBlock:/-initWithLayerBlock: cannot be unloaded. Node: %@", self);
   ASDN::MutexLocker l(__instanceLock__);
 
   if (_flags.layerBacked)
@@ -596,13 +616,18 @@ static ASDisplayNodeMethodOverrides GetASDisplayNodeMethodOverrides(Class c)
   if (ASDisplayNodeThreadIsMain()) {
     // Because the view and layer can only be created and destroyed on Main, that is also the only thread
     // where the state of this property can change. As an optimization, we can avoid locking.
-    return (_view != nil || (_layer != nil && _flags.layerBacked));
+    return [self _isNodeLoaded];
   } else {
     ASDN::MutexLocker l(__instanceLock__);
-    return (_view != nil || (_layer != nil && _flags.layerBacked));
+    return [self _isNodeLoaded];
   }
 }
 
+- (BOOL)_isNodeLoaded
+{
+  return (_view != nil || (_layer != nil && _flags.layerBacked));
+}
+
 - (NSString *)name
 {
   ASDN::MutexLocker l(__instanceLock__);
@@ -2481,10 +2506,11 @@ void recursivelyTriggerDisplayForLayer(CALayer *layer, BOOL shouldBlock)
 - (void)__didLoad
 {
   ASDN::MutexLocker l(__instanceLock__);
-  if (_nodeLoadedBlock) {
-    _nodeLoadedBlock(self);
-    _nodeLoadedBlock = nil;
+  
+  for (ASDisplayNodeDidLoadBlock block in _onDidLoadBlocks) {
+    block(self);
   }
+  _onDidLoadBlocks = nil;
   [self didLoad];
 }
 
@@ -2815,6 +2841,11 @@ void recursivelyTriggerDisplayForLayer(CALayer *layer, BOOL shouldBlock)
     _hierarchyState = newState;
   }
   
+  // Entered rasterization state.
+  if (newState & ASHierarchyStateRasterized) {
+    ASDisplayNodeAssert(_flags.synchronous == NO, @"Node created using -initWithViewBlock:/-initWithLayerBlock: cannot be added to subtree of node with shouldRasterizeDescendants=YES. Node: %@", self);
+  }
+  
   // Entered or exited contents rendering state.
   if ((newState & ASHierarchyStateRangeManaged) != (oldState & ASHierarchyStateRangeManaged)) {
     if (newState & ASHierarchyStateRangeManaged) {

AsyncDisplayKit/Private/ASDisplayNodeInternal.h

@@ -133,7 +133,7 @@ FOUNDATION_EXPORT NSString * const ASRenderingEngineDidDisplayNodesScheduledBefo
   
   ASDisplayNodeViewBlock _viewBlock;
   ASDisplayNodeLayerBlock _layerBlock;
-  ASDisplayNodeDidLoadBlock _nodeLoadedBlock;
+  NSMutableArray<ASDisplayNodeDidLoadBlock> *_onDidLoadBlocks;
   Class _viewClass;
   Class _layerClass;
   

AsyncDisplayKitTests/ASDisplayNodeTests.m

@@ -2004,4 +2004,29 @@ static bool stringContainsPointer(NSString *description, id p) {
   XCTAssertNoThrow([node.view layoutIfNeeded]);
 }
 
+- (void)testThatOnDidLoadThrowsIfCalledOnLoaded
+{
+  ASTestDisplayNode *node = [[[ASTestDisplayNode alloc] init] autorelease];
+  [node view];
+  XCTAssertThrows([node onDidLoad:^(ASDisplayNode * _Nonnull node) { }]);
+}
+
+- (void)testThatOnDidLoadWorks
+{
+  ASTestDisplayNode *node = [[[ASTestDisplayNode alloc] init] autorelease];
+  NSMutableArray *calls = [NSMutableArray array];
+  [node onDidLoad:^(ASTestDisplayNode * _Nonnull node) {
+    [calls addObject:@0];
+  }];
+  [node onDidLoad:^(ASTestDisplayNode * _Nonnull node) {
+    [calls addObject:@1];
+  }];
+  [node onDidLoad:^(ASTestDisplayNode * _Nonnull node) {
+    [calls addObject:@2];
+  }];
+  [node view];
+  NSArray *expected = @[ @0, @1, @2 ];
+  XCTAssertEqualObjects(calls, expected);
+}
+
 @end