[ASCellNode] If viewControllerBlock returns an ASViewController, use its node directly instead of wrapping

AsyncDisplayKit/ASCellNode.m

@@ -13,6 +13,7 @@
 #import <AsyncDisplayKit/ASDisplayNode+Subclasses.h>
 #import <AsyncDisplayKit/ASTextNode.h>
 
+#import <AsyncDisplayKit/ASViewController.h>
 #import <AsyncDisplayKit/ASInsetLayoutSpec.h>
 
 #pragma mark -
@@ -20,6 +21,8 @@
 
 @interface ASCellNode ()
 {
+  ASDisplayNodeViewControllerBlock _viewControllerBlock;
+  ASDisplayNodeDidLoadBlock _viewControllerDidLoadBlock;
   ASDisplayNode *_viewControllerNode;
 }
 
@@ -45,20 +48,41 @@
     return nil;
   
   ASDisplayNodeAssertNotNil(viewControllerBlock, @"should initialize with a valid block that returns a UIViewController");
-  
-  if (viewControllerBlock) {
-    
-    _viewControllerNode = [[ASDisplayNode alloc] initWithViewBlock:^UIView *{
-      _viewController = viewControllerBlock();
-      return _viewController.view;
-    } didLoadBlock:didLoadBlock];
-    
-    [self addSubnode:_viewControllerNode];
-  }
+  _viewControllerBlock = viewControllerBlock;
+  _viewControllerDidLoadBlock = didLoadBlock;
 
   return self;
 }
 
+- (void)didLoad
+{
+  [super didLoad];
+
+  if (_viewControllerBlock != nil) {
+
+    UIViewController *viewController = _viewControllerBlock();
+    _viewControllerBlock = nil;
+
+    if ([viewController isKindOfClass:[ASViewController class]]) {
+      ASViewController *asViewController = (ASViewController *)viewController;
+      _viewControllerNode = asViewController.node;
+    } else {
+      _viewControllerNode = [[ASDisplayNode alloc] initWithViewBlock:^{
+        return viewController.view;
+      }];
+    }
+    [self addSubnode:_viewControllerNode];
+
+    // Since we just loaded our node, and added _viewControllerNode as a subnode,
+    // _viewControllerNode must have just loaded its view, so now is an appropriate
+    // time to execute our didLoadBlock, if we were given one.
+    if (_viewControllerDidLoadBlock != nil) {
+      _viewControllerDidLoadBlock(self);
+      _viewControllerDidLoadBlock = nil;
+    }
+  }
+}
+
 - (void)layout
 {
   [super layout];