From f8474c0e2f4bf50e5398ba68b7686ce0e367c16b Mon Sep 17 00:00:00 2001
From: appleguy <scottgoodson@gmail.com>
Date: Sun, 2 Jul 2017 13:01:07 -0700
Subject: [PATCH] [Yoga] Refine the handling of measurement functions when Yoga
 is used. (#408)

This ensures that measure funcs are not set for container / empty spacer
nodes, because Yoga has more internal capabilities than layoutThatFits:
knows about.

Avoid need for the main layout pass to directly call setup for measure
funcs, by making it an implicit part of .yogaLayoutInProgress =.

Tear down the measure func after the layout pass to avoid retain cycle.
---
 Source/ASDisplayNode+Layout.mm |  2 +-
 Source/ASDisplayNode+Yoga.mm   | 25 ++++++++++---------------
 2 files changed, 11 insertions(+), 16 deletions(-)

diff --git a/Source/ASDisplayNode+Layout.mm b/Source/ASDisplayNode+Layout.mm
index 1506c7824..8006b13ff 100644
--- a/Source/ASDisplayNode+Layout.mm
+++ b/Source/ASDisplayNode+Layout.mm
@@ -143,7 +143,7 @@ - (NSString *)asciiArtName
 {
   NSString *string = NSStringFromClass([self class]);
   if (_debugName) {
-    string = [string stringByAppendingString:[NSString stringWithFormat:@"\"%@\"", _debugName]];
+    string = [string stringByAppendingString:[NSString stringWithFormat:@"\"%@\"",_debugName]];
   }
   return string;
 }
diff --git a/Source/ASDisplayNode+Yoga.mm b/Source/ASDisplayNode+Yoga.mm
index c47b64fe7..cc41593a6 100644
--- a/Source/ASDisplayNode+Yoga.mm
+++ b/Source/ASDisplayNode+Yoga.mm
@@ -40,7 +40,7 @@ @implementation ASDisplayNode (Yoga)
 
 - (void)setYogaChildren:(NSArray *)yogaChildren
 {
-  for (ASDisplayNode *child in _yogaChildren) {
+  for (ASDisplayNode *child in [_yogaChildren copy]) {
     // Make sure to un-associate the YGNodeRef tree before replacing _yogaChildren
     // If this becomes a performance bottleneck, it can be optimized by not doing the NSArray removals here.
     [self removeYogaChild:child];
@@ -68,14 +68,8 @@ - (void)addYogaChild:(ASDisplayNode *)child
   // Clean up state in case this child had another parent.
   [self removeYogaChild:child];
 
-  BOOL hadZeroChildren = (_yogaChildren.count == 0);
-
   [_yogaChildren addObject:child];
 
-  // Ensure any measure function is removed before inserting the YGNodeRef child.
-  if (hadZeroChildren) {
-    [self updateYogaMeasureFuncIfNeeded];
-  }
   // YGNodeRef insertion is done in setParent:
   child.yogaParent = self;
 }
@@ -86,15 +80,10 @@ - (void)removeYogaChild:(ASDisplayNode *)child
     return;
   }
   
-  BOOL hadChildren = (_yogaChildren.count > 0);
   [_yogaChildren removeObjectIdenticalTo:child];
 
   // YGNodeRef removal is done in setParent:
   child.yogaParent = nil;
-  // Ensure any measure function is re-added after removing the YGNodeRef child.
-  if (hadChildren && _yogaChildren.count == 0) {
-    [self updateYogaMeasureFuncIfNeeded];
-  }
 }
 
 - (void)semanticContentAttributeDidChange:(UISemanticContentAttribute)attribute
@@ -144,6 +133,7 @@ - (ASLayout *)yogaCalculatedLayout
 - (void)setYogaLayoutInProgress:(BOOL)yogaLayoutInProgress
 {
   setFlag(YogaLayoutInProgress, yogaLayoutInProgress);
+  [self updateYogaMeasureFuncIfNeeded];
 }
 
 - (BOOL)yogaLayoutInProgress
@@ -184,8 +174,14 @@ - (void)updateYogaMeasureFuncIfNeeded
 {
   // Size calculation via calculateSizeThatFits: or layoutSpecThatFits:
   // This will be used for ASTextNode, as well as any other node that has no Yoga children
-  id <ASLayoutElement> layoutElementToMeasure = (self.yogaChildren.count == 0 ? self : nil);
-  ASLayoutElementYogaUpdateMeasureFunc(self.style.yogaNode, layoutElementToMeasure);
+  BOOL isLeafNode = (self.yogaChildren.count == 0);
+  BOOL definesCustomLayout = [self implementsLayoutMethod];
+
+  // We set the measure func only during layout. Otherwise, a cycle is created:
+  // The YGNodeRef Context will retain the ASDisplayNode, which retains the style, which owns the YGNodeRef.
+  BOOL shouldHaveMeasureFunc = (isLeafNode && definesCustomLayout && checkFlag(YogaLayoutInProgress));
+
+  ASLayoutElementYogaUpdateMeasureFunc(self.style.yogaNode, shouldHaveMeasureFunc ? self : nil);
 }
 
 - (void)invalidateCalculatedYogaLayout
@@ -214,7 +210,6 @@ - (void)calculateLayoutFromYogaRoot:(ASSizeRange)rootConstrainedSize
   // Prepare all children for the layout pass with the current Yoga tree configuration.
   ASDisplayNodePerformBlockOnEveryYogaChild(self, ^(ASDisplayNode * _Nonnull node) {
     node.yogaLayoutInProgress = YES;
-    [node updateYogaMeasureFuncIfNeeded];
   });
 
   if (ASSizeRangeEqualToSizeRange(rootConstrainedSize, ASSizeRangeUnconstrained)) {