[Layout] Don't crash if layout elements are created in layoutSpecThatFits: (#2694)

* Fix crash if layout elements are created with no owner and referenced in layoutSpecThatFits:

* Add failing test for nodes deallocated while creating in layoutSpecThatFits:

* Some more

* Some cleanup

* Added more complexity to tests

* Only cache sublayouts if the layout get’s flattened

* Address comments

* Address comments

AsyncDisplayKit/Layout/ASLayout.mm

@@ -47,6 +47,16 @@ static inline NSString * descriptionIndents(NSUInteger indents)
  */
 @property (nonatomic, getter=isFlattened) BOOL flattened;
 
+/*
+ * Caches all sublayouts if set to YES or destroys the sublayout cache if set to NO. Defaults to YES
+ */
+@property (nonatomic, assign) BOOL retainSublayoutLayoutElements;
+
+/**
+ * Array for explicitly retain sublayout layout elements in case they are created and references in layoutSpecThatFits: and no one else will hold a strong reference on it
+ */
+@property (nonatomic, strong) NSMutableArray<id<ASLayoutElement>> *sublayoutLayoutElements;
+
 @end
 
 @implementation ASLayout
@@ -69,6 +79,7 @@ static inline NSString * descriptionIndents(NSUInteger indents)
 #endif
     
     _layoutElement = layoutElement;
+    
     // Read this now to avoid @c weak overhead later.
     _layoutElementType = layoutElement.layoutElementType;
     
@@ -88,7 +99,9 @@ static inline NSString * descriptionIndents(NSUInteger indents)
 
     _sublayouts = sublayouts != nil ? [sublayouts copy] : @[];
     _flattened = NO;
+    _retainSublayoutLayoutElements = NO;
   }
+  
   return self;
 }
 
@@ -137,6 +150,28 @@ static inline NSString * descriptionIndents(NSUInteger indents)
                             sublayouts:layout.sublayouts];
 }
 
+#pragma mark - Sublayout Elements Caching
+
+- (void)setRetainSublayoutLayoutElements:(BOOL)retainSublayoutLayoutElements
+{
+  if (_retainSublayoutLayoutElements != retainSublayoutLayoutElements) {
+    _retainSublayoutLayoutElements = retainSublayoutLayoutElements;
+    
+    if (retainSublayoutLayoutElements == NO) {
+      _sublayoutLayoutElements = nil;
+    } else {
+      // Add sublayouts layout elements to an internal array to retain it while the layout lives
+      NSUInteger sublayoutCount = _sublayouts.count;
+      if (sublayoutCount > 0) {
+        _sublayoutLayoutElements = [NSMutableArray arrayWithCapacity:sublayoutCount];
+        for (ASLayout *sublayout in _sublayouts) {
+          [_sublayoutLayoutElements addObject:sublayout.layoutElement];
+        }
+      }
+    }
+  }
+}
+
 #pragma mark - Layout Flattening
 
 - (ASLayout *)filteredNodeLayoutTree
@@ -170,8 +205,10 @@ static inline NSString * descriptionIndents(NSUInteger indents)
     }
     queue.insert(queue.cbegin(), sublayoutContexts.begin(), sublayoutContexts.end());
   }
-
-  return [ASLayout layoutWithLayoutElement:_layoutElement size:_size sublayouts:flattenedSublayouts];
+  
+  ASLayout *layout = [ASLayout layoutWithLayoutElement:_layoutElement size:_size position:CGPointZero sublayouts:flattenedSublayouts];
+  layout.retainSublayoutLayoutElements = YES;
+  return layout;
 }
 
 #pragma mark - Accessors

AsyncDisplayKitTests/ASDisplayNodeLayoutTests.mm

@@ -121,4 +121,54 @@
   XCTAssertThrows([ASLayout layoutWithLayoutElement:displayNode size:CGSizeMake(INFINITY, INFINITY)]);
 }
 
+- (void)testThatLayoutElementCreatedInLayoutSpecThatFitsDoNotGetDeallocated
+{
+  const CGSize kSize = CGSizeMake(300, 300);
+  
+  ASDisplayNode *subNode = [[ASDisplayNode alloc] init];
+  subNode.automaticallyManagesSubnodes = YES;
+  subNode.layoutSpecBlock = ^(ASDisplayNode * _Nonnull node, ASSizeRange constrainedSize) {
+    ASTextNode *textNode = [ASTextNode new];
+    textNode.attributedText = [[NSAttributedString alloc] initWithString:@"Test Test Test Test Test Test Test Test"];
+    ASInsetLayoutSpec *insetSpec = [ASInsetLayoutSpec insetLayoutSpecWithInsets:UIEdgeInsetsZero child:textNode];
+    return [ASInsetLayoutSpec insetLayoutSpecWithInsets:UIEdgeInsetsZero child:insetSpec];
+  };
+  
+  ASDisplayNode *rootNode = [[ASDisplayNode alloc] init];
+  rootNode.automaticallyManagesSubnodes = YES;
+  rootNode.layoutSpecBlock = ^(ASDisplayNode * _Nonnull node, ASSizeRange constrainedSize) {
+    ASTextNode *textNode = [ASTextNode new];
+    textNode.attributedText = [[NSAttributedString alloc] initWithString:@"Test Test Test Test Test"];
+    ASInsetLayoutSpec *insetSpec = [ASInsetLayoutSpec insetLayoutSpecWithInsets:UIEdgeInsetsZero child:textNode];
+    
+    return [ASStackLayoutSpec
+            stackLayoutSpecWithDirection:ASStackLayoutDirectionVertical
+            spacing:0.0
+            justifyContent:ASStackLayoutJustifyContentStart
+            alignItems:ASStackLayoutAlignItemsStretch
+            children:@[insetSpec, subNode]];
+  };
+
+  rootNode.frame = CGRectMake(0, 0, kSize.width, kSize.height);
+  [rootNode view];
+  
+  XCTestExpectation *expectation = [self expectationWithDescription:@"Execute measure and layout pass"];
+  
+  dispatch_async(dispatch_get_global_queue(DISPATCH_QUEUE_PRIORITY_DEFAULT, 0), ^{
+    
+    [rootNode layoutThatFits:ASSizeRangeMake(kSize)];
+    
+    dispatch_async(dispatch_get_main_queue(), ^{
+      XCTAssertNoThrow([rootNode.view layoutIfNeeded]);
+      [expectation fulfill];
+    });
+  });
+  
+  [self waitForExpectationsWithTimeout:5.0 handler:^(NSError *error) {
+    if (error) {
+      XCTFail(@"Expectation failed: %@", error);
+    }
+  }];
+}
+
 @end