[ASLayoutSpec] Use childrenMap directly to prevent creating an NSArray within ASDK Part 2 (#2021)

* Revert "Revert "[ASLayoutSpec] Use childrenMap directly to prevent creating an NSArray within ASDK (#1937)""

This reverts commit 735b4eb.

* Fix crash and add exception for mutating while using fast enumeration of ASLayoutSpec children

NSFastEnumeration is potentially quite dangerous in the wrong hands. In particular, it does not provide a safe mechanism for you to return temporary objects directly, and it does not provide any guarantee that you will be called when the enumeration has completed; therefore if we generate temporaries and store them in an instance variable, we will not necessarily be able to clean them up! This means fast enumeration methods should never be called within an autorelease pool or the autorelease pool be drained within the fast enumeration loop.

The reason is we store references to objects in the stackBuf struct by casting the child pointer to __autoreleasing id. If we pop the autorelease pool between calls to -countByEnumeratingWithState:objects:count:, it will die in a messy explosion of pointer dereferences and EXC_BAD_ACCESS.

* Add tests for ASDisplayNode and ASLayoutSpec fast enumeration

Author: maicki

AsyncDisplayKit.xcodeproj/project.pbxproj

@@ -214,6 +214,7 @@
 		69E1006F1CA89CB600D88C1B /* ASEnvironmentInternal.mm in Sources */ = {isa = PBXBuildFile; fileRef = 69E1006A1CA89CB600D88C1B /* ASEnvironmentInternal.mm */; };
 		69E100701CA89CB600D88C1B /* ASEnvironmentInternal.mm in Sources */ = {isa = PBXBuildFile; fileRef = 69E1006A1CA89CB600D88C1B /* ASEnvironmentInternal.mm */; };
 		69F10C871C84C35D0026140C /* ASRangeControllerUpdateRangeProtocol+Beta.h in Headers */ = {isa = PBXBuildFile; fileRef = 69F10C851C84C35D0026140C /* ASRangeControllerUpdateRangeProtocol+Beta.h */; settings = {ATTRIBUTES = (Public, ); }; };
+		69F6058D1D3DA27E00C8CA38 /* ASLayoutSpec+Private.h in Headers */ = {isa = PBXBuildFile; fileRef = 69F6058C1D3DA27E00C8CA38 /* ASLayoutSpec+Private.h */; };
 		7630FFA81C9E267E007A7C0E /* ASVideoNode.h in Headers */ = {isa = PBXBuildFile; fileRef = AEEC47DF1C20C2DD00EC1693 /* ASVideoNode.h */; settings = {ATTRIBUTES = (Public, ); }; };
 		764D83D51C8EA515009B4FB8 /* AsyncDisplayKit+Debug.h in Headers */ = {isa = PBXBuildFile; fileRef = 764D83D21C8EA515009B4FB8 /* AsyncDisplayKit+Debug.h */; settings = {ATTRIBUTES = (Public, ); }; };
 		764D83D61C8EA515009B4FB8 /* AsyncDisplayKit+Debug.m in Sources */ = {isa = PBXBuildFile; fileRef = 764D83D31C8EA515009B4FB8 /* AsyncDisplayKit+Debug.m */; };
@@ -965,6 +966,7 @@
 		69E100691CA89CB600D88C1B /* ASEnvironmentInternal.h */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.h; path = ASEnvironmentInternal.h; sourceTree = "<group>"; };
 		69E1006A1CA89CB600D88C1B /* ASEnvironmentInternal.mm */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.cpp.objcpp; path = ASEnvironmentInternal.mm; sourceTree = "<group>"; };
 		69F10C851C84C35D0026140C /* ASRangeControllerUpdateRangeProtocol+Beta.h */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.h; path = "ASRangeControllerUpdateRangeProtocol+Beta.h"; sourceTree = "<group>"; };
+		69F6058C1D3DA27E00C8CA38 /* ASLayoutSpec+Private.h */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.h; path = "ASLayoutSpec+Private.h"; sourceTree = "<group>"; };
 		6BDC61F51978FEA400E50D21 /* AsyncDisplayKit.h */ = {isa = PBXFileReference; explicitFileType = sourcecode.c.h; path = AsyncDisplayKit.h; sourceTree = "<group>"; };
 		764D83D21C8EA515009B4FB8 /* AsyncDisplayKit+Debug.h */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.h; path = "AsyncDisplayKit+Debug.h"; sourceTree = "<group>"; };
 		764D83D31C8EA515009B4FB8 /* AsyncDisplayKit+Debug.m */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.objc; path = "AsyncDisplayKit+Debug.m"; sourceTree = "<group>"; };
