Fix memory leaks, add section-object support to new test harness (#360)

AsyncDisplayKit.xcodeproj/project.pbxproj

@@ -370,8 +370,8 @@
 		CCA282CD1E9EB73E0037E8B7 /* ASTipNode.m in Sources */ = {isa = PBXBuildFile; fileRef = CCA282CB1E9EB73E0037E8B7 /* ASTipNode.m */; };
 		CCA282D01E9EBF6C0037E8B7 /* ASTipsWindow.h in Headers */ = {isa = PBXBuildFile; fileRef = CCA282CE1E9EBF6C0037E8B7 /* ASTipsWindow.h */; };
 		CCA282D11E9EBF6C0037E8B7 /* ASTipsWindow.m in Sources */ = {isa = PBXBuildFile; fileRef = CCA282CF1E9EBF6C0037E8B7 /* ASTipsWindow.m */; };
-		CCA5F62E1EECC2A80060C137 /* ASAssert.m in Sources */ = {isa = PBXBuildFile; fileRef = CCA5F62D1EECC2A80060C137 /* ASAssert.m */; };
 		CCA5F62C1EEC9E9B0060C137 /* NSInvocation+ASTestHelpers.m in Sources */ = {isa = PBXBuildFile; fileRef = CCA5F62B1EEC9E9B0060C137 /* NSInvocation+ASTestHelpers.m */; };
+		CCA5F62E1EECC2A80060C137 /* ASAssert.m in Sources */ = {isa = PBXBuildFile; fileRef = CCA5F62D1EECC2A80060C137 /* ASAssert.m */; };
 		CCB2F34D1D63CCC6004E6DE9 /* ASDisplayNodeSnapshotTests.m in Sources */ = {isa = PBXBuildFile; fileRef = CCB2F34C1D63CCC6004E6DE9 /* ASDisplayNodeSnapshotTests.m */; };
 		CCB338E41EEE11160081F21A /* OCMockObject+ASAdditions.m in Sources */ = {isa = PBXBuildFile; fileRef = CCB338E31EEE11160081F21A /* OCMockObject+ASAdditions.m */; };
 		CCB338E71EEE27760081F21A /* ASTestCase.m in Sources */ = {isa = PBXBuildFile; fileRef = CCB338E61EEE27760081F21A /* ASTestCase.m */; };
@@ -832,9 +832,9 @@
 		CCA282CB1E9EB73E0037E8B7 /* ASTipNode.m */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.objc; path = ASTipNode.m; sourceTree = "<group>"; };
 		CCA282CE1E9EBF6C0037E8B7 /* ASTipsWindow.h */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.h; path = ASTipsWindow.h; sourceTree = "<group>"; };
 		CCA282CF1E9EBF6C0037E8B7 /* ASTipsWindow.m */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.objc; path = ASTipsWindow.m; sourceTree = "<group>"; };
-		CCA5F62D1EECC2A80060C137 /* ASAssert.m */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.objc; path = ASAssert.m; sourceTree = "<group>"; };
 		CCA5F62A1EEC9E9B0060C137 /* NSInvocation+ASTestHelpers.h */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.h; path = "NSInvocation+ASTestHelpers.h"; sourceTree = "<group>"; };
 		CCA5F62B1EEC9E9B0060C137 /* NSInvocation+ASTestHelpers.m */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.objc; path = "NSInvocation+ASTestHelpers.m"; sourceTree = "<group>"; };
+		CCA5F62D1EECC2A80060C137 /* ASAssert.m */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.objc; path = ASAssert.m; sourceTree = "<group>"; };
 		CCB2F34C1D63CCC6004E6DE9 /* ASDisplayNodeSnapshotTests.m */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.objc; path = ASDisplayNodeSnapshotTests.m; sourceTree = "<group>"; };
 		CCB338E21EEE11160081F21A /* OCMockObject+ASAdditions.h */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.h; path = "OCMockObject+ASAdditions.h"; sourceTree = "<group>"; };
 		CCB338E31EEE11160081F21A /* OCMockObject+ASAdditions.m */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.objc; path = "OCMockObject+ASAdditions.m"; sourceTree = "<group>"; };

