[ASDisplayNode] Trigger a layout pass whenever a node enters preload state (#3263)

* Add a thread-safe layoutIfNeeded implementation to ASDisplayNode

* Trigger a layout pass when a display node enters preload state
- This ensures that all the subnodes have the correct size to preload their content.

* ASCollectionNode to trigger its initial data load when it enters preload state

* Minor change in _ASCollectionViewCell

* Layout sublayouts before dispatch to main for subclass hooks

* Update comments

* Don't wait until updates are committed when the collection node enters display state

* Same deal for table node

* Explain the layout trigger in ASDisplayNode

Source/ASCollectionNode.mm

@@ -193,18 +193,20 @@
   [self.rangeController clearContents];
 }
 
-- (void)didExitPreloadState
-{
-  [super didExitPreloadState];
-  [self.rangeController clearPreloadedData];
-}
-
 - (void)interfaceStateDidChange:(ASInterfaceState)newState fromState:(ASInterfaceState)oldState
 {
   [super interfaceStateDidChange:newState fromState:oldState];
   [ASRangeController layoutDebugOverlayIfNeeded];
 }
 
+- (void)didEnterPreloadState
+{
+  // Intentionally allocate the view here so that super will trigger a layout pass on it which in turn will trigger the intial data load.
+  // We can get rid of this call later when ASDataController, ASRangeController and ASCollectionLayout can operate without the view.
+  [self view];
+  [super didEnterPreloadState];
+}
+
 #if ASRangeControllerLoggingEnabled
 - (void)didEnterVisibleState
 {
@@ -219,6 +221,12 @@
 }
 #endif
 
+- (void)didExitPreloadState
+{
+  [super didExitPreloadState];
+  [self.rangeController clearPreloadedData];
+}
+
 #pragma mark Setter / Getter
 
 // TODO: Implement this without the view. Then revisit ASLayoutElementCollectionTableSetTraitCollection

Source/ASDisplayNode.h

@@ -639,6 +639,11 @@ extern NSInteger const ASDefaultDrawingPriority;
  */
 - (void)setNeedsLayout;
 
+/**
+ * Performs a layout pass on the node. Convenience for use whether the view / layer is loaded or not. Safe to call from a background thread.
+ */
+- (void)layoutIfNeeded;
+
 @property (nonatomic, strong, nullable) id contents;                           // default=nil
 @property (nonatomic, assign)           BOOL clipsToBounds;                    // default==NO
 @property (nonatomic, getter=isOpaque)  BOOL opaque;                           // default==YES

Source/ASDisplayNode.mm

@@ -985,7 +985,7 @@ ASLayoutElementFinalLayoutElementDefault
 
 - (void)__layout
 {
-  ASDisplayNodeAssertMainThread();
+  ASDisplayNodeAssertThreadAffinity(self);
   ASDisplayNodeAssertLockUnownedByCurrentThread(__instanceLock__);
   
   {
@@ -1014,8 +1014,12 @@ ASLayoutElementFinalLayoutElementDefault
     [self _locked_layoutPlaceholderIfNecessary];
   }
   
+  [self _layoutSublayouts];
+  
+  ASPerformBlockOnMainThread(^{
     [self layout];
     [self layoutDidFinish];
+  });
 }
 
 /// Needs to be called with lock held
@@ -1054,7 +1058,7 @@ ASLayoutElementFinalLayoutElementDefault
   std::shared_ptr<ASDisplayNodeLayout> nextLayout = _pendingDisplayNodeLayout;
   #define layoutSizeDifferentFromBounds !CGSizeEqualToSize(nextLayout->layout.size, boundsSizeForLayout)
   
