Fix issue where supplementary elements don't track section changes (#404)

* Fix collection view supplementary issue

* Add a fast-path

* Remove the old index path entry unconditionally

* Handle overwriting bug

AsyncDisplayKit.xcodeproj/project.pbxproj

@@ -325,7 +325,7 @@
 		CC0F886D1E4286FA00576FED /* ReferenceImages_iOS_10 in Resources */ = {isa = PBXBuildFile; fileRef = CC0F886A1E4286FA00576FED /* ReferenceImages_iOS_10 */; };
 		CC11F97A1DB181180024D77B /* ASNetworkImageNodeTests.m in Sources */ = {isa = PBXBuildFile; fileRef = CC11F9791DB181180024D77B /* ASNetworkImageNodeTests.m */; };
 		CC2F65EE1E5FFB1600DA57C9 /* ASMutableElementMap.h in Headers */ = {isa = PBXBuildFile; fileRef = CC2F65EC1E5FFB1600DA57C9 /* ASMutableElementMap.h */; };
-		CC2F65EF1E5FFB1600DA57C9 /* ASMutableElementMap.m in Sources */ = {isa = PBXBuildFile; fileRef = CC2F65ED1E5FFB1600DA57C9 /* ASMutableElementMap.m */; };
+		CC2F65EF1E5FFB1600DA57C9 /* ASMutableElementMap.mm in Sources */ = {isa = PBXBuildFile; fileRef = CC2F65ED1E5FFB1600DA57C9 /* ASMutableElementMap.mm */; };
 		CC3B20841C3F76D600798563 /* ASPendingStateController.h in Headers */ = {isa = PBXBuildFile; fileRef = CC3B20811C3F76D600798563 /* ASPendingStateController.h */; settings = {ATTRIBUTES = (Private, ); }; };
 		CC3B20861C3F76D600798563 /* ASPendingStateController.mm in Sources */ = {isa = PBXBuildFile; fileRef = CC3B20821C3F76D600798563 /* ASPendingStateController.mm */; };
 		CC3B208A1C3F7A5400798563 /* ASWeakSet.h in Headers */ = {isa = PBXBuildFile; fileRef = CC3B20871C3F7A5400798563 /* ASWeakSet.h */; settings = {ATTRIBUTES = (Public, ); }; };
@@ -784,7 +784,7 @@
 		CC11F9791DB181180024D77B /* ASNetworkImageNodeTests.m */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.objc; path = ASNetworkImageNodeTests.m; sourceTree = "<group>"; };
 		CC2E317F1DAC353700EEE891 /* ASCollectionView+Undeprecated.h */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.h; path = "ASCollectionView+Undeprecated.h"; sourceTree = "<group>"; };
 		CC2F65EC1E5FFB1600DA57C9 /* ASMutableElementMap.h */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.h; path = ASMutableElementMap.h; sourceTree = "<group>"; };
-		CC2F65ED1E5FFB1600DA57C9 /* ASMutableElementMap.m */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.objc; path = ASMutableElementMap.m; sourceTree = "<group>"; };
+		CC2F65ED1E5FFB1600DA57C9 /* ASMutableElementMap.mm */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.cpp.objcpp; path = ASMutableElementMap.mm; sourceTree = "<group>"; };
 		CC3B20811C3F76D600798563 /* ASPendingStateController.h */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.h; path = ASPendingStateController.h; sourceTree = "<group>"; };
 		CC3B20821C3F76D600798563 /* ASPendingStateController.mm */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.cpp.objcpp; path = ASPendingStateController.mm; sourceTree = "<group>"; };
 		CC3B20871C3F7A5400798563 /* ASWeakSet.h */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.h; path = ASWeakSet.h; sourceTree = "<group>"; };
@@ -1314,7 +1314,7 @@
 				CCA282B21E9EA7310037E8B7 /* ASTipsController.h */,
 				CCA282B31E9EA7310037E8B7 /* ASTipsController.m */,
 				CC2F65EC1E5FFB1600DA57C9 /* ASMutableElementMap.h */,
-				CC2F65ED1E5FFB1600DA57C9 /* ASMutableElementMap.m */,
+				CC2F65ED1E5FFB1600DA57C9 /* ASMutableElementMap.mm */,
 				E5ABAC791E8564EE007AC15C /* ASRectTable.h */,
 				E5ABAC7A1E8564EE007AC15C /* ASRectTable.m */,
 				CC55A70F1E52A0F200594372 /* ASResponderChainEnumerator.h */,
