Implement simple, in-order add/remove subnode support when changing layout specs

AsyncDisplayKit/ASDisplayNode.mm

@@ -30,6 +30,36 @@
 
 NSInteger const ASDefaultDrawingPriority = ASDefaultTransactionPriority;
 
+@interface _ASDisplayNodePosition : NSObject
+
+@property (nonatomic, assign) NSUInteger index;
+@property (nonatomic, strong) ASDisplayNode *node;
+
++ (instancetype)positionWithNode:(ASDisplayNode *)node atIndex:(NSUInteger)index;
+
+- (instancetype)initWithNode:(ASDisplayNode *)node atIndex:(NSUInteger)index;
+
+@end
+
+@implementation _ASDisplayNodePosition
+
++ (instancetype)positionWithNode:(ASDisplayNode *)node atIndex:(NSUInteger)index
+{
+  return [[self alloc] initWithNode:node atIndex:index];
+}
+
+- (instancetype)initWithNode:(ASDisplayNode *)node atIndex:(NSUInteger)index
+{
+  self = [super init];
+  if (self) {
+    _node = node;
+    _index = index;
+  }
+  return self;
+}
+
+@end
+
 @interface ASDisplayNode () <UIGestureRecognizerDelegate>
 
 /**
@@ -609,18 +639,20 @@ static ASDisplayNodeMethodOverrides GetASDisplayNodeMethodOverrides(Class c)
   //  - the constrained size range is different
   if (!_flags.isMeasured || !ASSizeRangeEqualToSizeRange(constrainedSize, _constrainedSize)) {
     ASLayout *newLayout = [self calculateLayoutThatFits:constrainedSize];
-    
+
     if (_layout) {
       NSIndexSet *insertions, *deletions;
-      [_layout.sublayouts asdk_diffWithArray:newLayout.sublayouts insertions:&insertions deletions:&deletions];
-      _insertedSubnodes = [self _filterLayouts:newLayout.sublayouts withIndexes:insertions];
-      _deletedSubnodes = [self _filterLayouts:newLayout.sublayouts withIndexes:deletions];
+      [_layout.sublayouts asdk_diffWithArray:newLayout.sublayouts insertions:&insertions deletions:&deletions compareBlock:^BOOL(ASLayout *lhs, ASLayout *rhs) {
+        return ASObjectIsEqual(lhs.layoutableObject, rhs.layoutableObject);
+      }];
+      _insertedSubnodes = [self _filterSublayouts:newLayout.sublayouts withIndexes:insertions];
+      _deletedSubnodes = [self _filterSublayouts:_layout.sublayouts withIndexes:deletions];
     } else {
       NSIndexSet *indexes = [NSIndexSet indexSetWithIndexesInRange:NSMakeRange(0, [newLayout.sublayouts count])];
-      _insertedSubnodes = [self _filterLayouts:newLayout.sublayouts withIndexes:indexes];
+      _insertedSubnodes = [self _filterSublayouts:newLayout.sublayouts withIndexes:indexes];
       _deletedSubnodes = @[];
     }
-    
+
     _layout = newLayout;
     _constrainedSize = constrainedSize;
     _flags.isMeasured = YES;
@@ -641,13 +673,13 @@ static ASDisplayNodeMethodOverrides GetASDisplayNodeMethodOverrides(Class c)
   return _layout;
 }
 
-- (NSArray<ASDisplayNode *> *)_filterLayouts:(NSArray<ASLayout *> *)layouts withIndexes:(NSIndexSet *)indexes
+- (NSArray<_ASDisplayNodePosition *> *)_filterSublayouts:(NSArray<ASLayout *> *)layouts withIndexes:(NSIndexSet *)indexes
 {
-  NSMutableArray<ASDisplayNode *> *result = [NSMutableArray array];
+  NSMutableArray<_ASDisplayNodePosition *> *result = [NSMutableArray array];
   [indexes enumerateIndexesUsingBlock:^(NSUInteger idx, BOOL * _Nonnull stop) {
     ASDisplayNode *node = (ASDisplayNode *)layouts[idx].layoutableObject;
     ASDisplayNodeAssertNotNil(node, @"A flattened layout must consist exclusively of node sublayouts");
-    [result addObject:node];
+    [result addObject:[_ASDisplayNodePosition positionWithNode:node atIndex:idx]];
   }];
   return result;
 }
@@ -2051,11 +2083,38 @@ static BOOL ShouldUseNewRenderingRange = YES;
     [subnode setFrame:subnodeFrame];
   }
   
-  for (ASDisplayNode *node in _insertedSubnodes) {
-    [self addSubnode:node];
+  if ([[self class] usesImplicitHierarchyManagement]) {
+    if (!_managedSubnodes) {
+      _managedSubnodes = [NSMutableArray array];
+    }
+    
+    for (_ASDisplayNodePosition *position in _deletedSubnodes) {
+      [self _implicitlyRemoveSubnode:position.node atIndex:position.index];
+    }
+    
+    for (_ASDisplayNodePosition *position in _insertedSubnodes) {
+      [self _implicitlyInsertSubnode:position.node atIndex:position.index];
+    }
   }
 }
 
+- (void)_implicitlyInsertSubnode:(ASDisplayNode *)node atIndex:(NSUInteger)idx
+{
+  ASDisplayNodeAssert(idx <= [_managedSubnodes count], @"index needs to be in range of the current managed subnodes");
+  if (idx == [_managedSubnodes count]) {
+    [_managedSubnodes addObject:node];
+  } else {
+    [_managedSubnodes insertObject:node atIndex:idx];
+  }
+  [self addSubnode:node];
+}
+
+- (void)_implicitlyRemoveSubnode:(ASDisplayNode *)node atIndex:(NSUInteger)idx
+{
+  [_managedSubnodes removeObjectAtIndex:idx];
+  [node removeFromSupernode];
+}
+
 - (void)displayWillStart
 {
   // in case current node takes longer to display than it's subnodes, treat it as a dependent node

AsyncDisplayKit/Private/ASDisplayNodeInternal.h

@@ -35,6 +35,7 @@ typedef NS_OPTIONS(NSUInteger, ASDisplayNodeMethodOverrides)
 };
 
 @class _ASPendingState;
+@class _ASDisplayNodePosition;
 
 // Allow 2^n increments of begin disabling hierarchy notifications
 #define VISIBILITY_NOTIFICATIONS_DISABLED_BITS 4
@@ -59,8 +60,12 @@ typedef NS_OPTIONS(NSUInteger, ASDisplayNodeMethodOverrides)
   ASSizeRange _constrainedSize;
   UIEdgeInsets _hitTestSlop;
   NSMutableArray *_subnodes;
-  NSArray<ASDisplayNode *> *_insertedSubnodes;
-  NSArray<ASDisplayNode *> *_deletedSubnodes;
+
+  // Subnodes implicitly managed by layout changes
+  NSMutableArray<ASDisplayNode *> *_managedSubnodes;
+  
+  NSArray<_ASDisplayNodePosition *> *_insertedSubnodes;
+  NSArray<_ASDisplayNodePosition *> *_deletedSubnodes;
 
   ASDisplayNodeViewBlock _viewBlock;
   ASDisplayNodeLayerBlock _layerBlock;

AsyncDisplayKitTests/ASDisplayNodeImplicitHierarchyTests.m

@@ -11,19 +11,35 @@
 #import "ASDisplayNode.h"
 #import "ASDisplayNode+Beta.h"
 #import "ASDisplayNode+Subclasses.h"
+
 #import "ASStaticLayoutSpec.h"
+#import "ASStackLayoutSpec.h"
 
 @interface ASSpecTestDisplayNode : ASDisplayNode
 
-@property (copy, nonatomic) ASLayoutSpec * (^layoutSpecBlock)(ASSizeRange constrainedSize);
+@property (copy, nonatomic) ASLayoutSpec * (^layoutSpecBlock)(ASSizeRange constrainedSize, NSNumber *layoutState);
+
+/**
+ Simple state identifier to allow control of current spec inside of the layoutSpecBlock
+ */
+@property (strong, nonatomic) NSNumber *layoutState;
 
 @end
 
 @implementation ASSpecTestDisplayNode
 
