From 7ec3607446ba91a25801e60232fe8b33a12f7ab6 Mon Sep 17 00:00:00 2001
From: Dustin Shahidehpour <dust@instagram.com>
Date: Wed, 21 Dec 2016 10:54:31 -0800
Subject: [PATCH] Fix diffing algorithm for reattaching views.

Summary:
There is a bug currently where we don't traverse the entire tree to detect view hierarchy changes.

Currently, if you had a hierachy like this:
```
container = UIView
container.subviews = @[subview1, subview2];

subview1.subviews = @[sub11, sub12, sub13];
subview2.subviews = @[sub21, sub22, sub23];
```

and then modified via:
```
subview2.subviews = @[newView1, newView2, newView3];
```

our algorithm wouldn't identify that we had new views that needed their layout calculated, and would cause a crash later on.

Reviewed By: emilsjolander

Differential Revision: D4357662

fbshipit-source-id: 2f61f213c5f1a62948a653f3b1fa3d874c5075f7
---
 YogaKit/Tests/YogaKitTests.m | 54 ++++++++++++++++++++++++++++++++++++
 YogaKit/UIView+Yoga.m        | 43 +++++++++++++++-------------
 2 files changed, 77 insertions(+), 20 deletions(-)

diff --git a/YogaKit/Tests/YogaKitTests.m b/YogaKit/Tests/YogaKitTests.m
index b54e308b..68d97dd5 100644
--- a/YogaKit/Tests/YogaKitTests.m
+++ b/YogaKit/Tests/YogaKitTests.m
@@ -289,4 +289,58 @@
   XCTAssertFalse(view.yg_isLeaf);
 }
 
+- (void)testThatWeCorrectlyAttachNestedViews
+{
+    UIView *container = [[UIView alloc] initWithFrame:CGRectMake(0, 0, 300, 50)];
+    [container yg_setUsesYoga:YES];
+    [container yg_setFlexDirection:YGFlexDirectionColumn];
+
+    UIView *subview1 = [[UIView alloc] initWithFrame:CGRectZero];
+    [subview1 yg_setUsesYoga:YES];
+    [subview1 yg_setWidth:100];
+    [subview1 yg_setFlexGrow:1];
+    [subview1 yg_setFlexDirection:YGFlexDirectionColumn];
+    [container addSubview:subview1];
+
+    UIView *subview2 = [[UIView alloc] initWithFrame:CGRectZero];
+    [subview2 yg_setUsesYoga:YES];
+    [subview2 yg_setWidth:150];
+    [subview2 yg_setFlexGrow:1];
+    [subview2 yg_setFlexDirection:YGFlexDirectionColumn];
+    [container addSubview:subview2];
+
+    for (UIView *view in @[subview1, subview2]) {
+        UIView *someView = [[UIView alloc] initWithFrame:CGRectZero];
+        [someView yg_setUsesYoga:YES];
+        [someView yg_setFlexGrow:1];
+        [view addSubview:someView];
+    }
+    [container yg_applyLayout];
+
+    // Add the same amount of new views, reapply layout.
+    for (UIView *view in @[subview1, subview2]) {
+        UIView *someView = [[UIView alloc] initWithFrame:CGRectZero];
+        [someView yg_setUsesYoga:YES];
+        [someView yg_setFlexGrow:1];
+        [view addSubview:someView];
+    }
+    [container yg_applyLayout];
+
+    XCTAssertTrue(CGSizeEqualToSize(CGSizeMake(100, 25), subview1.bounds.size), @"Actual size is %@", NSStringFromCGSize(subview1.bounds.size));
+    for (UIView *subview in subview1.subviews) {
+        const CGSize subviewSize = subview.bounds.size;
+        XCTAssertFalse(CGSizeEqualToSize(CGSizeZero, subviewSize));
+        XCTAssertFalse(isnan(subviewSize.height));
+        XCTAssertFalse(isnan(subviewSize.width));
+    }
+
+    XCTAssertTrue(CGSizeEqualToSize(CGSizeMake(150, 25), subview2.bounds.size), @"Actual size is %@", NSStringFromCGSize(subview2.bounds.size));
+    for (UIView *subview in subview2.subviews) {
+        const CGSize subviewSize = subview.bounds.size;
+        XCTAssertFalse(CGSizeEqualToSize(CGSizeZero, subview.bounds.size));
+        XCTAssertFalse(isnan(subviewSize.height));
+        XCTAssertFalse(isnan(subviewSize.width));
+    }
+}
+
 @end