Source/ASRunLoopQueue.h

@@ -51,6 +51,8 @@ AS_SUBCLASSING_RESTRICTED
 
 + (instancetype)sharedDeallocationQueue;
 
+- (void)test_drain;
+
 - (void)releaseObjectInBackground:(id)object;
 
 @end

Source/ASRunLoopQueue.mm

@@ -143,6 +143,29 @@ - (void)stop
   _thread = nil;
 }
 
+- (void)test_drain
+{
+  [self performSelector:@selector(_test_drain) onThread:_thread withObject:nil waitUntilDone:YES];
+}
+
+- (void)_test_drain
+{
+  while (true) {
+    @autoreleasepool {
+      _queueLock.lock();
+      std::deque<id> currentQueue = _queue;
+      _queue = std::deque<id>();
+      _queueLock.unlock();
+
+      if (currentQueue.empty()) {
+        return;
+      } else {
+        currentQueue.clear();
+      }
+    }
+  }
+}
+
 - (void)_stop
 {
   CFRunLoopStop(CFRunLoopGetCurrent());

Source/Details/ASWeakSet.m

@@ -27,7 +27,7 @@ - (instancetype)init
 {
   self = [super init];
   if (self) {
-    _hashTable = [NSHashTable hashTableWithOptions:NSPointerFunctionsWeakMemory | NSPointerFunctionsObjectPointerPersonality];
+    _hashTable = [NSHashTable hashTableWithOptions:NSHashTableWeakMemory | NSHashTableObjectPointerPersonality];
   }
   return self;
 }

Tests/ASCollectionModernDataSourceTests.m

@@ -23,21 +23,28 @@ @interface ASCollectionModernDataSourceTests : ASTestCase
 @interface ASTestCellNode : ASCellNode
 @end
 
+@interface ASTestSection : NSObject <ASSectionContext>
+@property (nonatomic, readonly) NSMutableArray *viewModels;
+@end
+
 @implementation ASCollectionModernDataSourceTests {
 @private
   id mockDataSource;
   UIWindow *window;
   UIViewController *viewController;
   ASCollectionNode *collectionNode;
-  NSMutableArray<NSMutableArray *> *sections;
+  NSMutableArray<ASTestSection *> *sections;
 }
 
 - (void)setUp {
   [super setUp];
   // Default is 2 sections: 2 items in first, 1 item in second.
   sections = [NSMutableArray array];
-  [sections addObject:[NSMutableArray arrayWithObjects:[NSObject new], [NSObject new], nil]];
-  [sections addObject:[NSMutableArray arrayWithObjects:[NSObject new], nil]];
+  [sections addObject:[ASTestSection new]];
+  [sections[0].viewModels addObject:[NSObject new]];
+  [sections[0].viewModels addObject:[NSObject new]];
+  [sections addObject:[ASTestSection new]];
+  [sections[1].viewModels addObject:[NSObject new]];
   window = [[UIWindow alloc] initWithFrame:CGRectMake(0, 0, 100, 100)];
   viewController = [[UIViewController alloc] init];
 
@@ -54,6 +61,7 @@ - (void)setUp {
    @selector(collectionNode:numberOfItemsInSection:),
    @selector(collectionNode:nodeBlockForItemAtIndexPath:),
    @selector(collectionNode:viewModelForItemAtIndexPath:),
+   @selector(collectionNode:contextForSection:),
    nil];
   [mockDataSource setExpectationOrderMatters:YES];
 
@@ -64,7 +72,6 @@ - (void)setUp {
 - (void)tearDown
 {
   [collectionNode waitUntilAllUpdatesAreCommitted];
-  OCMVerifyAll(mockDataSource);
   [super tearDown];
 }
 
@@ -82,11 +89,12 @@ - (void)testReloadingAnItem
   // Reload at (0, 0)
   NSIndexPath *reloadedPath = [NSIndexPath indexPathForItem:0 inSection:0];
 
-  [self performUpdateReloadingItems:@{ reloadedPath: [NSObject new] }
-                     reloadMappings:@{ reloadedPath: reloadedPath }
-                     insertingItems:nil
-                      deletingItems:nil
-            skippedReloadIndexPaths:nil];
+  [self performUpdateReloadingSections:nil
+                        reloadingItems:@{ reloadedPath: [NSObject new] }
+                        reloadMappings:@{ reloadedPath: reloadedPath }
+                        insertingItems:nil
+                         deletingItems:nil
+               skippedReloadIndexPaths:nil];
 }
 
 - (void)testInsertingAnItem
@@ -96,11 +104,12 @@ - (void)testInsertingAnItem
   // Insert at (1, 0)
   NSIndexPath *insertedPath = [NSIndexPath indexPathForItem:0 inSection:1];
 
-  [self performUpdateReloadingItems:nil
-                     reloadMappings:nil
-                     insertingItems:@{ insertedPath: [NSObject new] }
-                      deletingItems:nil
-            skippedReloadIndexPaths:nil];
+  [self performUpdateReloadingSections:nil
+                        reloadingItems:nil
+                        reloadMappings:nil
+                        insertingItems:@{ insertedPath: [NSObject new] }
+                         deletingItems:nil
+               skippedReloadIndexPaths:nil];
 }
 
 - (void)testReloadingAnItemWithACompatibleViewModel
@@ -115,17 +124,27 @@ - (void)testReloadingAnItemWithACompatibleViewModel
 
   // Cell node should get -canUpdateToViewModel:
   id mockCellNode = [collectionNode nodeForItemAtIndexPath:reloadedPath];
-  [mockCellNode setExpectationOrderMatters:YES];
   OCMExpect([mockCellNode canUpdateToViewModel:viewModel])
   .andReturn(YES);
 
-  [self performUpdateReloadingItems:@{ reloadedPath: viewModel }
-                     reloadMappings:@{ reloadedPath: [NSIndexPath indexPathForItem:0 inSection:0] }
-                     insertingItems:nil
-                      deletingItems:@[ deletedPath ]
-            skippedReloadIndexPaths:@[ reloadedPath ]];
-
-  OCMVerifyAll(mockCellNode);
+  [self performUpdateReloadingSections:nil
+                        reloadingItems:@{ reloadedPath: viewModel }
+                        reloadMappings:@{ reloadedPath: [NSIndexPath indexPathForItem:0 inSection:0] }
+                        insertingItems:nil
+                         deletingItems:@[ deletedPath ]
+               skippedReloadIndexPaths:@[ reloadedPath ]];
+}
+
+- (void)testReloadingASection
+{
+  [self loadInitialData];
+
+  [self performUpdateReloadingSections:@{ @0: [ASTestSection new] }
+                        reloadingItems:nil
+                        reloadMappings:nil
+                        insertingItems:nil
+                         deletingItems:nil
+               skippedReloadIndexPaths:nil];
 }
 
 #pragma mark - Helpers