@@ -1470,6 +1472,8 @@
 				044285051BAA63FE00D16268 /* ASBatchFetching.h */,
 				044285061BAA63FE00D16268 /* ASBatchFetching.m */,
 				251B8EF61BBB3D690087C538 /* ASDataController+Subclasses.h */,
+				8B0768B11CE752EC002E1453 /* ASDefaultPlaybackButton.h */,
+				8B0768B21CE752EC002E1453 /* ASDefaultPlaybackButton.m */,
 				AEB7B0181C5962EA00662EF4 /* ASDefaultPlayButton.h */,
 				AEB7B0191C5962EA00662EF4 /* ASDefaultPlayButton.m */,
 				058D0A08195D050800B7D73C /* ASDisplayNode+AsyncDisplay.mm */,
@@ -1478,16 +1482,17 @@
 				DE6EA3211C14000600183B10 /* ASDisplayNode+FrameworkPrivate.h */,
 				058D0A0B195D050800B7D73C /* ASDisplayNode+UIViewBridge.mm */,
 				058D0A0C195D050800B7D73C /* ASDisplayNodeInternal.h */,
-				E52405B41C8FEF16004DC8E7 /* ASLayoutTransition.h */,
-				E52405B21C8FEF03004DC8E7 /* ASLayoutTransition.mm */,
 				69E100691CA89CB600D88C1B /* ASEnvironmentInternal.h */,
 				69E1006A1CA89CB600D88C1B /* ASEnvironmentInternal.mm */,
+				68B8A4DB1CBD911D007E4543 /* ASImageNode+AnimatedImagePrivate.h */,
 				058D0A0D195D050800B7D73C /* ASImageNode+CGExtras.h */,
 				058D0A0E195D050800B7D73C /* ASImageNode+CGExtras.m */,
-				68B8A4DB1CBD911D007E4543 /* ASImageNode+AnimatedImagePrivate.h */,
 				ACF6ED431B17847A00DA7C62 /* ASInternalHelpers.h */,
 				ACF6ED441B17847A00DA7C62 /* ASInternalHelpers.m */,
+				69F6058C1D3DA27E00C8CA38 /* ASLayoutSpec+Private.h */,
 				ACF6ED451B17847A00DA7C62 /* ASLayoutSpecUtilities.h */,
+				E52405B41C8FEF16004DC8E7 /* ASLayoutTransition.h */,
+				E52405B21C8FEF03004DC8E7 /* ASLayoutTransition.mm */,
 				0442850B1BAA64EC00D16268 /* ASMultidimensionalArrayUtils.h */,
 				0442850C1BAA64EC00D16268 /* ASMultidimensionalArrayUtils.mm */,
 				CC3B20811C3F76D600798563 /* ASPendingStateController.h */,
@@ -1505,8 +1510,6 @@
 				CC3B20881C3F7A5400798563 /* ASWeakSet.m */,
 				DBC452D91C5BF64600B16017 /* NSArray+Diffing.h */,
 				DBC452DA1C5BF64600B16017 /* NSArray+Diffing.m */,
-				8B0768B11CE752EC002E1453 /* ASDefaultPlaybackButton.h */,
-				8B0768B21CE752EC002E1453 /* ASDefaultPlaybackButton.m */,
 			);
 			path = Private;
 			sourceTree = "<group>";
@@ -1740,6 +1743,7 @@
 				B350625B1B010F070018CF92 /* ASEqualityHelpers.h in Headers */,
 				680346941CE4052A0009FEB4 /* ASNavigationController.h in Headers */,
 				B350621B1B010EFD0018CF92 /* ASFlowLayoutController.h in Headers */,
+				69F6058D1D3DA27E00C8CA38 /* ASLayoutSpec+Private.h in Headers */,
 				B350621D1B010EFD0018CF92 /* ASHighlightOverlayLayer.h in Headers */,
 				C78F7E2B1BF7809800CDEAFC /* ASTableNode.h in Headers */,
 				AC7A2C181BDE11DF0093FE1A /* ASTableViewInternal.h in Headers */,

AsyncDisplayKit/ASDisplayNode.mm

@@ -2981,6 +2981,15 @@ - (void)setEnvironmentTraitCollection:(ASEnvironmentTraitCollection)environmentT
   }
 }
 
