[ASCollectionDataController] Repopulate deleted supplementary nodes if necessary (#1773)

* Repopulate deleted supplementary nodes if necessary

Fixes 1771

* Fix warning

Author: erichoracek

AsyncDisplayKit/Details/ASCollectionDataController.mm

@@ -154,11 +154,32 @@ - (void)willInsertRowsAtIndexPaths:(NSArray<NSIndexPath *> *)indexPaths
   }
 }
 
+- (void)prepareForDeleteRowsAtIndexPaths:(NSArray<NSIndexPath *> *)indexPaths
+{
+  for (NSString *kind in [self supplementaryKinds]) {
+    NSMutableArray<ASIndexedNodeContext *> *contexts = [NSMutableArray array];
+    [self _populateSupplementaryNodesOfKind:kind atIndexPaths:indexPaths mutableContexts:contexts];
+    _pendingContexts[kind] = contexts;
+  }
+}
+
 - (void)willDeleteRowsAtIndexPaths:(NSArray<NSIndexPath *> *)indexPaths
 {
   for (NSString *kind in [self supplementaryKinds]) {
     NSArray<NSIndexPath *> *deletedIndexPaths = ASIndexPathsInMultidimensionalArrayIntersectingIndexPaths([self editingNodesOfKind:kind], indexPaths);
+
     [self deleteNodesOfKind:kind atIndexPaths:deletedIndexPaths completion:nil];
+
+    // If any of the contexts remain after the deletion, re-insert them, e.g.
+    // UICollectionElementKindSectionHeader remains even if item 0 is deleted.
+    NSArray<ASIndexedNodeContext *> *contexts = [_pendingContexts[kind] filteredArrayUsingPredicate:[NSPredicate predicateWithBlock:^(ASIndexedNodeContext *context, NSDictionary *_) {
+      return [deletedIndexPaths containsObject:context.indexPath];
+    }]];
+
+    [self batchLayoutNodesFromContexts:contexts ofKind:kind completion:^(NSArray<ASCellNode *> *nodes, NSArray<NSIndexPath *> *indexPaths) {
+      [self insertNodes:nodes ofKind:kind atIndexPaths:indexPaths completion:nil];
+    }];
+    [_pendingContexts removeObjectForKey:kind];
   }
 }