diff --git a/YogaKit/UIView+Yoga.m b/YogaKit/UIView+Yoga.m
index 776304f0..980a7683 100644
--- a/YogaKit/UIView+Yoga.m
+++ b/YogaKit/UIView+Yoga.m
@@ -290,9 +290,24 @@ static CGFloat YGSanitizeMeasurement(
   return result;
 }
 
-static void YGAttachNodesFromViewHierachy(UIView *view)
+static BOOL YGNodeHasExactSameChildren(const YGNodeRef node, NSArray<UIView *> *subviews)
 {
-  YGNodeRef node = [view ygNode];
+  if (YGNodeGetChildCount(node) != subviews.count) {
+    return NO;
+  }
+
+  for (int i=0; i<subviews.count; i++) {
+    if (YGNodeGetChild(node, i) != subviews[i].ygNode) {
+      return NO;
+    }
+  }
+
+  return YES;
+}
+
+static void YGAttachNodesFromViewHierachy(UIView *const view)
+{
+  const YGNodeRef node = [view ygNode];
 
   // Only leaf nodes should have a measure function
   if (view.yg_isLeaf) {
@@ -301,7 +316,6 @@ static void YGAttachNodesFromViewHierachy(UIView *view)
   } else {
     YGNodeSetMeasureFunc(node, NULL);
 
-    // Create a list of all the subviews that we are going to use for layout.
     NSMutableArray<UIView *> *subviewsToInclude = [[NSMutableArray alloc] initWithCapacity:view.subviews.count];
     for (UIView *subview in view.subviews) {
       if ([subview yg_includeInLayout]) {
@@ -309,26 +323,15 @@ static void YGAttachNodesFromViewHierachy(UIView *view)
       }
     }
 
-    BOOL shouldReconstructChildList = NO;
-    if (YGNodeGetChildCount(node) != subviewsToInclude.count) {
-      shouldReconstructChildList = YES;
-    } else {
-      for (int i = 0; i < subviewsToInclude.count; i++) {
-        if (YGNodeGetChild(node, i) != [subviewsToInclude[i] ygNode]) {
-          shouldReconstructChildList = YES;
-          break;
+    if (!YGNodeHasExactSameChildren(node, subviewsToInclude)) {
+        YGRemoveAllChildren(node);
+        for (int i=0; i<subviewsToInclude.count; i++) {
+            YGNodeInsertChild(node, [subviewsToInclude[i] ygNode], i);
         }
-      }
     }
 
-    if (shouldReconstructChildList) {
-      YGRemoveAllChildren(node);
-
-      for (int i = 0 ; i < subviewsToInclude.count; i++) {
-        UIView *const subview = subviewsToInclude[i];
-        YGNodeInsertChild(node, [subview ygNode], i);
+    for (UIView *const subview in subviewsToInclude) {
         YGAttachNodesFromViewHierachy(subview);
-      }
     }
   }
 }
@@ -385,7 +388,7 @@ static void YGApplyLayoutToViewHierarchy(UIView *view)
   };
 
   if (!view.yg_isLeaf) {
-    for (NSUInteger i = 0; i < view.subviews.count; i++) {
+    for (NSUInteger i=0; i<view.subviews.count; i++) {
       YGApplyLayoutToViewHierarchy(view.subviews[i]);
     }
   }