From 3efd01a781b60d3f4913960de3d71770ede14782 Mon Sep 17 00:00:00 2001
From: Eric Horacek <horacek.eric@gmail.com>
Date: Sat, 9 Jul 2016 22:46:29 -0700
Subject: [PATCH] [ASCollectionDataController] Repopulate deleted supplementary
 nodes if necessary (#1773)

* Repopulate deleted supplementary nodes if necessary

Fixes 1771

* Fix warning
---
 .../Details/ASCollectionDataController.mm     | 21 +++++++++++++++++++
 1 file changed, 21 insertions(+)

diff --git a/AsyncDisplayKit/Details/ASCollectionDataController.mm b/AsyncDisplayKit/Details/ASCollectionDataController.mm
index 87552536..2a0978cb 100644
--- a/AsyncDisplayKit/Details/ASCollectionDataController.mm
+++ b/AsyncDisplayKit/Details/ASCollectionDataController.mm
@@ -154,11 +154,32 @@
   }
 }
 
+- (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];
   }
 }