[Yoga] Rewrite YOGA_TREE_CONTIGUOUS mode with improved behavior and cleaner integration

Source/ASDisplayNode+Beta.h

@@ -178,6 +178,7 @@ extern void ASDisplayNodePerformBlockOnEveryYogaChild(ASDisplayNode * _Nullable
 - (void)semanticContentAttributeDidChange:(UISemanticContentAttribute)attribute;
 
 #if YOGA_TREE_CONTIGUOUS
+@property (nonatomic, assign) BOOL yogaLayoutInProgress;
 @property (nonatomic, strong, nullable) ASLayout *yogaCalculatedLayout;
 // These methods should not normally be called directly.
 - (void)invalidateCalculatedYogaLayout;
@@ -188,6 +189,9 @@ extern void ASDisplayNodePerformBlockOnEveryYogaChild(ASDisplayNode * _Nullable
 
 @interface ASLayoutElementStyle (Yoga)
 
+- (YGNodeRef)yogaNodeCreateIfNeeded;
+@property (nonatomic, assign, readonly) YGNodeRef yogaNode;
+
 @property (nonatomic, assign, readwrite) ASStackLayoutDirection flexDirection;
 @property (nonatomic, assign, readwrite) YGDirection direction;
 @property (nonatomic, assign, readwrite) CGFloat spacing;

Source/ASDisplayNode+Yoga.mm

@@ -23,7 +23,7 @@
 #import <AsyncDisplayKit/ASYogaUtilities.h>
 #import <AsyncDisplayKit/ASDisplayNode+Beta.h>
 #import <AsyncDisplayKit/ASDisplayNode+FrameworkPrivate.h>
-#import <AsyncDisplayKit/ASDisplayNode+Subclasses.h>
+#import <AsyncDisplayKit/ASDisplayNode+FrameworkSubclasses.h>
 #import <AsyncDisplayKit/ASDisplayNodeInternal.h>
 #import <AsyncDisplayKit/ASLayout.h>
 
@@ -35,7 +35,7 @@
 
 @interface ASDisplayNode (YogaInternal)
 @property (nonatomic, weak) ASDisplayNode *yogaParent;
-@property (nonatomic, assign) YGNodeRef yogaNode;
+- (ASSizeRange)_locked_constrainedSizeForLayoutPass;
 @end
 
 #endif /* YOGA_TREE_CONTIGUOUS */
@@ -76,7 +76,6 @@ - (void)addYogaChild:(ASDisplayNode *)child
 #if YOGA_TREE_CONTIGUOUS
   // YGNodeRef insertion is done in setParent:
   child.yogaParent = self;
-  self.hierarchyState |= ASHierarchyStateYogaLayoutEnabled;
 #else
   // When using non-contiguous Yoga layout, each level in the node hierarchy independently uses an ASYogaLayoutSpec
   __weak ASDisplayNode *weakSelf = self;
@@ -99,9 +98,6 @@ - (void)removeYogaChild:(ASDisplayNode *)child
 #if YOGA_TREE_CONTIGUOUS
   // YGNodeRef removal is done in setParent:
   child.yogaParent = nil;
-  if (_yogaChildren.count == 0 && self.yogaParent == nil) {
-    self.hierarchyState &= ~ASHierarchyStateYogaLayoutEnabled;
-  }
 #else
   if (_yogaChildren.count == 0) {
     self.layoutSpecBlock = nil;
@@ -121,38 +117,22 @@ - (void)semanticContentAttributeDidChange:(UISemanticContentAttribute)attribute
 
 #if YOGA_TREE_CONTIGUOUS /* YOGA_TREE_CONTIGUOUS */
 
-- (void)setYogaNode:(YGNodeRef)yogaNode
-{
-  _yogaNode = yogaNode;
-}
-
-- (YGNodeRef)yogaNode
-{
-  if (_yogaNode == NULL) {
-    _yogaNode = YGNodeNew();
-  }
-  return _yogaNode;
-}
-
 - (void)setYogaParent:(ASDisplayNode *)yogaParent
 {
   if (_yogaParent == yogaParent) {
     return;
   }
 
-  YGNodeRef yogaNode = self.yogaNode; // Use property to assign Ref if needed.
+  YGNodeRef yogaNode = [self.style yogaNodeCreateIfNeeded];
   YGNodeRef oldParentRef = YGNodeGetParent(yogaNode);
   if (oldParentRef != NULL) {
     YGNodeRemoveChild(oldParentRef, yogaNode);
   }
 
   _yogaParent = yogaParent;
   if (yogaParent) {
-    self.hierarchyState |= ASHierarchyStateYogaLayoutEnabled;
-    YGNodeRef newParentRef = yogaParent.yogaNode;
+    YGNodeRef newParentRef = [yogaParent.style yogaNodeCreateIfNeeded];
     YGNodeInsertChild(newParentRef, yogaNode, YGNodeGetChildCount(newParentRef));
-  } else {
-    self.hierarchyState &= ~ASHierarchyStateYogaLayoutEnabled;
   }
 }
 
@@ -171,9 +151,19 @@ - (ASLayout *)yogaCalculatedLayout
   return _yogaCalculatedLayout;
 }
 
+- (void)setYogaLayoutInProgress:(BOOL)yogaLayoutInProgress
+{
+  _yogaLayoutInProgress = yogaLayoutInProgress;
+}
+
+- (BOOL)yogaLayoutInProgress
+{
+  return _yogaLayoutInProgress;
+}
+
 - (ASLayout *)layoutForYogaNode
 {
-  YGNodeRef yogaNode = self.yogaNode;
+  YGNodeRef yogaNode = self.style.yogaNode;
 
   CGSize  size     = CGSizeMake(YGNodeLayoutGetWidth(yogaNode), YGNodeLayoutGetHeight(yogaNode));
   CGPoint position = CGPointMake(YGNodeLayoutGetLeft(yogaNode), YGNodeLayoutGetTop(yogaNode));
@@ -184,7 +174,7 @@ - (ASLayout *)layoutForYogaNode
 
 - (void)setupYogaCalculatedLayout
 {
-  YGNodeRef yogaNode = self.yogaNode; // Use property to assign Ref if needed.
+  YGNodeRef yogaNode = self.style.yogaNode;
   uint32_t childCount = YGNodeGetChildCount(yogaNode);
   ASDisplayNodeAssert(childCount == self.yogaChildren.count,
                       @"Yoga tree should always be in sync with .yogaNodes array! %@", self.yogaChildren);
@@ -197,6 +187,8 @@ - (void)setupYogaCalculatedLayout
   // The layout for self should have position CGPointNull, but include the calculated size.
   CGSize size = CGSizeMake(YGNodeLayoutGetWidth(yogaNode), YGNodeLayoutGetHeight(yogaNode));
   ASLayout *layout = [ASLayout layoutWithLayoutElement:self size:size sublayouts:sublayouts];
+
+  // TODO(appleguy): Set up _pendingDisplayNodeLayout, with a placeholder constrainedSize.
   self.yogaCalculatedLayout = layout;
 }
 
@@ -205,104 +197,59 @@ - (void)setYogaMeasureFuncIfNeeded
   // Size calculation via calculateSizeThatFits: or layoutSpecThatFits:
   // This will be used for ASTextNode, as well as any other node that has no Yoga children
   if (self.yogaChildren.count == 0) {
-    YGNodeRef yogaNode = self.yogaNode; // Use property to assign Ref if needed.
-    YGNodeSetContext(yogaNode, (__bridge void *)self);
-    YGNodeSetMeasureFunc(yogaNode, &ASLayoutElementYogaMeasureFunc);
+    // TODO(appleguy): Add override detection for calculateSizeThatFits: and calculateLayoutThatFits:,
+    // then we can set the MeasureFunc only for nodes that override one of the trio of measurement methods.
+    // if (_layoutSpecBlock == NULL && (_methodOverrides & ASDisplayNodeMethodOverrideLayoutSpecThatFits) == 0 && ...) {
+
+    YGNodeRef yogaNode = self.style.yogaNode;
+    ASDisplayNodeAssert(yogaNode, @"-[ASDisplayNode setYogaMeasureFuncIfNeeded] called with a NULL yogaNode! %@", self);
+    if (yogaNode) {
+      // TODO(appleguy): Use __bridge_transfer to retain nodes and clean up appropriately.
+      YGNodeSetContext(yogaNode, (__bridge void *)self);
+      YGNodeSetMeasureFunc(yogaNode, &ASLayoutElementYogaMeasureFunc);
+    }
   }
 }
 
 - (void)invalidateCalculatedYogaLayout
 {
   // Yoga internally asserts that this method may only be called on nodes with a measurement function.
-  YGNodeRef yogaNode = self.yogaNode;
-  if (YGNodeGetMeasureFunc(yogaNode)) {
+  YGNodeRef yogaNode = self.style.yogaNode;
+  if (yogaNode && YGNodeGetMeasureFunc(yogaNode)) {
     YGNodeMarkDirty(yogaNode);
   }
   self.yogaCalculatedLayout = nil;
 }
 
 - (void)calculateLayoutFromYogaRoot:(ASSizeRange)rootConstrainedSize
 {
-  if (self.yogaParent) {
-    if (self.yogaCalculatedLayout == nil) {
-      [self _setNeedsLayoutFromAbove];
-    }
-    return;
-  }
-  if (ASHierarchyStateIncludesYogaLayoutMeasuring(self.hierarchyState)) {
-    ASDisplayNodeAssert(NO, @"A Yoga layout is being performed by a parent; children must not perform their own until it is done! %@", [self displayNodeRecursiveDescription]);
+  ASDisplayNode *yogaParent = self.yogaParent;
+  self.yogaLayoutInProgress = YES;
+
+  if (yogaParent) {
+    // TODO(appleguy): Consider how to get the constrainedSize for the yogaRoot when escalating manually.
+    [yogaParent calculateLayoutFromYogaRoot:ASSizeRangeUnconstrained];
     return;
   }
 
-  ASDisplayNodePerformBlockOnEveryYogaChild(self, ^(ASDisplayNode * _Nonnull node) {
-    node.hierarchyState |= ASHierarchyStateYogaLayoutMeasuring;
-  });
+  ASDN::MutexLocker l(__instanceLock__);
 
-  YGNodeRef rootYogaNode = self.yogaNode;
+  if (ASSizeRangeEqualToSizeRange(rootConstrainedSize, ASSizeRangeUnconstrained)) {
+    rootConstrainedSize = [self _locked_constrainedSizeForLayoutPass];
+  }
+
+  YGNodeRef rootYogaNode = self.style.yogaNode;
 
   // Apply the constrainedSize as a base, known frame of reference.
   // If the root node also has style.*Size set, these will be overridden below.
   // YGNodeCalculateLayout currently doesn't offer the ability to pass a minimum size (max is passed there).
+
+  // TODO(appleguy): Reconcile the self.style.*Size properties with rootConstrainedSize
   YGNodeStyleSetMinWidth (rootYogaNode, yogaFloatForCGFloat(rootConstrainedSize.min.width));
   YGNodeStyleSetMinHeight(rootYogaNode, yogaFloatForCGFloat(rootConstrainedSize.min.height));
 
   ASDisplayNodePerformBlockOnEveryYogaChild(self, ^(ASDisplayNode * _Nonnull node) {
-    ASLayoutElementStyle *style = node.style;
-    YGNodeRef yogaNode = node.yogaNode;
-
-    YGNodeStyleSetDirection     (yogaNode, style.direction);
-
-    YGNodeStyleSetFlexWrap      (yogaNode, style.flexWrap);
-    YGNodeStyleSetFlexGrow      (yogaNode, style.flexGrow);
-    YGNodeStyleSetFlexShrink    (yogaNode, style.flexShrink);
-    YGNODE_STYLE_SET_DIMENSION  (yogaNode, FlexBasis, style.flexBasis);
-
-    YGNodeStyleSetFlexDirection (yogaNode, yogaFlexDirection(style.flexDirection));
-    YGNodeStyleSetJustifyContent(yogaNode, yogaJustifyContent(style.justifyContent));
-    YGNodeStyleSetAlignSelf     (yogaNode, yogaAlignSelf(style.alignSelf));
-    ASStackLayoutAlignItems alignItems = style.alignItems;
-    if (alignItems != ASStackLayoutAlignItemsNotSet) {
-      YGNodeStyleSetAlignItems(yogaNode, yogaAlignItems(alignItems));
-    }
-
-    YGNodeStyleSetPositionType  (yogaNode, style.positionType);
-    ASEdgeInsets position = style.position;
-    ASEdgeInsets margin   = style.margin;
-    ASEdgeInsets padding  = style.padding;
-    ASEdgeInsets border   = style.border;
-
-    YGEdge edge = YGEdgeLeft;
-    for (int i = 0; i < YGEdgeAll + 1; ++i) {
-      YGNODE_STYLE_SET_DIMENSION_WITH_EDGE(yogaNode, Position, dimensionForEdgeWithEdgeInsets(edge, position), edge);
-      YGNODE_STYLE_SET_DIMENSION_WITH_EDGE(yogaNode, Margin, dimensionForEdgeWithEdgeInsets(edge, margin), edge);
-      YGNODE_STYLE_SET_DIMENSION_WITH_EDGE(yogaNode, Padding, dimensionForEdgeWithEdgeInsets(edge, padding), edge);
-      YGNODE_STYLE_SET_FLOAT_WITH_EDGE(yogaNode, Border, dimensionForEdgeWithEdgeInsets(edge, border), edge);
-      edge = (YGEdge)(edge + 1);
-    }
-
-    CGFloat aspectRatio = style.aspectRatio;
-    if (aspectRatio > FLT_EPSILON && aspectRatio < CGFLOAT_MAX / 2.0) {
-      YGNodeStyleSetAspectRatio(yogaNode, aspectRatio);
-    }
-
-    // For the root node, we use rootConstrainedSize above. For children, consult the style for their size.
-    if (node != self) {
-      YGNODE_STYLE_SET_DIMENSION(yogaNode, Width, style.width);
-      YGNODE_STYLE_SET_DIMENSION(yogaNode, Height, style.height);
-
-      YGNODE_STYLE_SET_DIMENSION(yogaNode, MinWidth, style.minWidth);
-      YGNODE_STYLE_SET_DIMENSION(yogaNode, MinHeight, style.minHeight);
-
-      YGNODE_STYLE_SET_DIMENSION(yogaNode, MaxWidth, style.maxWidth);
-      YGNODE_STYLE_SET_DIMENSION(yogaNode, MaxHeight, style.maxHeight);
-    }
-
     [node setYogaMeasureFuncIfNeeded];
-
-    /* TODO(appleguy): STYLE SETTER METHODS LEFT TO IMPLEMENT
-     void YGNodeStyleSetOverflow(YGNodeRef node, YGOverflow overflow);
-     void YGNodeStyleSetFlex(YGNodeRef node, float flex);
-     */
   });
 
   // It is crucial to use yogaFloat... to convert CGFLOAT_MAX into YGUndefined here.
@@ -313,7 +260,7 @@ - (void)calculateLayoutFromYogaRoot:(ASSizeRange)rootConstrainedSize
 
   ASDisplayNodePerformBlockOnEveryYogaChild(self, ^(ASDisplayNode * _Nonnull node) {
     [node setupYogaCalculatedLayout];
-    node.hierarchyState &= ~ASHierarchyStateYogaLayoutMeasuring;
+    node.yogaLayoutInProgress = NO;
   });
 
 #if YOGA_LAYOUT_LOGGING /* YOGA_LAYOUT_LOGGING */

Source/ASDisplayNode.mm

@@ -270,7 +270,6 @@ - (void)_initializeInstance
 #if ASEVENTLOG_ENABLE
   _eventLog = [[ASEventLog alloc] initWithObject:self];
 #endif
-
   _viewClass = [self.class viewClass];
   _layerClass = [self.class layerClass];
   _contentsScaleForDisplay = ASScreenScale();
@@ -424,12 +423,6 @@ - (void)dealloc
     [self _scheduleIvarsForMainDeallocation];
   }
 
-#if YOGA_TREE_CONTIGUOUS
-  if (_yogaNode != NULL) {
-    YGNodeFree(_yogaNode);
-  }
-#endif
-
   // TODO: Remove this? If supernode isn't already nil, this method isn't dealloc-safe anyway.
   [self _setSupernode:nil];
 }
@@ -966,22 +959,26 @@ - (ASLayout *)calculateLayoutThatFits:(ASSizeRange)constrainedSize
   ASDN::MutexLocker l(__instanceLock__);
 
 #if YOGA_TREE_CONTIGUOUS /* YOGA */
-  if (ASHierarchyStateIncludesYogaLayoutEnabled(_hierarchyState) == YES) {
-    if (ASHierarchyStateIncludesYogaLayoutMeasuring(_hierarchyState) == NO && self.yogaCalculatedLayout == nil) {
-      ASDN::MutexUnlocker ul(__instanceLock__);
+  // There are several cases where Yoga could arrive here:
+  // - This node is not in a Yoga tree: it has neither a yogaParent nor yogaChildren.
+  // - This node is a Yoga tree root: it has no yogaParent, but has yogaChildren.
+  // - This node is a Yoga tree node: it has both a yogaParent and yogaChildren.
+  // - This node is a Yoga tree leaf: it has a yogaParent, but no yogaChidlren.
+  // If we're a leaf node, we are probably being called by a measure function and proceed as normal.
+  // If we're a yoga root or tree node, initiate a new Yoga calculation pass from root.
+  YGNodeRef yogaNode = _style.yogaNode;
+  if (yogaNode && (_yogaParent == nil || _yogaChildren.count > 0)) {
+    // This node has some connection to a Yoga tree.
+    ASDN::MutexUnlocker ul(__instanceLock__);
+    if (self.yogaLayoutInProgress == NO) {
       [self calculateLayoutFromYogaRoot:constrainedSize];
     }
-
-    // The call above may set yogaCalculatedLayout, even if it tested as nil to enter it.
-    if (self.yogaCalculatedLayout && self.yogaChildren.count > 0) {
-      return self.yogaCalculatedLayout;
-    }
+    return _yogaCalculatedLayout;
   }
 #endif /* YOGA */
 
   // Manual size calculation via calculateSizeThatFits:
-  if (((_methodOverrides & ASDisplayNodeMethodOverrideLayoutSpecThatFits) ||
-      (_layoutSpecBlock != NULL)) == NO) {
+  if (_layoutSpecBlock == NULL && (_methodOverrides & ASDisplayNodeMethodOverrideLayoutSpecThatFits) == 0) {
     CGSize size = [self calculateSizeThatFits:constrainedSize.max];
     ASDisplayNodeLogEvent(self, @"calculatedSize: %@", NSStringFromCGSize(size));
     return [ASLayout layoutWithLayoutElement:self size:ASSizeRangeClamp(constrainedSize, size) sublayouts:nil];
@@ -2894,7 +2891,9 @@ - (void)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 (CGRectEqualToRect(self.bounds, CGRectZero) == NO) {
+    [self layoutIfNeeded];
+  }
 
   if (_methodOverrides & ASDisplayNodeMethodOverrideFetchData) {
 #pragma clang diagnostic push

Source/Details/ASDataController.mm

@@ -391,7 +391,7 @@ - (void)_insertElementsIntoMap:(ASMutableElementMap *)map
       nodeBlock = [dataSource dataController:self supplementaryNodeBlockOfKind:kind atIndexPath:indexPath];
     }
 
-    ASSizeRange constrainedSize;
+    ASSizeRange constrainedSize = ASSizeRangeUnconstrained;
     if (shouldFetchSizeRanges) {
       constrainedSize = [self constrainedSizeForNodeOfKind:kind atIndexPath:indexPath];
     }