diff --git a/AsyncDisplayKit/ASCellNode.m b/AsyncDisplayKit/ASCellNode.m index cfb329601f..5083edf12c 100644 --- a/AsyncDisplayKit/ASCellNode.m +++ b/AsyncDisplayKit/ASCellNode.m @@ -13,6 +13,7 @@ #import #import +#import #import #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];