-  // nextLayout was likely created by a call to layoutThatFits:, check if is valid and can be applied.
+  // nextLayout was likely created by a call to layoutThatFits:, check if it is valid and can be applied.
   // If our bounds size is different than it, or invalid, recalculate.  Use #define to avoid nullptr->
   if (nextLayout == nullptr || nextLayout->isDirty() == YES || layoutSizeDifferentFromBounds) {
     // Use the last known constrainedSize passed from a parent during layout (if never, use bounds).
@@ -1405,12 +1409,13 @@ ASLayoutElementFinalLayoutElementDefault
 
 - (void)layout
 {
-  [self _layoutSublayouts];
+  ASDisplayNodeAssertMainThread();
+  // Subclass hook
 }
 
 - (void)_layoutSublayouts
 {
-  ASDisplayNodeAssertMainThread();
+  ASDisplayNodeAssertThreadAffinity(self);
   ASDisplayNodeAssertLockUnownedByCurrentThread(__instanceLock__);
   
   ASLayout *layout;
@@ -3720,6 +3725,10 @@ ASDISPLAYNODE_INLINE BOOL nodeIsInRasterizedTree(ASDisplayNode *node) {
   ASDisplayNodeAssertLockUnownedByCurrentThread(__instanceLock__);
   [_interfaceStateDelegate didEnterPreloadState];
   
+  // Trigger a layout pass to ensure all subnodes have the correct size to preload their content.
+  // This is important for image nodes, as well as collection and table nodes.
+  [self layoutIfNeeded];
+  
   if (_methodOverrides & ASDisplayNodeMethodOverrideFetchData) {
 #pragma clang diagnostic push
 #pragma clang diagnostic ignored "-Wdeprecated-declarations"

Source/ASTableNode.mm

@@ -122,18 +122,20 @@
   [self.rangeController clearContents];
 }
 
-- (void)didExitPreloadState
-{
-  [super didExitPreloadState];
-  [self.rangeController clearPreloadedData];
-}
-
 - (void)interfaceStateDidChange:(ASInterfaceState)newState fromState:(ASInterfaceState)oldState
 {
   [super interfaceStateDidChange:newState fromState:oldState];
   [ASRangeController layoutDebugOverlayIfNeeded];
 }
 
+- (void)didEnterPreloadState
+{
+  // Intentionally allocate the view here so that super will trigger a layout pass on it which in turn will trigger the intial data load.
+  // We can get rid of this call later when ASDataController, ASRangeController and ASCollectionLayout can operate without the view.
+  [self view];
+  [super didEnterPreloadState];
+}
+
 #if ASRangeControllerLoggingEnabled
 - (void)didEnterVisibleState
 {
@@ -148,6 +150,12 @@
 }
 #endif
 
+- (void)didExitPreloadState
+{
+  [super didExitPreloadState];
+  [self.rangeController clearPreloadedData];
+}
+
 #pragma mark Setter / Getter
 
 // TODO: Implement this without the view. Then revisit ASLayoutElementCollectionTableSetTraitCollection

Source/Details/UIView+ASConvenience.h

@@ -42,6 +42,7 @@ NS_ASSUME_NONNULL_BEGIN
 
 - (void)setNeedsDisplay;
 - (void)setNeedsLayout;
+- (void)layoutIfNeeded;
 
 @end
 

Source/Details/_ASCollectionViewCell.m

@@ -58,6 +58,7 @@
  */
 - (void)applyLayoutAttributes:(UICollectionViewLayoutAttributes *)layoutAttributes
 {
+  [super applyLayoutAttributes:layoutAttributes];
   self.layoutAttributes = layoutAttributes;
 }
 

Source/Private/ASDisplayNode+UIViewBridge.mm

@@ -40,7 +40,6 @@
 #if DISPLAYNODE_USE_LOCKS
 #define _bridge_prologue_read ASDN::MutexLocker l(__instanceLock__); ASDisplayNodeAssertThreadAffinity(self)
 #define _bridge_prologue_write ASDN::MutexLocker l(__instanceLock__)
-#define _bridge_prologue_write_unlock ASDN::MutexUnlocker u(__instanceLock__)
 #else
 #define _bridge_prologue_read ASDisplayNodeAssertThreadAffinity(self)
 #define _bridge_prologue_write
@@ -79,8 +78,6 @@ if (shouldApply) { _view.viewAndPendingViewStateProperty = (viewAndPendingViewSt
 #define _setToLayer(layerProperty, layerValueExpr) BOOL shouldApply = ASDisplayNodeShouldApplyBridgedWriteToView(self); \
 if (shouldApply) { _layer.layerProperty = (layerValueExpr); } else { ASDisplayNodeGetPendingState(self).layerProperty = (layerValueExpr); }
 
-#define _messageToViewOrLayer(viewAndLayerSelector) (_view ? [_view viewAndLayerSelector] : [_layer viewAndLayerSelector])
-
 /**
  * This category implements certain frequently-used properties and methods of UIView and CALayer so that ASDisplayNode clients can just call the view/layer methods on the node,
  * with minimal loss in performance.  Unlike UIView and CALayer methods, these can be called from a non-main thread until the view or layer is created.
@@ -300,9 +297,18 @@ if (shouldApply) { _layer.layerProperty = (layerValueExpr); } else { ASDisplayNo
 }
 
 - (void)setNeedsDisplay
+{
+  BOOL isRasterized = NO;
+  BOOL shouldApply = NO;
+  id viewOrLayer = nil;
   {
     _bridge_prologue_write;
-  if (_hierarchyState & ASHierarchyStateRasterized) {
+    isRasterized = _hierarchyState & ASHierarchyStateRasterized;
+    shouldApply = ASDisplayNodeShouldApplyBridgedWriteToView(self);
+    viewOrLayer = _view ?: _layer;
+  }
+  
+  if (isRasterized) {
     ASPerformBlockOnMainThread(^{
       // The below operation must be performed on the main thread to ensure against an extremely rare deadlock, where a parent node
       // begins materializing the view / layer hierarchy (locking itself or a descendant) while this node walks up
@@ -319,13 +325,13 @@ if (shouldApply) { _layer.layerProperty = (layerValueExpr); } else { ASDisplayNo
       [rasterizedContainerNode setNeedsDisplay];
     });
   } else {
-    BOOL shouldApply = ASDisplayNodeShouldApplyBridgedWriteToView(self);
     if (shouldApply) {
       // If not rasterized, and the node is loaded (meaning we certainly have a view or layer), send a
       // message to the view/layer first. This is because __setNeedsDisplay calls as scheduleNodeForDisplay,
       // which may call -displayIfNeeded. We want to ensure the needsDisplay flag is set now, and then cleared.
-      _messageToViewOrLayer(setNeedsDisplay);
+      [viewOrLayer setNeedsDisplay];
     } else {
+      _bridge_prologue_write;
       [ASDisplayNodeGetPendingState(self) setNeedsDisplay];
     }
     [self __setNeedsDisplay];
@@ -333,30 +339,60 @@ if (shouldApply) { _layer.layerProperty = (layerValueExpr); } else { ASDisplayNo
 }
 
 - (void)setNeedsLayout
+{
+  BOOL shouldApply = NO;
+  BOOL loaded = NO;
+  id viewOrLayer = nil;
   {
     _bridge_prologue_write;
-  BOOL shouldApply = ASDisplayNodeShouldApplyBridgedWriteToView(self);
+    shouldApply = ASDisplayNodeShouldApplyBridgedWriteToView(self);
+    loaded = __loaded(self);
+    viewOrLayer = _view ?: _layer;
+  }
+  
   if (shouldApply) {
     // The node is loaded and we're on main.
     // Quite the opposite of setNeedsDisplay, we must call __setNeedsLayout before messaging
     // the view or layer to ensure that measurement and implicitly added subnodes have been handled.
-    
-    // Calling __setNeedsLayout while holding the property lock can cause deadlocks
-    _bridge_prologue_write_unlock;
     [self __setNeedsLayout];
-    _bridge_prologue_write;
-    _messageToViewOrLayer(setNeedsLayout);
-  } else if (__loaded(self)) {
+    [viewOrLayer setNeedsLayout];
+  } else if (loaded) {
     // The node is loaded but we're not on main.
-    // We will call [self __setNeedsLayout] when we apply
-    // the pending state. We need to call it on main if the node is loaded
-    // to support automatic subnode management.
+    // We will call [self __setNeedsLayout] when we apply the pending state.
+    // We need to call it on main if the node is loaded to support automatic subnode management.
+    _bridge_prologue_write;
     [ASDisplayNodeGetPendingState(self) setNeedsLayout];
   } else {
     // The node is not loaded and we're not on main.
-    _bridge_prologue_write_unlock;
     [self __setNeedsLayout];
+  }
+}
+
+- (void)layoutIfNeeded
+{
+  BOOL shouldApply = NO;
+  BOOL loaded = NO;
+  id viewOrLayer = nil;
+  {
     _bridge_prologue_write;
+    shouldApply = ASDisplayNodeShouldApplyBridgedWriteToView(self);
+    loaded = __loaded(self);
+    viewOrLayer = _view ?: _layer;
+  }
+  
+  if (shouldApply) {
+    // The node is loaded and we're on main.
+    // Message the view or layer which in turn will call __layout on us (see -[_ASDisplayLayer layoutSublayers]).
+    [viewOrLayer layoutIfNeeded];
+  } else if (loaded) {
+    // The node is loaded but we're not on main.
+    // We will call layoutIfNeeded on the view or layer when we apply the pending state. __layout will in turn be called on us (see -[_ASDisplayLayer layoutSublayers]).
+    // We need to call it on main if the node is loaded to support automatic subnode management.
+    _bridge_prologue_write;
+    [ASDisplayNodeGetPendingState(self) layoutIfNeeded];
+  } else {
+    // The node is not loaded and we're not on main.
+    [self __layout];
   }
 }
 

Source/Private/_ASPendingState.mm

@@ -24,6 +24,7 @@ typedef struct {
   // Properties
   int needsDisplay:1;
   int needsLayout:1;
+  int layoutIfNeeded:1;
   
   // Flags indicating that a given property should be applied to the view at creation
   int setClipsToBounds:1;
@@ -272,6 +273,11 @@ static BOOL defaultAllowsEdgeAntialiasing = NO;
   _flags.needsLayout = YES;
 }
 
+- (void)layoutIfNeeded
+{
+  _flags.layoutIfNeeded = YES;
+}
+
 - (void)setClipsToBounds:(BOOL)flag
 {
   clipsToBounds = flag;
@@ -761,9 +767,6 @@ static BOOL defaultAllowsEdgeAntialiasing = NO;
   if (flags.setEdgeAntialiasingMask)
     layer.edgeAntialiasingMask = edgeAntialiasingMask;
 
-  if (flags.needsLayout)
-    [layer setNeedsLayout];
-
   if (flags.setAsyncTransactionContainer)
     layer.asyncdisplaykit_asyncTransactionContainer = asyncTransactionContainer;
 
@@ -771,6 +774,12 @@ static BOOL defaultAllowsEdgeAntialiasing = NO;
     ASDisplayNodeAssert(layer.opaque == opaque, @"Didn't set opaque as desired");
 
   ASPendingStateApplyMetricsToLayer(self, layer);
+  
+  if (flags.needsLayout)
+    [layer setNeedsLayout];
+  
+  if (flags.layoutIfNeeded)
+    [layer layoutIfNeeded];
 }
 
 - (void)applyToView:(UIView *)view withSpecialPropertiesHandling:(BOOL)specialPropertiesHandling
@@ -889,9 +898,6 @@ static BOOL defaultAllowsEdgeAntialiasing = NO;
   if (flags.setEdgeAntialiasingMask)
     layer.edgeAntialiasingMask = edgeAntialiasingMask;
 
-  if (flags.needsLayout)
-    [view setNeedsLayout];
-
   if (flags.setAsyncTransactionContainer)
     view.asyncdisplaykit_asyncTransactionContainer = asyncTransactionContainer;
 
@@ -955,6 +961,12 @@ static BOOL defaultAllowsEdgeAntialiasing = NO;
   } else {
     ASPendingStateApplyMetricsToLayer(self, layer);
   }
+  
+  if (flags.needsLayout)
+    [view setNeedsLayout];
+  
+  if (flags.layoutIfNeeded)
+    [view layoutIfNeeded];
 }
 
 // FIXME: Make this more efficient by tracking which properties are set rather than reading everything.