+- (instancetype)init
+{
+  self = [super init];
+  if (self) {
+    _layoutState = @1;
+  }
+  return self;
+}
+
 - (ASLayoutSpec *)layoutSpecThatFits:(ASSizeRange)constrainedSize
 {
-  return self.layoutSpecBlock(constrainedSize);
+  return self.layoutSpecBlock(constrainedSize, _layoutState);
 }
 
 @end
@@ -49,18 +65,65 @@
   XCTAssert([ASDisplayNode usesImplicitHierarchyManagement]);
 }
 
-- (void)testInitialNodeInsertion
+- (void)testInitialNodeInsertionWithOrdering
 {
   ASDisplayNode *node1 = [[ASDisplayNode alloc] init];
   ASDisplayNode *node2 = [[ASDisplayNode alloc] init];
+  ASDisplayNode *node3 = [[ASDisplayNode alloc] init];
+  ASDisplayNode *node4 = [[ASDisplayNode alloc] init];
+  ASDisplayNode *node5 = [[ASDisplayNode alloc] init];
+
   ASSpecTestDisplayNode *node = [[ASSpecTestDisplayNode alloc] init];
-  node.layoutSpecBlock = ^(ASSizeRange constrainedSize){
-    return [ASStaticLayoutSpec staticLayoutSpecWithChildren:@[node1, node2]];
+  node.layoutSpecBlock = ^(ASSizeRange constrainedSize, NSNumber *layoutState) {
+    ASStaticLayoutSpec *staticLayout = [ASStaticLayoutSpec staticLayoutSpecWithChildren:@[node4]];
+    
+    ASStackLayoutSpec *stack1 = [[ASStackLayoutSpec alloc] init];
+    [stack1 setChildren:@[node1, node2]];
+
+    ASStackLayoutSpec *stack2 = [[ASStackLayoutSpec alloc] init];
+    [stack2 setChildren:@[node3, staticLayout]];
+    
+    return [ASStaticLayoutSpec staticLayoutSpecWithChildren:@[stack1, stack2, node5]];
   };
+  [node measureWithSizeRange:ASSizeRangeMake(CGSizeZero, CGSizeZero)];
+  [node layout]; // Layout immediately
+  XCTAssertEqual(node.subnodes[0], node5);
+  XCTAssertEqual(node.subnodes[1], node1);
+  XCTAssertEqual(node.subnodes[2], node2);
+  XCTAssertEqual(node.subnodes[3], node3);
+  XCTAssertEqual(node.subnodes[4], node4);
+}
+
+- (void)testCalculatedLayoutHierarchyTransitions
+{
+  ASDisplayNode *node1 = [[ASDisplayNode alloc] init];
+  ASDisplayNode *node2 = [[ASDisplayNode alloc] init];
+  ASDisplayNode *node3 = [[ASDisplayNode alloc] init];
+  
+  ASSpecTestDisplayNode *node = [[ASSpecTestDisplayNode alloc] init];
+  node.layoutSpecBlock = ^(ASSizeRange constrainedSize, NSNumber *layoutState){
+    if ([layoutState isEqualToNumber:@1]) {
+      return [ASStaticLayoutSpec staticLayoutSpecWithChildren:@[node1, node2]];
+    } else {
+      ASStackLayoutSpec *stackLayout = [[ASStackLayoutSpec alloc] init];
+      [stackLayout setChildren:@[node3, node2]];
+      return [ASStaticLayoutSpec staticLayoutSpecWithChildren:@[node1, stackLayout]];
+    }
+  };
+  
   [node measureWithSizeRange:ASSizeRangeMake(CGSizeZero, CGSizeZero)];
   [node layout]; // Layout immediately
   XCTAssertEqual(node.subnodes[0], node1);
   XCTAssertEqual(node.subnodes[1], node2);
+  
+  node.layoutState = @2;
+  [node invalidateCalculatedLayout]; // TODO(levi): Look into a way where measureWithSizeRange resizes when a new hierarchy is introduced but the size has not changed
+  [node measureWithSizeRange:ASSizeRangeMake(CGSizeZero, CGSizeZero)];
+  [node layout]; // Layout immediately
+
+  XCTAssertEqual(node.subnodes[0], node1);
+  XCTAssertEqual(node.subnodes[1], node3);
+  XCTAssertEqual(node.subnodes[2], node2);
 }
 
 @end