@@ -137,14 +156,19 @@ - (void)loadInitialData
     // It reads all the counts
     [self expectDataSourceCountMethods];
 
+    // It reads each section object.
+    for (NSInteger section = 0; section < sections.count; section++) {
+      [self expectContextMethodForSection:section];
+    }
+
     // It reads the contents for each item.
     for (NSInteger section = 0; section < sections.count; section++) {
-      NSArray *items = sections[section];
+      NSArray *viewModels = sections[section].viewModels;
 
       // For each item:
-      for (NSInteger i = 0; i < items.count; i++) {
+      for (NSInteger i = 0; i < viewModels.count; i++) {
         NSIndexPath *indexPath = [NSIndexPath indexPathForItem:i inSection:section];
-        [self expectViewModelMethodForItemAtIndexPath:indexPath viewModel:items[i]];
+        [self expectViewModelMethodForItemAtIndexPath:indexPath viewModel:viewModels[i]];
         [self expectNodeBlockMethodForItemAtIndexPath:indexPath];
       }
     }
@@ -172,9 +196,8 @@ - (void)expectDataSourceCountMethods
   // For each section:
   // Note: Skip fast enumeration for readability.
   for (NSInteger section = 0; section < sections.count; section++) {
-    NSInteger itemCount = sections[section].count;
     OCMExpect([mockDataSource collectionNode:collectionNode numberOfItemsInSection:section])
-    .andReturn(itemCount);
+    .andReturn(sections[section].viewModels.count);
   }
 }
 
