[Layout] Check Layout Spec Tree for Duplicated Elements (#2408)

* Check layout spec tree for duplicated elements

* Enable layout spec tree dedupe only in debug

AsyncDisplayKit/ASDisplayNode.mm

@@ -34,6 +34,11 @@
 #import "ASLayoutSpec.h"
 #import "ASCellNode+Internal.h"
 #import "ASWeakProxy.h"
+#import "ASLayoutSpecPrivate.h"
+
+#if DEBUG
+  #define AS_DEDUPE_LAYOUT_SPEC_TREE 1
+#endif
 
 NSInteger const ASDefaultDrawingPriority = ASDefaultTransactionPriority;
 NSString * const ASRenderingEngineDidDisplayScheduledNodesNotification = @"ASRenderingEngineDidDisplayScheduledNodes";
@@ -2426,6 +2431,15 @@ void recursivelyTriggerDisplayForLayer(CALayer *layer, BOOL shouldBlock)
       [self layoutSpecThatFits:constrainedSize];
     });
 
+#if AS_DEDUPE_LAYOUT_SPEC_TREE
+    NSSet *duplicateElements = [layoutSpec findDuplicatedElementsInSubtree];
+    if (duplicateElements.count > 0) {
+      ASDisplayNodeFailAssert(@"Node %@ returned a layout spec that contains the same elements in multiple positions. Elements: %@", self, duplicateElements);
+      // Use an empty layout spec to avoid crash.
+      layoutSpec = [[ASLayoutSpec alloc] init];
+    }
+#endif
+
     ASDisplayNodeAssert(layoutSpec.isMutable, @"Node %@ returned layout spec %@ that has already been used. Layout specs should always be regenerated.", self, layoutSpec);
 
     layoutSpec.parent = self; // This causes upward propogation of any non-default layoutElement values.

AsyncDisplayKit/Layout/ASLayoutSpec.mm

@@ -148,11 +148,7 @@
 {
   ASDisplayNodeAssert(_childrenArray.count < 2, @"This layout spec does not support more than one child. Use the setChildren: or the setChild:AtIndex: API");
   
-  if (_childrenArray.count) {
-    return _childrenArray[0];
-  }
-  
-  return nil;
+  return _childrenArray.firstObject;
 }
 
 #pragma mark - Children
@@ -235,6 +231,51 @@
 
 ASEnvironmentLayoutExtensibilityForwarding
 
+#pragma mark - Framework Private
+
+- (nullable NSSet<id<ASLayoutElement>> *)findDuplicatedElementsInSubtree
+{
+  NSMutableSet *result = nil;
+  NSUInteger count = 0;
+  [self _findDuplicatedElementsInSubtreeWithWorkingSet:[[NSMutableSet alloc] init] workingCount:&count result:&result];
+  return result;
+}
+
+/**
+ * This method is extremely performance-sensitive, so we do some strange things.
+ *
+ * @param workingSet A working set of elements for use in the recursion.
+ * @param workingCount The current count of the set for use in the recursion.
+ * @param result The set into which to put the result. This initially points to @c nil to save time if no duplicates exist.
+ */
+- (void)_findDuplicatedElementsInSubtreeWithWorkingSet:(NSMutableSet<id<ASLayoutElement>> *)workingSet workingCount:(NSUInteger *)workingCount result:(NSMutableSet<id<ASLayoutElement>>  * _Nullable *)result
+{
+  Class layoutSpecClass = [ASLayoutSpec class];
+
+  for (id<ASLayoutElement> child in self) {
+    // Add the object into the set.
+    [workingSet addObject:child];
+
+    // Check that addObject: caused the count to increase.
+    // This is faster than using containsObject.
+    NSUInteger oldCount = *workingCount;
+    NSUInteger newCount = workingSet.count;
+    BOOL objectAlreadyExisted = (newCount != oldCount + 1);
+    if (objectAlreadyExisted) {
+      if (*result == nil) {
+        *result = [[NSMutableSet alloc] init];
+      }
+      [*result addObject:child];
+    } else {
+      *workingCount = newCount;
+      // If child is a layout spec we haven't visited, recurse its children.
+      if ([child isKindOfClass:layoutSpecClass]) {
+        [(ASLayoutSpec *)child _findDuplicatedElementsInSubtreeWithWorkingSet:workingSet workingCount:workingCount result:result];
+      }
+    }
+  }
+}
+
 @end
 
 #pragma mark - ASWrapperLayoutSpec

AsyncDisplayKit/Private/ASLayoutSpecPrivate.h

@@ -15,11 +15,20 @@
 #import "ASEnvironmentInternal.h"
 #import "ASThread.h"
 
+NS_ASSUME_NONNULL_BEGIN
+
 @interface ASLayoutSpec() {
   ASDN::RecursiveMutex __instanceLock__;
   ASEnvironmentState _environmentState;
   ASLayoutElementStyle *_style;
   NSMutableArray *_childrenArray;
 }
+
+/**
+ * Recursively search the subtree for elements that occur more than once.
+ */
+- (nullable NSSet<id<ASLayoutElement>> *)findDuplicatedElementsInSubtree;
+
 @end
 
+NS_ASSUME_NONNULL_END

AsyncDisplayKitTests/ASDisplayNodeTests.m

@@ -21,6 +21,8 @@
 #import "UIView+ASConvenience.h"
 #import "ASCellNode.h"
 #import "ASImageNode.h"
+#import "ASOverlayLayoutSpec.h"
+#import "ASInsetLayoutSpec.h"
 
 // Conveniences for making nodes named a certain way
 #define DeclareNodeNamed(n) ASDisplayNode *n = [[ASDisplayNode alloc] init]; n.name = @#n
@@ -2139,4 +2141,14 @@ static bool stringContainsPointer(NSString *description, id p) {
   }
 }
 
+- (void)testThatHavingTheSameNodeTwiceInALayoutSpecCausesExceptionOnLayoutCalculation
+{
+  ASDisplayNode *node = [[ASDisplayNode alloc] init];
+  ASDisplayNode *subnode = [[ASDisplayNode alloc] init];
+  node.layoutSpecBlock = ^ASLayoutSpec *(ASDisplayNode *node, ASSizeRange constrainedSize) {
+    return [ASOverlayLayoutSpec overlayLayoutSpecWithChild:[ASInsetLayoutSpec insetLayoutSpecWithInsets:UIEdgeInsetsZero child:subnode] overlay:subnode];
+  };
+  XCTAssertThrowsSpecificNamed([node calculateLayoutThatFits:ASSizeRangeMake(CGSizeMake(100, 100))], NSException, NSInternalInconsistencyException);
+}
+
 @end