From dc10fdd9585e12a172cc14650e002fc6f42140f5 Mon Sep 17 00:00:00 2001
From: Dustin Shahidehpour <dust@instagram.com>
Date: Tue, 22 Nov 2016 08:27:03 -0800
Subject: [PATCH] Fix bug where swapping views in hierarchy throws an assert.

Summary: I found that if you move a subview to a different index, it would cause a crash in subsequent layout passes because the `CSSNodeRef` representing it was being re-inserted into the list. I spent a lot of time figuring out the best way to "diff" the view hierarchy from before and after, and I found the easiest implementation is also the most reliable. We can continue to optimize, but this is a great start.

Reviewed By: emilsjolander

Differential Revision: D4218623

fbshipit-source-id: 06253089492ac37ae4b94b7c30140ce6ed680ed2
---
 CSSLayoutKit/Tests/CSSLayoutKitTests.m | 40 +++++++++++++++++
 CSSLayoutKit/UIView+CSSLayout.m        | 60 ++++++++++++++++----------
 2 files changed, 77 insertions(+), 23 deletions(-)

diff --git a/CSSLayoutKit/Tests/CSSLayoutKitTests.m b/CSSLayoutKit/Tests/CSSLayoutKitTests.m
index c6773895..22b159b2 100644
--- a/CSSLayoutKit/Tests/CSSLayoutKitTests.m
+++ b/CSSLayoutKit/Tests/CSSLayoutKitTests.m
@@ -125,6 +125,46 @@
   XCTAssertEqual(containerSize.width, totalWidth, @"The container's width is %.6f, the subviews take up %.6f", containerSize.width, totalWidth);
 }
 
+- (void)testThatLayoutIsCorrectWhenWeSwapViewOrder
+{
+  const CGSize containerSize = CGSizeMake(300, 50);
+
+  UIView *container = [[UIView alloc] initWithFrame:CGRectMake(0, 0, containerSize.width, containerSize.height)];
+  [container css_setUsesFlexbox:YES];
+  [container css_setFlexDirection:CSSFlexDirectionRow];
+
+  UIView *subview1 = [[UIView alloc] initWithFrame:CGRectZero];
+  [subview1 css_setUsesFlexbox:YES];
+  [subview1 css_setFlexGrow:1];
+  [container addSubview:subview1];
+
+  UIView *subview2 = [[UIView alloc] initWithFrame:CGRectZero];
+  [subview2 css_setUsesFlexbox:YES];
+  [subview2 css_setFlexGrow:1];
+  [container addSubview:subview2];
+
+  UIView *subview3 = [[UIView alloc] initWithFrame:CGRectZero];
+  [subview3 css_setUsesFlexbox:YES];
+  [subview3 css_setFlexGrow:1];
+  [container addSubview:subview3];
+
+  [container css_applyLayout];
+
+  XCTAssertTrue(CGRectEqualToRect(subview1.frame, CGRectMake(0, 0, 100, 50)));
+  XCTAssertTrue(CGRectEqualToRect(subview2.frame, CGRectMake(100, 0, 100, 50)), @"It's actually %@", NSStringFromCGRect(subview2.frame));
+  XCTAssertTrue(CGRectEqualToRect(subview3.frame, CGRectMake(200, 0, 100, 50)));
+
+  [container exchangeSubviewAtIndex:2 withSubviewAtIndex:0];
+  [subview2 css_setIncludeInLayout:NO];
+  [container css_applyLayout];
+
+  XCTAssertTrue(CGRectEqualToRect(subview3.frame, CGRectMake(0, 0, 150, 50)));
+  XCTAssertTrue(CGRectEqualToRect(subview1.frame, CGRectMake(150, 0, 150, 50)));
+
+  // this frame shouldn't have been modified since last time.
+  XCTAssertTrue(CGRectEqualToRect(subview2.frame, CGRectMake(100, 0, 100, 50)));
+}
+
 - (void)testThatWeRespectIncludeInLayoutFlag
 {
   const CGSize containerSize = CGSizeMake(300, 50);
diff --git a/CSSLayoutKit/UIView+CSSLayout.m b/CSSLayoutKit/UIView+CSSLayout.m
index 57c9e305..56ddfcf7 100644
--- a/CSSLayoutKit/UIView+CSSLayout.m
+++ b/CSSLayoutKit/UIView+CSSLayout.m
@@ -272,40 +272,54 @@ static CGFloat CSSSanitizeMeasurement(
 
 static void CSSAttachNodesFromViewHierachy(UIView *view) {
   CSSNodeRef node = [view cssNode];
-  const BOOL usesFlexbox = [view css_usesFlexbox];
-  const BOOL isLeaf = !usesFlexbox || view.subviews.count == 0;
 
   // Only leaf nodes should have a measure function
-  if (isLeaf) {
+  if (![view css_usesFlexbox] || view.subviews.count == 0) {
     CSSNodeSetMeasureFunc(node, CSSMeasureView);
-
-    // Clear any children
-    while (CSSNodeChildCount(node) > 0) {
-      CSSNodeRemoveChild(node, CSSNodeGetChild(node, 0));
-    }
+    CSSRemoveAllChildren(node);
   } else {
     CSSNodeSetMeasureFunc(node, NULL);
 
-    NSUInteger subviewIndex = 0;
-    // Add any children which were added since the last call to css_applyLayout
+    // 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 css_includeInLayout]) {
-        continue;
+      if ([subview css_includeInLayout]) {
+        [subviewsToInclude addObject:subview];
       }
-
-      CSSNodeRef childNode = [subview cssNode];
-      if (CSSNodeGetChild(node, subviewIndex) != childNode) {
-        CSSNodeInsertChild(node, childNode, subviewIndex);
-      }
-
-      CSSAttachNodesFromViewHierachy(subview);
-      subviewIndex++;
     }
 
-    // Remove any children which were removed since the last call to css_applyLayout
-    while (subviewIndex < CSSNodeChildCount(node)) {
-      CSSNodeRemoveChild(node, CSSNodeGetChild(node, CSSNodeChildCount(node) - 1));
+    BOOL shouldReconstructChildList = NO;
+    if (CSSNodeChildCount(node) != subviewsToInclude.count) {
+      shouldReconstructChildList = YES;
+    } else {
+      for (int i = 0; i < subviewsToInclude.count; i++) {
+        if (CSSNodeGetChild(node, i) != [subviewsToInclude[i] cssNode]) {
+          shouldReconstructChildList = YES;
+          break;
+        }
+      }
     }
+
+    if (shouldReconstructChildList) {
+      CSSRemoveAllChildren(node);
+
+      for (int i = 0 ; i < subviewsToInclude.count; i++) {
+        UIView *const subview = subviewsToInclude[i];
+        CSSNodeInsertChild(node, [subview cssNode], i);
+        CSSAttachNodesFromViewHierachy(subview);
+      }
+    }
+  }
+}
+
+static void CSSRemoveAllChildren(const CSSNodeRef node)
+{
+  if (node == NULL) {
+      return;
+  }
+
+  while (CSSNodeChildCount(node) > 0) {
+    CSSNodeRemoveChild(node, CSSNodeGetChild(node, CSSNodeChildCount(node) - 1));
   }
 }