@@ -184,15 +207,24 @@ - (void)expectViewModelMethodForItemAtIndexPath:(NSIndexPath *)indexPath viewMod
   .andReturn(viewModel);
 }
 
+- (void)expectContextMethodForSection:(NSInteger)section
+{
+  OCMExpect([mockDataSource collectionNode:collectionNode contextForSection:section])
+  .andReturn(sections[section]);
+}
+
 - (void)expectNodeBlockMethodForItemAtIndexPath:(NSIndexPath *)indexPath
 {
   OCMExpect([mockDataSource collectionNode:collectionNode nodeBlockForItemAtIndexPath:indexPath])
   .andReturn((ASCellNodeBlock)^{
     ASCellNode *node = [ASTestCellNode new];
     // Generating multiple partial mocks of the same class is not thread-safe.
+    id mockNode;
     @synchronized (NSNull.null) {
-      return OCMPartialMock(node);
+      mockNode = OCMPartialMock(node);
     }
+    [mockNode setExpectationOrderMatters:YES];
+    return mockNode;
   });
 }
 
@@ -203,14 +235,19 @@ - (void)assertCollectionNodeContent
   XCTAssertEqual(collectionNode.numberOfSections, sections.count);
 
   for (NSInteger section = 0; section < sections.count; section++) {
-    NSArray *items = sections[section];
+    ASTestSection *sectionObject = sections[section];
+    NSArray *viewModels = sectionObject.viewModels;
+
+    // Assert section object
+    XCTAssertEqualObjects([collectionNode contextForSection:section], sectionObject);
+
     // Assert item count
-    XCTAssertEqual([collectionNode numberOfItemsInSection:section], items.count);
-    for (NSInteger item = 0; item < items.count; item++) {
+    XCTAssertEqual([collectionNode numberOfItemsInSection:section], viewModels.count);
+    for (NSInteger item = 0; item < viewModels.count; item++) {
       // Assert view model
       // Could use pointer equality but the error message is less readable.
       NSIndexPath *indexPath = [NSIndexPath indexPathForItem:item inSection:section];
-      id viewModel = sections[indexPath.section][indexPath.item];
+      id viewModel = viewModels[indexPath.item];
       XCTAssertEqualObjects(viewModel, [collectionNode viewModelForItemAtIndexPath:indexPath]);
       ASCellNode *node = [collectionNode nodeForItemAtIndexPath:indexPath];
       XCTAssertEqualObjects(node.viewModel, viewModel);
@@ -224,30 +261,39 @@ - (void)assertCollectionNodeContent
  *
  * skippedReloadIndexPaths are the old index paths for nodes that should use -canUpdateToViewModel: instead of being refetched.
  */
-- (void)performUpdateReloadingItems:(NSDictionary<NSIndexPath *, id> *)reloadedItems
-                     reloadMappings:(NSDictionary<NSIndexPath *, NSIndexPath *> *)reloadMappings
-                     insertingItems:(NSDictionary<NSIndexPath *, id> *)insertedItems
-                      deletingItems:(NSArray<NSIndexPath *> *)deletedItems
-            skippedReloadIndexPaths:(NSArray<NSIndexPath *> *)skippedReloadIndexPaths
+- (void)performUpdateReloadingSections:(NSDictionary<NSNumber *, id> *)reloadedSections
+                        reloadingItems:(NSDictionary<NSIndexPath *, id> *)reloadedItems
+                        reloadMappings:(NSDictionary<NSIndexPath *, NSIndexPath *> *)reloadMappings
+                        insertingItems:(NSDictionary<NSIndexPath *, id> *)insertedItems
+                         deletingItems:(NSArray<NSIndexPath *> *)deletedItems
+               skippedReloadIndexPaths:(NSArray<NSIndexPath *> *)skippedReloadIndexPaths
 {
   [collectionNode performBatchUpdates:^{
     // First update our data source.
     [reloadedItems enumerateKeysAndObjectsUsingBlock:^(NSIndexPath * _Nonnull key, id  _Nonnull obj, BOOL * _Nonnull stop) {
-      sections[key.section][key.item] = obj;
+      sections[key.section].viewModels[key.item] = obj;
+    }];
+    [reloadedSections enumerateKeysAndObjectsUsingBlock:^(NSNumber * _Nonnull key, id  _Nonnull obj, BOOL * _Nonnull stop) {
+      sections[key.integerValue] = obj;
     }];
 
     // Deletion paths, sorted descending
     for (NSIndexPath *indexPath in [deletedItems sortedArrayUsingSelector:@selector(compare:)].reverseObjectEnumerator) {
-      [sections[indexPath.section] removeObjectAtIndex:indexPath.item];
+      [sections[indexPath.section].viewModels removeObjectAtIndex:indexPath.item];
     }
 
     // Insertion paths, sorted ascending.
     NSArray *insertionsSortedAcending = [insertedItems.allKeys sortedArrayUsingSelector:@selector(compare:)];
     for (NSIndexPath *indexPath in insertionsSortedAcending) {
-      [sections[indexPath.section] insertObject:insertedItems[indexPath] atIndex:indexPath.item];
+      [sections[indexPath.section].viewModels insertObject:insertedItems[indexPath] atIndex:indexPath.item];
     }
 
     // Then update the collection node.
+    NSMutableIndexSet *reloadedSectionIndexes = [NSMutableIndexSet indexSet];
+    for (NSNumber *i in reloadedSections) {
+      [reloadedSectionIndexes addIndex:i.integerValue];
+    }
+    [collectionNode reloadSections:reloadedSectionIndexes];
     [collectionNode reloadItemsAtIndexPaths:reloadedItems.allKeys];
     [collectionNode deleteItemsAtIndexPaths:deletedItems];
     [collectionNode insertItemsAtIndexPaths:insertedItems.allKeys];
@@ -260,6 +306,17 @@ - (void)performUpdateReloadingItems:(NSDictionary<NSIndexPath *, id> *)reloadedI
     // Combine reloads + inserts and expect them to load content for all of them, in ascending order.
     NSMutableDictionary<NSIndexPath *, id> *insertsPlusReloads = [NSMutableDictionary dictionary];
     [insertsPlusReloads addEntriesFromDictionary:insertedItems];
+
+    // Go through reloaded sections and add all their items into `insertsPlusReloads`
+    [reloadedSectionIndexes enumerateIndexesUsingBlock:^(NSUInteger section, BOOL * _Nonnull stop) {
+      [self expectContextMethodForSection:section];
+      NSArray *viewModels = sections[section].viewModels;
+      for (NSInteger i = 0; i < viewModels.count; i++) {
+        NSIndexPath *indexPath = [NSIndexPath indexPathForItem:i inSection:section];
+        insertsPlusReloads[indexPath] = viewModels[i];
+      }
+    }];
+
     [reloadedItems enumerateKeysAndObjectsUsingBlock:^(NSIndexPath * _Nonnull key, id  _Nonnull obj, BOOL * _Nonnull stop) {
       insertsPlusReloads[reloadMappings[key]] = obj;
     }];
@@ -280,6 +337,8 @@ - (void)performUpdateReloadingItems:(NSDictionary<NSIndexPath *, id> *)reloadedI
 
 @end
 
+#pragma mark - Other Objects
+
 @implementation ASTestCellNode
 
 - (BOOL)canUpdateToViewModel:(id)viewModel
@@ -289,3 +348,17 @@ - (BOOL)canUpdateToViewModel:(id)viewModel
 }
 
 @end
+
+@implementation ASTestSection
+@synthesize collectionView;
+@synthesize sectionName;
+
+- (instancetype)init
+{
+  if (self = [super init]) {
+    _viewModels = [NSMutableArray array];
+  }
+  return self;
+}
+
+@end

Tests/ASTestCase.h

@@ -12,6 +12,12 @@
 
 #import <XCTest/XCTest.h>
 
+NS_ASSUME_NONNULL_BEGIN
+
 @interface ASTestCase : XCTestCase
 
+@property (class, nonatomic, nullable, readonly) ASTestCase *currentTestCase;
+
 @end
+
+NS_ASSUME_NONNULL_END