+#pragma mark - NSFastEnumeration
+
+- (NSUInteger)countByEnumeratingWithState:(NSFastEnumerationState *)state
+                                  objects:(id __unsafe_unretained [])stackbuf
+                                    count:(NSUInteger)stackbufLength
+{
+  return [self.children countByEnumeratingWithState:state objects:stackbuf count:stackbufLength];
+}
+
 ASEnvironmentLayoutOptionsForwarding
 ASEnvironmentLayoutExtensibilityForwarding
 

AsyncDisplayKit/ASViewController.mm

@@ -289,8 +289,7 @@ - (void)progagateNewEnvironmentTraitCollection:(ASEnvironmentTraitCollection)env
     self.node.environmentState = environmentState;
     [self.node setNeedsLayout];
 
-    NSArray<id<ASEnvironment>> *children = [self.node children];
-    for (id<ASEnvironment> child in children) {
+    for (id<ASEnvironment> child in self.node) {
       ASEnvironmentStatePropagateDown(child, environmentState.environmentTraitCollection);
     }
   }

AsyncDisplayKit/Details/ASEnvironment.h

@@ -99,7 +99,7 @@ ASDISPLAYNODE_EXTERN_C_END
  * defined in an ASEnvironmentState up and down the ASEnvironment tree. To be able to define how merges of
  * States should happen, specific merge functions can be provided
  */
-@protocol ASEnvironment <NSObject>
+@protocol ASEnvironment <NSObject, NSFastEnumeration>
 
 /// The environment collection of an object which class conforms to the ASEnvironment protocol
 - (ASEnvironmentState)environmentState;
@@ -126,6 +126,7 @@ ASDISPLAYNODE_EXTERN_C_END
 
 /// sets a trait collection on this environment state.
 - (void)setEnvironmentTraitCollection:(ASEnvironmentTraitCollection)environmentTraitCollection;
+
 @end
 
 // ASCollection/TableNodes don't actually have ASCellNodes as subnodes. Because of this we can't rely on display trait

AsyncDisplayKit/Layout/ASLayoutSpec.mm

@@ -8,7 +8,7 @@
 //  of patent rights can be found in the PATENTS file in the same directory.
 //
 
-#import "ASLayoutSpec.h"
+#import "ASLayoutSpec+Private.h"
 
 #import "ASAssert.h"
 #import "ASEnvironmentInternal.h"
@@ -17,16 +17,13 @@
 #import "ASThread.h"
 #import "ASTraitCollection.h"
 
-#import <objc/runtime.h>
-#import <map>
 #import <vector>
 
-typedef std::map<unsigned long, id<ASLayoutable>, std::less<unsigned long>> ASChildMap;
-
 @interface ASLayoutSpec() {
   ASEnvironmentState _environmentState;
   ASDN::RecursiveMutex __instanceLock__;
-  ASChildMap _children;
+  ASChildrenMap _childrenMap;
+  unsigned long _mutations;
 }
 @end
 
@@ -35,6 +32,7 @@ @implementation ASLayoutSpec
 // these dynamic properties all defined in ASLayoutOptionsPrivate.m
 @dynamic spacingAfter, spacingBefore, flexGrow, flexShrink, flexBasis,
          alignSelf, ascender, descender, sizeRange, layoutPosition, layoutableType;
+@synthesize parent = _parent;
 @synthesize isFinalLayoutable = _isFinalLayoutable;
 
 - (instancetype)init
@@ -44,6 +42,7 @@ - (instancetype)init
   }
   _isMutable = YES;
   _environmentState = ASEnvironmentStateMakeDefault();
+  _mutations = 0;
   return self;
 }
 
@@ -105,6 +104,8 @@ - (ASLayout *)measureWithSizeRange:(ASSizeRange)constrainedSize
   return child;
 }
 
+#pragma mark - Parent
+
 - (void)setParent:(id<ASLayoutable>)parent
 {
   // FIXME: Locking should be evaluated here.  _parent is not widely used yet, though.
@@ -115,30 +116,33 @@ - (void)setParent:(id<ASLayoutable>)parent
   }
 }
 
