From 835bb1cd9c30884f28c5c68b43dd6748b56c5961 Mon Sep 17 00:00:00 2001
From: Dustin Shahidehpour <dust@instagram.com>
Date: Tue, 20 Dec 2016 11:31:55 -0800
Subject: [PATCH] Add isLeaf flag to Views.

Summary: While building some views with YogaKit, I found an interesting problem. If you are apply `yg_` properties to a view which is constructed of subclassed views (that don't use Yoga.) we still treat view (and it's subviews) like they are using YogaKit. This fixes that behavior.

Reviewed By: emilsjolander

Differential Revision: D4348344

fbshipit-source-id: 37b8648918529f55afea1d1883a03af398aac849
---
 YogaKit/Tests/YogaKitTests.m | 21 +++++++++++++++++++++
 YogaKit/UIView+Yoga.h        |  5 +++++
 YogaKit/UIView+Yoga.m        | 25 ++++++++++++++++++++-----
 3 files changed, 46 insertions(+), 5 deletions(-)

diff --git a/YogaKit/Tests/YogaKitTests.m b/YogaKit/Tests/YogaKitTests.m
index 99927ade..b54e308b 100644
--- a/YogaKit/Tests/YogaKitTests.m
+++ b/YogaKit/Tests/YogaKitTests.m
@@ -268,4 +268,25 @@
   }
 }
 
+- (void)testyg_isLeafFlag
+{
+  UIView *view = [[UIView alloc] initWithFrame:CGRectZero];
+  XCTAssertTrue(view.yg_isLeaf);
+
+  for (int i=0; i<10; i++) {
+    UIView *subview = [[UIView alloc] initWithFrame:CGRectZero];
+    [view addSubview:subview];
+  }
+  XCTAssertTrue(view.yg_isLeaf);
+
+  [view yg_setUsesYoga:YES];
+  [view yg_setWidth:50.0];
+  XCTAssertTrue(view.yg_isLeaf);
+
+  UIView *const subview = view.subviews[0];
+  [subview yg_setUsesYoga:YES];
+  [subview yg_setWidth:50.0];
+  XCTAssertFalse(view.yg_isLeaf);
+}
+
 @end
diff --git a/YogaKit/UIView+Yoga.h b/YogaKit/UIView+Yoga.h
index 386416b2..734c13d4 100644
--- a/YogaKit/UIView+Yoga.h
+++ b/YogaKit/UIView+Yoga.h
@@ -69,4 +69,9 @@
  */
 - (NSUInteger)yg_numberOfChildren;
 
+/**
+ Return a BOOL indiciating whether or not we this node contains any subviews that are included in Yoga's layout.
+ */
+- (BOOL)yg_isLeaf;
+
 @end
diff --git a/YogaKit/UIView+Yoga.m b/YogaKit/UIView+Yoga.m
index 241ad33e..776304f0 100644
--- a/YogaKit/UIView+Yoga.m
+++ b/YogaKit/UIView+Yoga.m
@@ -56,6 +56,20 @@
   return YGNodeGetChildCount([self ygNode]);
 }
 
+- (BOOL)yg_isLeaf
+{
+  NSAssert([NSThread isMainThread], @"This method must be called on the main thread.");
+  if ([self yg_usesYoga]) {
+      for (UIView *subview in self.subviews) {
+          if ([subview yg_usesYoga] && [subview yg_includeInLayout]) {
+              return NO;
+          }
+      }
+  }
+
+  return YES;
+}
+
 #pragma mark - Setters
 
 - (void)yg_setIncludeInLayout:(BOOL)includeInLayout
@@ -276,11 +290,12 @@ static CGFloat YGSanitizeMeasurement(
   return result;
 }
 
-static void YGAttachNodesFromViewHierachy(UIView *view) {
+static void YGAttachNodesFromViewHierachy(UIView *view)
+{
   YGNodeRef node = [view ygNode];
 
   // Only leaf nodes should have a measure function
-  if (![view yg_usesYoga] || view.subviews.count == 0) {
+  if (view.yg_isLeaf) {
     YGNodeSetMeasureFunc(node, YGMeasureView);
     YGRemoveAllChildren(node);
   } else {
@@ -340,7 +355,8 @@ static CGFloat YGRoundPixelValue(CGFloat value)
   return round(value * scale) / scale;
 }
 
-static void YGApplyLayoutToViewHierarchy(UIView *view) {
+static void YGApplyLayoutToViewHierarchy(UIView *view)
+{
   NSCAssert([NSThread isMainThread], @"Framesetting should only be done on the main thread.");
   if (![view yg_includeInLayout]) {
      return;
@@ -368,8 +384,7 @@ static void YGApplyLayoutToViewHierarchy(UIView *view) {
     },
   };
 
-  const BOOL isLeaf = ![view yg_usesYoga] || view.subviews.count == 0;
-  if (!isLeaf) {
+  if (!view.yg_isLeaf) {
     for (NSUInteger i = 0; i < view.subviews.count; i++) {
       YGApplyLayoutToViewHierarchy(view.subviews[i]);
     }