[Yoga] Minimize number of nodes that have MeasureFunc set on them.

Source/ASDisplayNode+Layout.mm

@@ -25,6 +25,15 @@ @implementation ASDisplayNode (ASLayoutElement)
 
 #pragma mark <ASLayoutElement>
 
+- (BOOL)implementsLayoutMethod
+{
+  ASDN::MutexLocker l(__instanceLock__);
+  return (_methodOverrides & (ASDisplayNodeMethodOverrideLayoutSpecThatFits |
+                              ASDisplayNodeMethodOverrideCalcLayoutThatFits |
+                              ASDisplayNodeMethodOverrideCalcSizeThatFits)) != 0 || _layoutSpecBlock != nil;
+}
+
+
 - (ASLayoutElementStyle *)style
 {
   ASDN::MutexLocker l(__instanceLock__);
@@ -148,9 +157,9 @@ @implementation ASDisplayNode (ASLayout)
 
 - (void)setLayoutSpecBlock:(ASLayoutSpecBlock)layoutSpecBlock
 {
-  // For now there should never be an override of layoutSpecThatFits: / layoutElementThatFits: and a layoutSpecBlock
-  ASDisplayNodeAssert(!(_methodOverrides & ASDisplayNodeMethodOverrideLayoutSpecThatFits), @"Overwriting layoutSpecThatFits: and providing a layoutSpecBlock block is currently not supported");
-
+  // For now there should never be an override of layoutSpecThatFits: and a layoutSpecBlock together.
+  ASDisplayNodeAssert(!(_methodOverrides & ASDisplayNodeMethodOverrideLayoutSpecThatFits),
+                      @"Nodes with a .layoutSpecBlock must not also implement -layoutSpecThatFits:");
   ASDN::MutexLocker l(__instanceLock__);
   _layoutSpecBlock = layoutSpecBlock;
 }

Source/ASDisplayNode.mm

@@ -177,6 +177,15 @@ static ASDisplayNodeMethodOverrides GetASDisplayNodeMethodOverrides(Class c)
   if (ASDisplayNodeSubclassOverridesSelector(c, @selector(layoutSpecThatFits:))) {
     overrides |= ASDisplayNodeMethodOverrideLayoutSpecThatFits;
   }
+  if (ASDisplayNodeSubclassOverridesSelector(c, @selector(calculateLayoutThatFits:)) ||
+      ASDisplayNodeSubclassOverridesSelector(c, @selector(calculateLayoutThatFits:
+                                                                 restrictedToSize:
+                                                             relativeToParentSize:))) {
+    overrides |= ASDisplayNodeMethodOverrideCalcLayoutThatFits;
+  }
+  if (ASDisplayNodeSubclassOverridesSelector(c, @selector(calculateSizeThatFits:))) {
+    overrides |= ASDisplayNodeMethodOverrideCalcSizeThatFits;
+  }
   if (ASDisplayNodeSubclassOverridesSelector(c, @selector(fetchData))) {
     overrides |= ASDisplayNodeMethodOverrideFetchData;
   }

Source/Layout/ASLayoutElement.h

@@ -144,6 +144,7 @@ typedef NS_ENUM(NSUInteger, ASLayoutElementType) {
                      restrictedToSize:(ASLayoutElementSize)size
                  relativeToParentSize:(CGSize)parentSize;
 
+- (BOOL)implementsLayoutMethod;
 
 #pragma mark - Deprecated
 

Source/Layout/ASLayoutSpec.mm

@@ -71,6 +71,11 @@ - (BOOL)canLayoutAsynchronous
   return YES;
 }
 
+- (BOOL)implementsLayoutMethod
+{
+  return YES;
+}
+
 #pragma mark - Style
 
 - (ASLayoutElementStyle *)style

Source/Layout/ASYogaUtilities.h

@@ -22,8 +22,9 @@
 @interface ASDisplayNode (YogaHelpers)
 
 + (ASDisplayNode *)yogaNode;
-+ (ASDisplayNode *)verticalYogaStack;
-+ (ASDisplayNode *)horizontalYogaStack;
++ (ASDisplayNode *)yogaSpacerNode;
++ (ASDisplayNode *)yogaVerticalStack;
++ (ASDisplayNode *)yogaHorizontalStack;
 
 @end
 

Source/Layout/ASYogaUtilities.mm

@@ -24,18 +24,25 @@ + (ASDisplayNode *)yogaNode
   return node;
 }
 
-+ (ASDisplayNode *)verticalYogaStack
++ (ASDisplayNode *)yogaSpacerNode
 {
-  ASDisplayNode *stack = [self yogaNode];
-  stack.style.flexDirection = ASStackLayoutDirectionVertical;
-  return stack;
+  ASDisplayNode *node = [ASDisplayNode yogaNode];
+  node.style.flexGrow = 1.0f;
+  return node;
 }
 
-+ (ASDisplayNode *)horizontalYogaStack
++ (ASDisplayNode *)yogaVerticalStack
 {
-  ASDisplayNode *stack = [self yogaNode];
-  stack.style.flexDirection = ASStackLayoutDirectionHorizontal;
-  return stack;
+  ASDisplayNode *node = [self yogaNode];
+  node.style.flexDirection = ASStackLayoutDirectionVertical;
+  return node;
+}
+
++ (ASDisplayNode *)yogaHorizontalStack
+{
+  ASDisplayNode *node = [self yogaNode];
+  node.style.flexDirection = ASStackLayoutDirectionHorizontal;
+  return node;
 }
 
 @end