+- (id<ASLayoutable>)parent
+{
+  return _parent;
+}
+
+#pragma mark - Children
+
 - (void)setChild:(id<ASLayoutable>)child
 {
-  ASDisplayNodeAssert(self.isMutable, @"Cannot set properties when layout spec is not mutable");
-  if (child) {
-    id<ASLayoutable> finalLayoutable = [self layoutableToAddFromLayoutable:child];
-    if (finalLayoutable) {
-      _children[0] = finalLayoutable;
-      [self propagateUpLayoutable:finalLayoutable];
-    }
-  } else {
-    _children.erase(0);
-  }
+  [self setChild:child forIndex:0];
 }
 
 - (void)setChild:(id<ASLayoutable>)child forIndex:(NSUInteger)index
 {
   ASDisplayNodeAssert(self.isMutable, @"Cannot set properties when layout spec is not mutable");
   if (child) {
     id<ASLayoutable> finalLayoutable = [self layoutableToAddFromLayoutable:child];
-    _children[index] = finalLayoutable;
+    if (finalLayoutable) {
+      _childrenMap[index] = finalLayoutable;
+      [self propagateUpLayoutable:finalLayoutable];
+    }
   } else {
-    _children.erase(index);
+    _childrenMap.erase(index);
   }
-  // TODO: Should we propagate up the layoutable at it could happen that multiple children will propagated up their
+  _mutations++;
+  
+  // TODO: Should we propagate up the layoutable as it could happen that multiple children will propagated up their
   //       layout options and one child will overwrite values from another child
   // [self propagateUpLayoutable:finalLayoutable];
 }
