From 3911fb1c34fa20b0f2bf038774be21fb8efc0f2d Mon Sep 17 00:00:00 2001
From: Bin Liu <binl@pinterest.com>
Date: Thu, 11 Feb 2016 14:11:10 -0800
Subject: [PATCH] Adjust reloads' indexPaths in ChangeSet

---
 .../Private/_ASHierarchyChangeSet.h           |  4 +-
 .../Private/_ASHierarchyChangeSet.m           | 54 +++++++++++++++++++
 2 files changed, 57 insertions(+), 1 deletion(-)

diff --git a/AsyncDisplayKit/Private/_ASHierarchyChangeSet.h b/AsyncDisplayKit/Private/_ASHierarchyChangeSet.h
index a2547ff017..33d183a79d 100644
--- a/AsyncDisplayKit/Private/_ASHierarchyChangeSet.h
+++ b/AsyncDisplayKit/Private/_ASHierarchyChangeSet.h
@@ -30,8 +30,10 @@ typedef NS_ENUM(NSInteger, _ASHierarchyChangeType) {
 @property (nonatomic, readonly) ASDataControllerAnimationOptions animationOptions;
 
 /// Index paths are sorted descending for changeType .Delete, ascending otherwise
-@property (nonatomic, strong, readonly) NSArray *indexPaths;
+@property (nonatomic, strong) NSArray *indexPaths;
 @property (nonatomic, readonly) _ASHierarchyChangeType changeType;
+
++ (NSDictionary *)sectionToIndexSetMapFromChanges:(NSArray *)changes ofType:(_ASHierarchyChangeType)changeType;
 @end
 
 @interface _ASHierarchyChangeSet : NSObject
diff --git a/AsyncDisplayKit/Private/_ASHierarchyChangeSet.m b/AsyncDisplayKit/Private/_ASHierarchyChangeSet.m
index bd820e9cd6..39047f867e 100644
--- a/AsyncDisplayKit/Private/_ASHierarchyChangeSet.m
+++ b/AsyncDisplayKit/Private/_ASHierarchyChangeSet.m
@@ -198,9 +198,41 @@
 
     // Ignore item inserts in reloaded(new)/inserted sections.
     [_ASHierarchyItemChange sortAndCoalesceChanges:_insertItemChanges ignoringChangesInSections:insertedOrReloaded];
+    
+
+    // reload itemsChanges need to be adjusted so that we access the correct indexPaths in the datasource
+    NSDictionary *insertedIndexPathsMap = [_ASHierarchyItemChange sectionToIndexSetMapFromChanges:_insertItemChanges ofType:_ASHierarchyChangeTypeInsert];
+    NSDictionary *deletedIndexPathsMap = [_ASHierarchyItemChange sectionToIndexSetMapFromChanges:_deleteItemChanges ofType:_ASHierarchyChangeTypeDelete];
+
+    for (_ASHierarchyItemChange *change in _reloadItemChanges) {
+      NSAssert(change.changeType == _ASHierarchyChangeTypeReload, @"It must be a reload change to be in here");
+      NSMutableArray *newIndexPaths = [NSMutableArray array];
+      
+      // Every indexPaths in the change need to update its section and/or row
+      // depending on all the deletions and insertions
+      for (NSIndexPath *indexPath in change.indexPaths) {
+        NSUInteger section = indexPath.section;
+        NSUInteger row = indexPath.row;
+        
+        section -= [_deletedSections countOfIndexesInRange:NSMakeRange(0, section)];
+        section += [_insertedSections countOfIndexesInRange:NSMakeRange(0, section)];
+        
+        NSIndexSet *indicesInsertedInSection = insertedIndexPathsMap[@(section).stringValue];
+        row += [indicesInsertedInSection countOfIndexesInRange:NSMakeRange(0, row)];
+        NSIndexSet *indicesDeletedInSection = deletedIndexPathsMap[@(section).stringValue];
+        row += [indicesDeletedInSection countOfIndexesInRange:NSMakeRange(0, row)];
+        
+        //TODO: reuse the old indexPath object if section and row aren't changed
+        NSIndexPath *newIndexPath = [NSIndexPath indexPathForRow:row inSection:section];
+        [newIndexPaths addObject:newIndexPath];
+      }
+      change.indexPaths = newIndexPaths;
+    }
   }
 }
 
+
+
 @end
 
 @implementation _ASHierarchySectionChange
@@ -301,6 +333,28 @@
   return self;
 }
 
+// Create a mapping out of changes indexPaths to a {@section : [indexSet]} fashion
+// e.g. changes: (0 - 0), (0 - 1), (2 - 5)
+//  will become: {@0 : [0, 1], @2 : [5]}
++ (NSDictionary *)sectionToIndexSetMapFromChanges:(NSArray *)changes ofType:(_ASHierarchyChangeType)changeType
+{
+  NSMutableDictionary *sectionToIndexSetMap = [NSMutableDictionary dictionary];
+  for (_ASHierarchyItemChange *change in changes) {
+    NSAssert(change.changeType == changeType, @"The map we created must all be of the same changeType as of now");
+    for (NSIndexPath *indexPath in change.indexPaths) {
+      NSString *sectionKey = @(indexPath.section).stringValue;
+      NSMutableIndexSet *indexSet = sectionToIndexSetMap[sectionKey];
+      if (indexSet) {
+        [indexSet addIndex:indexPath.row];
+      } else {
+        indexSet = [NSMutableIndexSet indexSetWithIndex:indexPath.row];
+        sectionToIndexSetMap[sectionKey] = indexSet;
+      }
+    }
+  }
+  return sectionToIndexSetMap;
+}
+
 + (void)sortAndCoalesceChanges:(NSMutableArray *)changes ignoringChangesInSections:(NSIndexSet *)sections
 {
   if (changes.count < 1) {