@@ -2203,7 +2203,7 @@
 				B350621E1B010EFD0018CF92 /* ASHighlightOverlayLayer.mm in Sources */,
 				9CC606651D24DF9E006581A0 /* NSIndexSet+ASHelpers.m in Sources */,
 				CC0F885F1E4280B800576FED /* _ASCollectionViewCell.m in Sources */,
-				CC2F65EF1E5FFB1600DA57C9 /* ASMutableElementMap.m in Sources */,
+				CC2F65EF1E5FFB1600DA57C9 /* ASMutableElementMap.mm in Sources */,
 				B35062541B010EFD0018CF92 /* ASImageNode+CGExtras.m in Sources */,
 				E58E9E4A1E941DA5004CFC59 /* ASCollectionLayout.mm in Sources */,
 				6947B0C01E36B4E30007C478 /* ASStackUnpositionedLayout.mm in Sources */,

CHANGELOG.md

@@ -13,6 +13,7 @@
 - Migrated unit tests to OCMock 3.4 (from 2.2) and improved the multiplex image node tests. [Adlai Holler](https://github.com/Adlai-Holler)
 - Fix CollectionNode double-load issue. This should significantly improve performance in cases where a collection node has content immediately available on first layout i.e. not fetched from the network. [Adlai Holler](https://github.com/Adlai-Holler)
 - Overhaul layout flattening algorithm [Huy Nguyen](https://github.com/nguyenhuy) [#395](https://github.com/TextureGroup/Texture/pull/395).
+- Fix an issue where inserting/deleting sections could lead to inconsistent supplementary element behavior. [Adlai Holler](https://github.com/Adlai-Holler)
 
 ## 2.3.3
 - [ASTextKitFontSizeAdjuster] Replace use of NSAttributedString's boundingRectWithSize:options:context: with NSLayoutManager's boundingRectForGlyphRange:inTextContainer: [Ricky Cancro](https://github.com/rcancro)

Source/Details/ASDataController.mm

@@ -298,7 +298,6 @@ typedef void (^ASDataControllerCompletionBlock)(NSArray<ASCollectionElement *> *
 
   // Remove all old supplementaries from these sections
   NSIndexSet *oldSections = [NSIndexSet as_sectionsFromIndexPaths:indexPaths];
-  [map removeSupplementaryElementsInSections:oldSections];
 
   // Add in new ones with the new kinds.
   NSIndexSet *newSections;
@@ -675,6 +674,9 @@ typedef void (^ASDataControllerCompletionBlock)(NSArray<ASCollectionElement *> *
     return;
   }
   
+  // Migrate old supplementary nodes to their new index paths.
+  [map migrateSupplementaryElementsWithChangeSet:changeSet];
+
   for (_ASHierarchyItemChange *change in [changeSet itemChangesOfType:_ASHierarchyChangeTypeDelete]) {
     [map removeItemsAtIndexPaths:change.indexPaths];
     // Aggressively repopulate supplementary nodes (#1773 & #1629)
@@ -688,7 +690,6 @@ typedef void (^ASDataControllerCompletionBlock)(NSArray<ASCollectionElement *> *
 
   for (_ASHierarchySectionChange *change in [changeSet sectionChangesOfType:_ASHierarchyChangeTypeDelete]) {
     NSIndexSet *sectionIndexes = change.indexSet;
-    [map removeSupplementaryElementsInSections:sectionIndexes];
     [map removeSectionsOfItems:sectionIndexes];
   }
   

Source/Private/ASMutableElementMap.h

@@ -21,7 +21,7 @@
 
 NS_ASSUME_NONNULL_BEGIN
 
-@class ASSection, ASCollectionElement;
+@class ASSection, ASCollectionElement, _ASHierarchyChangeSet;
 
 /**
  * This mutable version will be removed in the future. It's only here now to keep the diff small
@@ -47,12 +47,18 @@ AS_SUBCLASSING_RESTRICTED
 
 - (void)removeSectionsOfItems:(NSIndexSet *)itemSections;
 
-- (void)removeSupplementaryElementsInSections:(NSIndexSet *)sections;
-
 - (void)insertEmptySectionsOfItemsAtIndexes:(NSIndexSet *)sections;
 
 - (void)insertElement:(ASCollectionElement *)element atIndexPath:(NSIndexPath *)indexPath;
 
+/**
+ * Update the index paths for all supplementary elements to account for section-level
+ * deletes, moves, inserts. This must be called before adding new supplementary elements.
+ *
+ * This also deletes any supplementary elements in deleted sections.
+ */
+- (void)migrateSupplementaryElementsWithChangeSet:(_ASHierarchyChangeSet *)changeSet;
+
 @end
 
 @interface ASElementMap (MutableCopying) <NSMutableCopying>

Source/Private/ASMutableElementMap.mm

@@ -1,5 +1,5 @@
 //
-//  ASMutableElementMap.m
+//  ASMutableElementMap.mm
 //  Texture
 //
 //  Copyright (c) 2014-present, Facebook, Inc.  All rights reserved.
@@ -22,6 +22,7 @@
 #import <AsyncDisplayKit/ASElementMap.h>
 #import <AsyncDisplayKit/ASTwoDimensionalArrayUtils.h>
 #import <AsyncDisplayKit/NSIndexSet+ASHelpers.h>
+#import <AsyncDisplayKit/_ASHierarchyChangeSet.h>
 
 typedef NSMutableArray<NSMutableArray<ASCollectionElement *> *> ASMutableCollectionElementTwoDimensionalArray;
 
@@ -79,13 +80,6 @@ typedef NSMutableDictionary<NSString *, NSMutableDictionary<NSIndexPath *, ASCol
   [_sectionsOfItems removeObjectsAtIndexes:itemSections];
 }
 
-- (void)removeSupplementaryElementsInSections:(NSIndexSet *)sections
-{
-  [_supplementaryElements enumerateKeysAndObjectsUsingBlock:^(NSString * _Nonnull key, NSMutableDictionary<NSIndexPath *,ASCollectionElement *> * _Nonnull supplementariesForKind, BOOL * _Nonnull stop) {
-    [supplementariesForKind removeObjectsForKeys:[sections as_filterIndexPathsBySection:supplementariesForKind]];
-  }];
-}
-
 - (void)insertEmptySectionsOfItemsAtIndexes:(NSIndexSet *)sections
 {
   [sections enumerateIndexesUsingBlock:^(NSUInteger idx, BOOL * _Nonnull stop) {
@@ -108,6 +102,37 @@ typedef NSMutableDictionary<NSString *, NSMutableDictionary<NSIndexPath *, ASCol
   }
 }
 
+- (void)migrateSupplementaryElementsWithChangeSet:(_ASHierarchyChangeSet *)changeSet
+{
+  if (changeSet.deletedSections.count == 0 && changeSet.insertedSections.count == 0) {
+    return;
+  }
+
+  // For each element kind,
+  [_supplementaryElements enumerateKeysAndObjectsUsingBlock:^(NSString * _Nonnull key, NSMutableDictionary<NSIndexPath *,ASCollectionElement *> * _Nonnull supps, BOOL * _Nonnull stop) {
+    
+    // For each index path of that kind, move entries into a new dictionary.
+    // Note: it's tempting to update the dictionary in-place but because of the likely collision between old and new index paths,
+    // subtle bugs are possible. Note that this process is rare (only on section-level updates),
+    // that this work is done off-main, and that the typical supplementary element use case is just 1-per-section (header).
+    NSMutableDictionary *newSupps = [NSMutableDictionary dictionary];
+    [supps enumerateKeysAndObjectsUsingBlock:^(NSIndexPath * _Nonnull oldIndexPath, ASCollectionElement * _Nonnull obj, BOOL * _Nonnull stop) {
+      NSInteger oldSection = oldIndexPath.section;
+      NSInteger newSection = [changeSet newSectionForOldSection:oldSection];
+      
+      if (oldSection == newSection) {
+        // Index path stayed the same, just copy it over.
+        newSupps[oldIndexPath] = obj;
+      } else if (newSection != NSNotFound) {
+        // Section index changed, move it.
+        NSIndexPath *newIndexPath = [NSIndexPath indexPathForItem:oldIndexPath.item inSection:newSection];
+        newSupps[newIndexPath] = obj;
+      }
+    }];
+    [supps setDictionary:newSupps];
+  }];
+}
+
 #pragma mark - Helpers
 
 + (ASMutableSupplementaryElementDictionary *)deepMutableCopyOfElementsDictionary:(ASSupplementaryElementDictionary *)originalDict