@@ -147,37 +151,73 @@ - (void)setChildren:(NSArray<id<ASLayoutable>> *)children
 {
   ASDisplayNodeAssert(self.isMutable, @"Cannot set properties when layout spec is not mutable");
 
-  _children.clear();
+  _childrenMap.clear();
   NSUInteger i = 0;
   for (id<ASLayoutable> child in children) {
-    _children[i] = [self layoutableToAddFromLayoutable:child];
+    _childrenMap[i] = [self layoutableToAddFromLayoutable:child];
     i += 1;
+    
+    _mutations++;
   }
 }
 
 - (id<ASLayoutable>)childForIndex:(NSUInteger)index
 {
-  if (index < _children.size()) {
-    return _children[index];
+  if (index < _childrenMap.size()) {
+    return _childrenMap[index];
   }
   return nil;
 }
 
 - (id<ASLayoutable>)child
 {
-  return _children[0];
+  return _childrenMap[0];
 }
 
 - (NSArray *)children
 {
+  // If used inside ASDK, the childrenMap property should be preferred over the children array to prevent
+  // unecessary boxing
   std::vector<ASLayout *> children;
-  for (ASChildMap::iterator it = _children.begin(); it != _children.end(); ++it ) {
-    children.push_back(it->second);
+  for (auto const &entry : _childrenMap) {
+    children.push_back(entry.second);
   }
-  
+
   return [NSArray arrayWithObjects:&children[0] count:children.size()];
 }
 
+#pragma mark - NSFastEnumeration
+
+- (NSUInteger)countByEnumeratingWithState:(NSFastEnumerationState *)state
+                                  objects:(id __unsafe_unretained [])stackbuf
+                                    count:(NSUInteger)stackbufLength
+{
+  NSUInteger count = 0;
+  unsigned long countOfItemsAlreadyEnumerated = state->state;
+  
+  if (countOfItemsAlreadyEnumerated == 0) {
+    state->mutationsPtr = &_mutations;
+  }
+
+  if (countOfItemsAlreadyEnumerated < _childrenMap.size()) {
+    state->itemsPtr = stackbuf;
+        
+    while((countOfItemsAlreadyEnumerated < _childrenMap.size()) && (count < stackbufLength)) {
+      // Hold on for the object while enumerating
+      __autoreleasing id child = _childrenMap[countOfItemsAlreadyEnumerated];
+      stackbuf[count] = child;
+      countOfItemsAlreadyEnumerated++;
+      count++;
+    }
+  } else {
+    count = 0;
+  }
+  
+  state->state = countOfItemsAlreadyEnumerated;
+
+  return count;
+}
+
 #pragma mark - ASEnvironment
 
 - (ASEnvironmentState)environmentState
@@ -233,6 +273,15 @@ - (ASTraitCollection *)asyncTraitCollection
 
 @end
 
+@implementation ASLayoutSpec (Private)
+
+- (ASChildrenMap)childrenMap
+{
+  return _childrenMap;
+}
+
+@end
+
 @implementation ASLayoutSpec (Debugging)
 
 #pragma mark - ASLayoutableAsciiArtProtocol

AsyncDisplayKit/Layout/ASStackLayoutSpec.mm

@@ -13,6 +13,7 @@
 
 #import "ASInternalHelpers.h"
 
+#import "ASLayoutSpec+Private.h"
 #import "ASLayoutSpecUtilities.h"
 #import "ASStackBaselinePositionedLayout.h"
 #import "ASThread.h"
@@ -58,7 +59,7 @@ - (instancetype)initWithDirection:(ASStackLayoutDirection)direction spacing:(CGF
   _alignItems = alignItems;
   _justifyContent = justifyContent;
 
-  [self setChildren:children];
+  self.children = children;
   return self;
 }
 
@@ -120,7 +121,9 @@ - (void)setBaselineRelativeArrangement:(BOOL)baselineRelativeArrangement
 
 - (ASLayout *)measureWithSizeRange:(ASSizeRange)constrainedSize
 {
-  if (self.children.count == 0) {
+  auto const &children = self.childrenMap;
+  
+  if (children.size() == 0) {
     return [ASLayout layoutWithLayoutableObject:self
                            constrainedSizeRange:constrainedSize
                                            size:constrainedSize.min];
@@ -129,9 +132,9 @@ - (ASLayout *)measureWithSizeRange:(ASSizeRange)constrainedSize
   ASStackLayoutSpecStyle style = {.direction = _direction, .spacing = _spacing, .justifyContent = _justifyContent, .alignItems = _alignItems, .baselineRelativeArrangement = _baselineRelativeArrangement};
   BOOL needsBaselinePass = _baselineRelativeArrangement || _alignItems == ASStackLayoutAlignItemsBaselineFirst || _alignItems == ASStackLayoutAlignItemsBaselineLast;
 
-  std::vector<id<ASLayoutable>> stackChildren = std::vector<id<ASLayoutable>>();
-  for (id<ASLayoutable> child in self.children) {
-    stackChildren.push_back(child);
+  std::vector<id<ASLayoutable>> stackChildren;
+  for (auto const &entry : children) {
+    stackChildren.push_back(entry.second);
   }
 
   const auto unpositionedLayout = ASStackUnpositionedLayout::compute(stackChildren, style, constrainedSize);
@@ -143,14 +146,18 @@ - (ASLayout *)measureWithSizeRange:(ASSizeRange)constrainedSize
   // regardless of whether or not this stack aligns to baseline, we should let ASStackBaselinePositionedLayout::compute find the max ascender
   // and min descender in case this spec is a child in another spec that wants to align to a baseline.
   const auto baselinePositionedLayout = ASStackBaselinePositionedLayout::compute(positionedLayout, style, constrainedSize);
-  if (self.direction == ASStackLayoutDirectionVertical) {
-    ASDN::MutexLocker l(__instanceLock__);
-    self.ascender = [[self.children firstObject] ascender];
-    self.descender = [[self.children lastObject] descender];
-  } else {
+  {
     ASDN::MutexLocker l(__instanceLock__);
-    self.ascender = baselinePositionedLayout.ascender;
-    self.descender = baselinePositionedLayout.descender;
+    if (self.direction == ASStackLayoutDirectionVertical) {
+      if (stackChildren.size() > 0) {
+        self.ascender = stackChildren.front().ascender;
+        self.descender = stackChildren.back().descender;
+      }
+
+    } else {
+      self.ascender = baselinePositionedLayout.ascender;
+      self.descender = baselinePositionedLayout.descender;
+    }
   }
 
   if (needsBaselinePass) {

AsyncDisplayKit/Layout/ASStaticLayoutSpec.mm

@@ -10,6 +10,7 @@
 
 #import "ASStaticLayoutSpec.h"
 
+#import "ASLayoutSpec+Private.h"
 #import "ASLayoutSpecUtilities.h"
 #import "ASLayout.h"
 
@@ -38,10 +39,8 @@ - (ASLayout *)measureWithSizeRange:(ASSizeRange)constrainedSize
 {
   CGSize maxConstrainedSize = CGSizeMake(constrainedSize.max.width, constrainedSize.max.height);
 
-  NSArray *children = self.children;
-  NSMutableArray *sublayouts = [NSMutableArray arrayWithCapacity:children.count];
-
-  for (id<ASLayoutable> child in children) {
+  NSMutableArray *sublayouts = [NSMutableArray arrayWithCapacity:self.childrenMap.size()];
+  for (id<ASLayoutable> child in self) {
     CGPoint layoutPosition = child.layoutPosition;
     CGSize autoMaxSize = CGSizeMake(maxConstrainedSize.width  - layoutPosition.x,
                                     maxConstrainedSize.height - layoutPosition.y);