Fix stacked section controller OOB when cell ends display

Summary:
Fixes bug reported internally. When items are removed dynamically the stack internal store will attempt to access data that has already been removed. Instead use assoc objects.

We did change `IGListAdapter` to [use a map](https://github.com/Instagram/IGListKit/blob/master/Source/IGListAdapter.m#L681) instead of assoc objects. That could be a good cleanup.

cc cdoncarroll

- [x] All tests pass. Demo project builds and runs.
- [x] I added tests, an experiment, or detailed why my change isn't tested.
- [x] I have reviewed the [contributing guide](https://github.com/Instagram/IGListKit/blob/master/.github/CONTRIBUTING.md)
Closes https://github.com/Instagram/IGListKit/pull/358

Differential Revision: D4363840

Pulled By: jessesquires

fbshipit-source-id: ef73b4302f88a15cbf70378421d702f7e2bddbd5

CHANGELOG.md

@@ -21,6 +21,8 @@ This release closes the [2.1.0 milestone](https://github.com/Instagram/IGListKit
 
 - Prevent adapter data source from deallocating after queueing an update. [Ryan Nystrom](https://github.com/rnystrom) (tbd)
 
+- Fix out-of-bounds bug when child section controllers in a stack remove cells. [Ryan Nystrom](https://github.com/rnystrom) [(#358)](https://github.com/Instagram/IGListKit/pull/358)
+
 2.0.0
 -----
 

Source/IGListStackedSectionController.m

@@ -16,6 +16,32 @@
 
 #import "IGListSectionControllerInternal.h"
 
+@interface UICollectionViewCell (IGListStackedSectionController)
+@end
+@implementation UICollectionViewCell (IGListStackedSectionController)
+
+static void * kStackedSectionControllerKey = &kStackedSectionControllerKey;
+
+- (void)ig_setStackedSectionController:(id)stackedSectionController {
+    objc_setAssociatedObject(self, kStackedSectionControllerKey, stackedSectionController, OBJC_ASSOCIATION_RETAIN_NONATOMIC);
+}
+
+- (id)ig_stackedSectionController {
+    return objc_getAssociatedObject(self, kStackedSectionControllerKey);
+}
+
+static void * kStackedSectionControllerIndexKey = &kStackedSectionControllerIndexKey;
+
+- (void)ig_setStackedSectionControllerIndex:(NSInteger)stackedSectionControllerIndex {
+    objc_setAssociatedObject(self, kStackedSectionControllerIndexKey, @(stackedSectionControllerIndex), OBJC_ASSOCIATION_ASSIGN);
+}
+
+- (NSInteger)ig_stackedSectionControllerIndex {
+    return [objc_getAssociatedObject(self, kStackedSectionControllerIndexKey) integerValue];
+}
+
+@end
+
 @implementation IGListStackedSectionController
 
 - (instancetype)initWithSectionControllers:(NSArray <IGListSectionController<IGListSectionType> *> *)sectionControllers {
@@ -284,6 +310,10 @@
     IGListSectionController<IGListSectionType> *childSectionController = [self sectionControllerForObjectIndex:index];
     const NSUInteger localIndex = [self localIndexForSectionController:childSectionController index:index];
 
+    // update the assoc objects for use in didEndDisplay
+    [cell ig_setStackedSectionController:childSectionController];
+    [cell ig_setStackedSectionControllerIndex:localIndex];
+
     NSCountedSet *visibleSectionControllers = self.visibleSectionControllers;
     id<IGListDisplayDelegate> displayDelegate = [childSectionController displayDelegate];
 
@@ -296,8 +326,9 @@
 }
 
 - (void)listAdapter:(IGListAdapter *)listAdapter didEndDisplayingSectionController:(IGListSectionController<IGListSectionType> *)sectionController cell:(UICollectionViewCell *)cell atIndex:(NSInteger)index {
-    IGListSectionController<IGListSectionType> *childSectionController = [self sectionControllerForObjectIndex:index];
-    const NSUInteger localIndex = [self localIndexForSectionController:childSectionController index:index];
+    const NSUInteger localIndex = [cell ig_stackedSectionControllerIndex];
+    IGListSectionController<IGListSectionType> *childSectionController = [cell ig_stackedSectionController];
+
     NSCountedSet *visibleSectionControllers = self.visibleSectionControllers;
     id<IGListDisplayDelegate> displayDelegate = [childSectionController displayDelegate];
 

Tests/IGListStackSectionControllerTests.m

@@ -702,4 +702,27 @@ static const CGRect kStackTestFrame = (CGRect){{0.0, 0.0}, {100.0, 100.0}};
     XCTAssertFalse([[self.collectionView cellForItemAtIndexPath:path] isSelected]);
 }
 
+- (void)test_whenRemovingCellsFromChild_thatStackSendsDisplayEventsCorrectly {
+    IGTestObject *object = [[IGTestObject alloc] initWithKey:@0 value:@[@1, @2]];
+    [self setupWithObjects:@[object]];
+
+    IGListStackedSectionController *stack = [self.adapter sectionControllerForObject:object];
+    IGListTestSection *section = stack.sectionControllers.lastObject;
+
+    XCTAssertEqual([self.collectionView numberOfSections], 1);
+    XCTAssertEqual([self.collectionView numberOfItemsInSection:0], 3);
+
+    XCTestExpectation *expectation = [self expectationWithDescription:NSStringFromSelector(_cmd)];
+    [section.collectionContext performBatchAnimated:YES updates:^{
+        section.items = 1;
+        [section.collectionContext deleteInSectionController:section atIndexes:[NSIndexSet indexSetWithIndex:1]];
+    } completion:^(BOOL finished) {
+        XCTAssertEqual([self.collectionView numberOfSections], 1);
+        XCTAssertEqual([self.collectionView numberOfItemsInSection:0], 2);
+        [expectation fulfill];
+    }];
+
+    [self waitForExpectationsWithTimeout:15 handler:nil];
+}
+
 @end