@@ -141,14 +148,18 @@ void ASLayoutElementYogaUpdateMeasureFunc(YGNodeRef yogaNode, id <ASLayoutElemen
     return;
   }
   BOOL hasMeasureFunc = (YGNodeGetMeasureFunc(yogaNode) != NULL);
-  if (layoutElement != nil && hasMeasureFunc == NO) {
-    // 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 && ...) {
-    // Retain the Context object. This must be explicitly released with a __bridge_transfer; YGNodeFree() is not sufficient.
-    YGNodeSetContext(yogaNode, (__bridge_retained void *)layoutElement);
-    YGNodeSetMeasureFunc(yogaNode, &ASLayoutElementYogaMeasureFunc);
-  } else if (layoutElement == nil && hasMeasureFunc == YES){
+
+  if (layoutElement != nil && [layoutElement implementsLayoutMethod]) {
+    if (hasMeasureFunc == NO) {
+      // Retain the Context object. This must be explicitly released with a
+      // __bridge_transfer - YGNodeFree() is not sufficient.
+      YGNodeSetContext(yogaNode, (__bridge_retained void *)layoutElement);
+      YGNodeSetMeasureFunc(yogaNode, &ASLayoutElementYogaMeasureFunc);
+    }
+    ASDisplayNodeCAssert(YGNodeGetContext(yogaNode) == (__bridge void *)layoutElement,
+                         @"Yoga node context should contain layoutElement: %@", layoutElement);
+  } else if (hasMeasureFunc == YES) {
+    // If we lack any of the conditions above, and currently have a measure func, get rid of it.
     // Release the __bridge_retained Context object.
     __unused id <ASLayoutElement> element = (__bridge_transfer id)YGNodeGetContext(yogaNode);
     YGNodeSetContext(yogaNode, NULL);

Source/Private/ASDisplayNodeInternal.h

@@ -51,8 +51,10 @@ typedef NS_OPTIONS(NSUInteger, ASDisplayNodeMethodOverrides)
   ASDisplayNodeMethodOverrideTouchesEnded       = 1 << 2,
   ASDisplayNodeMethodOverrideTouchesMoved       = 1 << 3,
   ASDisplayNodeMethodOverrideLayoutSpecThatFits = 1 << 4,
-  ASDisplayNodeMethodOverrideFetchData          = 1 << 5,
-  ASDisplayNodeMethodOverrideClearFetchedData   = 1 << 6
+  ASDisplayNodeMethodOverrideCalcLayoutThatFits = 1 << 5,
+  ASDisplayNodeMethodOverrideCalcSizeThatFits   = 1 << 6,
+  ASDisplayNodeMethodOverrideFetchData          = 1 << 7,
+  ASDisplayNodeMethodOverrideClearFetchedData   = 1 << 8
 };
 
 typedef NS_OPTIONS(uint_least32_t, ASDisplayNodeAtomicFlags)

examples/ASDKgram/Sample/PhotoCellNode.m

@@ -316,7 +316,7 @@ - (void)setupYogaLayoutIfNeeded
   [_photoLocationLabel.style yogaNodeCreateIfNeeded];
   [_photoTimeIntervalSincePostLabel.style yogaNodeCreateIfNeeded];
 
-  ASDisplayNode *headerStack = [ASDisplayNode horizontalYogaStack];
+  ASDisplayNode *headerStack = [ASDisplayNode yogaHorizontalStack];
   headerStack.style.margin = ASEdgeInsetsMake(InsetForHeader);
   headerStack.style.alignItems = ASStackLayoutAlignItemsCenter;
   headerStack.style.flexGrow = 1.0;
@@ -327,7 +327,7 @@ - (void)setupYogaLayoutIfNeeded
   [headerStack addYogaChild:_userAvatarImageNode];
 
   // User Name and Photo Location stack is next
-  ASDisplayNode *userPhotoLocationStack = [ASDisplayNode verticalYogaStack];
+  ASDisplayNode *userPhotoLocationStack = [ASDisplayNode yogaVerticalStack];
   userPhotoLocationStack.style.flexShrink = 1.0;
   [headerStack addYogaChild:userPhotoLocationStack];
 
@@ -340,20 +340,15 @@ - (void)setupYogaLayoutIfNeeded
     [userPhotoLocationStack addYogaChild:_photoLocationLabel];
   }
 
-/* TODO: These parameters aren't working as expected. For now the timestamp is next to the username.
   // Add a spacer to allow a flexible space between the User Name / Location stack, and the Timestamp.
-  ASDisplayNode *spacer = [ASDisplayNode new];
-  spacer.style.flexShrink = 1.0;
-  spacer.style.width = ASDimensionMakeWithFraction(1.0);
-  [headerStack addYogaChild:spacer];
-*/
+  [headerStack addYogaChild:[ASDisplayNode yogaSpacerNode]];
 
   // Photo Timestamp Label.
   _photoTimeIntervalSincePostLabel.style.spacingBefore = HORIZONTAL_BUFFER;
   [headerStack addYogaChild:_photoTimeIntervalSincePostLabel];
 
   // Create the last stack before assembling everything: the Footer Stack contains the description and comments.
-  ASDisplayNode *footerStack = [ASDisplayNode verticalYogaStack];
+  ASDisplayNode *footerStack = [ASDisplayNode yogaVerticalStack];
   footerStack.style.margin = ASEdgeInsetsMake(InsetForFooter);
   footerStack.style.padding = ASEdgeInsetsMake(UIEdgeInsetsMake(0.0, 0.0, VERTICAL_BUFFER, 0.0));
   footerStack.yogaChildren = @[_photoLikesLabel, _photoDescriptionLabel, _photoCommentsNode];