Add support for keeping letting cell nodes update to new view models when reloaded. #trivial (#357)

* Add support for skipping reload if node decides it is compatible with new view model also

* Sort things right

* Put the order back

* No need for redundant expectation

* Fix license header

* Fix comment

AsyncDisplayKit.xcodeproj/project.pbxproj

@@ -326,6 +326,7 @@
 		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 */; };
+		CC311E071EEF81C400A8D7A6 /* ASDisplayNode+OCMock.m in Sources */ = {isa = PBXBuildFile; fileRef = CC311E061EEF81C400A8D7A6 /* ASDisplayNode+OCMock.m */; };
 		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,6 +785,7 @@
 		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>"; };
+		CC311E061EEF81C400A8D7A6 /* ASDisplayNode+OCMock.m */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.objc; path = "ASDisplayNode+OCMock.m"; 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>"; };
@@ -1117,6 +1119,7 @@
 		058D09C5195D04C000B7D73C /* Tests */ = {
 			isa = PBXGroup;
 			children = (
+				CC311E061EEF81C400A8D7A6 /* ASDisplayNode+OCMock.m */,
 				CCB338E51EEE27760081F21A /* ASTestCase.h */,
 				CCB338E61EEE27760081F21A /* ASTestCase.m */,
 				CCB338E21EEE11160081F21A /* OCMockObject+ASAdditions.h */,
@@ -2077,6 +2080,7 @@
 				058D0A3C195D057000B7D73C /* ASMutableAttributedStringBuilderTests.m in Sources */,
 				CC8B05D81D73979700F54286 /* ASTextNodePerformanceTests.m in Sources */,
 				697B315A1CFE4B410049936F /* ASEditableTextNodeTests.m in Sources */,
+				CC311E071EEF81C400A8D7A6 /* ASDisplayNode+OCMock.m in Sources */,
 				CCB338E41EEE11160081F21A /* OCMockObject+ASAdditions.m in Sources */,
 				ACF6ED611B178DC700DA7C62 /* ASOverlayLayoutSpecSnapshotTests.mm in Sources */,
 				CC8B05D61D73836400F54286 /* ASPerformanceTestContext.m in Sources */,

Source/ASCellNode.h

@@ -126,6 +126,13 @@ typedef NS_ENUM(NSUInteger, ASCellNodeVisibilityEvent) {
  */
 @property (atomic, nullable) id viewModel;
 
+/**
+ * Asks the node whether it can be updated to the given view model.
+ *
+ * The default implementation returns YES if the class matches that of the current view-model.
+ */
+- (BOOL)canUpdateToViewModel:(id)viewModel;
+
 /**
  * The backing view controller, or @c nil if the node wasn't initialized with backing view controller
  * @note This property must be accessed on the main thread.

Source/ASCellNode.mm

@@ -177,6 +177,11 @@
   }
 }
 
+- (BOOL)canUpdateToViewModel:(id)viewModel
+{
+  return [self.viewModel class] == [viewModel class];
+}
+
 - (NSIndexPath *)indexPath
 {
   return [self.owningNode indexPathForNode:self];

Source/Details/ASDataController.mm

@@ -303,6 +303,7 @@ typedef void (^ASDataControllerCompletionBlock)(NSArray<ASCollectionElement *> *
                              traitCollection:(ASPrimitiveTraitCollection)traitCollection
                             indexPathsAreNew:(BOOL)indexPathsAreNew
                        shouldFetchSizeRanges:(BOOL)shouldFetchSizeRanges
+                                 previousMap:(ASElementMap *)previousMap
 {
   ASDisplayNodeAssertMainThread();
 
@@ -325,7 +326,7 @@ typedef void (^ASDataControllerCompletionBlock)(NSArray<ASCollectionElement *> *
   }
 
   for (NSString *kind in [self supplementaryKindsInSections:newSections]) {
-    [self _insertElementsIntoMap:map kind:kind forSections:newSections traitCollection:traitCollection shouldFetchSizeRanges:shouldFetchSizeRanges];
+    [self _insertElementsIntoMap:map kind:kind forSections:newSections traitCollection:traitCollection shouldFetchSizeRanges:shouldFetchSizeRanges changeSet:changeSet previousMap:previousMap];
   }
 }
 
@@ -342,6 +343,8 @@ typedef void (^ASDataControllerCompletionBlock)(NSArray<ASCollectionElement *> *
                    forSections:(NSIndexSet *)sections
                traitCollection:(ASPrimitiveTraitCollection)traitCollection
          shouldFetchSizeRanges:(BOOL)shouldFetchSizeRanges
+                     changeSet:(_ASHierarchyChangeSet *)changeSet
+                   previousMap:(ASElementMap *)previousMap
 {
   ASDisplayNodeAssertMainThread();
   
@@ -350,7 +353,7 @@ typedef void (^ASDataControllerCompletionBlock)(NSArray<ASCollectionElement *> *
   }
   
   NSArray<NSIndexPath *> *indexPaths = [self _allIndexPathsForItemsOfKind:kind inSections:sections];
-  [self _insertElementsIntoMap:map kind:kind atIndexPaths:indexPaths traitCollection:traitCollection shouldFetchSizeRanges:shouldFetchSizeRanges];
+  [self _insertElementsIntoMap:map kind:kind atIndexPaths:indexPaths traitCollection:traitCollection shouldFetchSizeRanges:shouldFetchSizeRanges changeSet:changeSet previousMap:previousMap];
 }
 
 /**
@@ -367,6 +370,8 @@ typedef void (^ASDataControllerCompletionBlock)(NSArray<ASCollectionElement *> *
                   atIndexPaths:(NSArray<NSIndexPath *> *)indexPaths
                traitCollection:(ASPrimitiveTraitCollection)traitCollection
          shouldFetchSizeRanges:(BOOL)shouldFetchSizeRanges
+                     changeSet:(_ASHierarchyChangeSet *)changeSet
+                   previousMap:(ASElementMap *)previousMap
 {
   ASDisplayNodeAssertMainThread();
   
@@ -384,11 +389,28 @@ typedef void (^ASDataControllerCompletionBlock)(NSArray<ASCollectionElement *> *
   id<ASDataControllerSource> dataSource = self.dataSource;
   id<ASRangeManagingNode> node = self.node;
   for (NSIndexPath *indexPath in indexPaths) {
-    id viewModel = [dataSource dataController:self viewModelForItemAtIndexPath:indexPath];
-
     ASCellNodeBlock nodeBlock;
+    id viewModel;
     if (isRowKind) {
+      viewModel = [dataSource dataController:self viewModelForItemAtIndexPath:indexPath];
+      
+      // Get the prior element and attempt to update the existing cell node.
+      if (viewModel != nil && !changeSet.includesReloadData) {
+        NSIndexPath *oldIndexPath = [changeSet oldIndexPathForNewIndexPath:indexPath];
+        if (oldIndexPath != nil) {
+          ASCollectionElement *oldElement = [previousMap elementForItemAtIndexPath:oldIndexPath];
+          ASCellNode *oldNode = oldElement.node;
+          if ([oldNode canUpdateToViewModel:viewModel]) {
+            // Just wrap the node in a block. The collection element will -setViewModel:
+            nodeBlock = ^{
+              return oldNode;
+            };
+          }
+        }
+      }
+      if (nodeBlock == nil) {
         nodeBlock = [dataSource dataController:self nodeBlockAtIndexPath:indexPath];
+      }
     } else {
       nodeBlock = [dataSource dataController:self supplementaryNodeBlockOfKind:kind atIndexPath:indexPath];
     }
@@ -534,14 +556,15 @@ typedef void (^ASDataControllerCompletionBlock)(NSArray<ASCollectionElement *> *
   }
 
   // Mutable copy of current data.
-  ASMutableElementMap *mutableMap = [_pendingMap mutableCopy];
+  ASElementMap *previousMap = _pendingMap;
+  ASMutableElementMap *mutableMap = [previousMap mutableCopy];
   
   BOOL canDelegateLayout = (_layoutDelegate != nil);
 
   // Step 1: Update the mutable copies to match the data source's state
   [self _updateSectionContextsInMap:mutableMap changeSet:changeSet];
   ASPrimitiveTraitCollection existingTraitCollection = [self.node primitiveTraitCollection];
-  [self _updateElementsInMap:mutableMap changeSet:changeSet traitCollection:existingTraitCollection shouldFetchSizeRanges:(! canDelegateLayout)];
+  [self _updateElementsInMap:mutableMap changeSet:changeSet traitCollection:existingTraitCollection shouldFetchSizeRanges:(! canDelegateLayout) previousMap:previousMap];
   
   // Step 2: Clone the new data
   ASElementMap *newMap = [mutableMap copy];
@@ -644,6 +667,7 @@ typedef void (^ASDataControllerCompletionBlock)(NSArray<ASCollectionElement *> *
                    changeSet:(_ASHierarchyChangeSet *)changeSet
              traitCollection:(ASPrimitiveTraitCollection)traitCollection
        shouldFetchSizeRanges:(BOOL)shouldFetchSizeRanges
+                 previousMap:(ASElementMap *)previousMap
 {
   ASDisplayNodeAssertMainThread();
 
@@ -653,7 +677,7 @@ typedef void (^ASDataControllerCompletionBlock)(NSArray<ASCollectionElement *> *
     NSUInteger sectionCount = [self itemCountsFromDataSource].size();
     if (sectionCount > 0) {
       NSIndexSet *sectionIndexes = [NSIndexSet indexSetWithIndexesInRange:NSMakeRange(0, sectionCount)];
-      [self _insertElementsIntoMap:map sections:sectionIndexes traitCollection:traitCollection shouldFetchSizeRanges:shouldFetchSizeRanges];
+      [self _insertElementsIntoMap:map sections:sectionIndexes traitCollection:traitCollection shouldFetchSizeRanges:shouldFetchSizeRanges changeSet:changeSet previousMap:previousMap];
     }
     // Return immediately because reloadData can't be used in conjuntion with other updates.
     return;
@@ -666,7 +690,8 @@ typedef void (^ASDataControllerCompletionBlock)(NSArray<ASCollectionElement *> *
                                      changeSet:changeSet
                                traitCollection:traitCollection
                               indexPathsAreNew:NO
-                         shouldFetchSizeRanges:shouldFetchSizeRanges];
+                         shouldFetchSizeRanges:shouldFetchSizeRanges
+                                   previousMap:previousMap];
   }
 
   for (_ASHierarchySectionChange *change in [changeSet sectionChangesOfType:_ASHierarchyChangeTypeDelete]) {
@@ -676,17 +701,18 @@ typedef void (^ASDataControllerCompletionBlock)(NSArray<ASCollectionElement *> *
   }
   
   for (_ASHierarchySectionChange *change in [changeSet sectionChangesOfType:_ASHierarchyChangeTypeInsert]) {
-    [self _insertElementsIntoMap:map sections:change.indexSet traitCollection:traitCollection shouldFetchSizeRanges:shouldFetchSizeRanges];
+    [self _insertElementsIntoMap:map sections:change.indexSet traitCollection:traitCollection shouldFetchSizeRanges:shouldFetchSizeRanges changeSet:changeSet previousMap:previousMap];
   }
   
   for (_ASHierarchyItemChange *change in [changeSet itemChangesOfType:_ASHierarchyChangeTypeInsert]) {
-    [self _insertElementsIntoMap:map kind:ASDataControllerRowNodeKind atIndexPaths:change.indexPaths traitCollection:traitCollection shouldFetchSizeRanges:shouldFetchSizeRanges];
+    [self _insertElementsIntoMap:map kind:ASDataControllerRowNodeKind atIndexPaths:change.indexPaths traitCollection:traitCollection shouldFetchSizeRanges:shouldFetchSizeRanges changeSet:changeSet previousMap:previousMap];
     // Aggressively reload supplementary nodes (#1773 & #1629)
     [self _repopulateSupplementaryNodesIntoMap:map forSectionsContainingIndexPaths:change.indexPaths
                                      changeSet:changeSet
                                traitCollection:traitCollection
                               indexPathsAreNew:YES
-                         shouldFetchSizeRanges:shouldFetchSizeRanges];
+                         shouldFetchSizeRanges:shouldFetchSizeRanges
+                                   previousMap:previousMap];
   }
 }
 
@@ -694,6 +720,8 @@ typedef void (^ASDataControllerCompletionBlock)(NSArray<ASCollectionElement *> *
                       sections:(NSIndexSet *)sectionIndexes
                traitCollection:(ASPrimitiveTraitCollection)traitCollection
          shouldFetchSizeRanges:(BOOL)shouldFetchSizeRanges
+                     changeSet:(_ASHierarchyChangeSet *)changeSet
+                   previousMap:(ASElementMap *)previousMap
 {
   ASDisplayNodeAssertMainThread();
   
@@ -703,12 +731,12 @@ typedef void (^ASDataControllerCompletionBlock)(NSArray<ASCollectionElement *> *
 
   // Items
   [map insertEmptySectionsOfItemsAtIndexes:sectionIndexes];
-  [self _insertElementsIntoMap:map kind:ASDataControllerRowNodeKind forSections:sectionIndexes traitCollection:traitCollection shouldFetchSizeRanges:shouldFetchSizeRanges];
+  [self _insertElementsIntoMap:map kind:ASDataControllerRowNodeKind forSections:sectionIndexes traitCollection:traitCollection shouldFetchSizeRanges:shouldFetchSizeRanges changeSet:changeSet previousMap:previousMap];
 
   // Supplementaries
   for (NSString *kind in [self supplementaryKindsInSections:sectionIndexes]) {
     // Step 2: Populate new elements for all sections
-    [self _insertElementsIntoMap:map kind:kind forSections:sectionIndexes traitCollection:traitCollection shouldFetchSizeRanges:shouldFetchSizeRanges];
+    [self _insertElementsIntoMap:map kind:kind forSections:sectionIndexes traitCollection:traitCollection shouldFetchSizeRanges:shouldFetchSizeRanges changeSet:changeSet previousMap:previousMap];
   }
 }
 

Source/Private/_ASHierarchyChangeSet.h

@@ -147,6 +147,14 @@ NSString *NSStringFromASHierarchyChangeType(_ASHierarchyChangeType changeType);
  */
 - (NSUInteger)newSectionForOldSection:(NSUInteger)oldSection;
 
+/**
+ * Get the old item index path for the given new index path.
+ *
+ * @precondition The change set must be completed.
+ * @return The old index path, or nil if the given item was inserted.
+ */
+- (nullable NSIndexPath *)oldIndexPathForNewIndexPath:(NSIndexPath *)indexPath;
+
 /// Call this once the change set has been constructed to prevent future modifications to the changeset. Calling this more than once is a programmer error.
 /// NOTE: Calling this method will cause the changeset to convert all reloads into delete/insert pairs.
 - (void)markCompletedWithNewItemCounts:(std::vector<NSInteger>)newItemCounts;

Source/Private/_ASHierarchyChangeSet.mm

@@ -1,5 +1,5 @@
 //
-//  _ASHierarchyChangeSet.m
+//  _ASHierarchyChangeSet.mm
 //  Texture
 //
 //  Copyright (c) 2014-present, Facebook, Inc.  All rights reserved.
@@ -256,6 +256,45 @@ NSString *NSStringFromASHierarchyChangeType(_ASHierarchyChangeType changeType)
   return newIndex;
 }
 
+- (NSUInteger)oldSectionForNewSection:(NSUInteger)newSection
+{
+  [self _ensureCompleted];
+  if ([_insertedSections containsIndex:newSection]) {
+    return NSNotFound;
+  }
+  
+  NSInteger oldIndex = newSection - [_insertedSections as_indexChangeByInsertingItemsBelowIndex:newSection];
+  oldIndex += [_deletedSections countOfIndexesInRange:NSMakeRange(0, oldIndex)];
+  return oldIndex;
+}
+
+- (NSIndexPath *)oldIndexPathForNewIndexPath:(NSIndexPath *)indexPath
+{
+  [self _ensureCompleted];
+  // Inserted sections return nil.
+  NSInteger newSection = indexPath.section;
+  NSInteger newItem = indexPath.item;
+  NSInteger oldSection = [self oldSectionForNewSection:newSection];
+  if (oldSection == NSNotFound) {
+    return nil;
+  }
+  
+  // Inserted items return nil.
+  for (_ASHierarchyItemChange *change in _originalInsertItemChanges) {
+    if ([change.indexPaths containsObject:indexPath]) {
+      return nil;
+    }
+  }
+  
+  // TODO: This is a pretty inefficient way to do this.
+  NSIndexSet *insertsInSection = [_ASHierarchyItemChange sectionToIndexSetMapFromChanges:_insertItemChanges][@(newSection)];
+  NSIndexSet *deletesInSection = [_ASHierarchyItemChange sectionToIndexSetMapFromChanges:_deleteItemChanges][@(oldSection)];
+  
+  NSInteger oldIndex = newItem - [insertsInSection as_indexChangeByInsertingItemsBelowIndex:newItem];
+  oldIndex += [deletesInSection countOfIndexesInRange:NSMakeRange(0, oldIndex)];
+  return [NSIndexPath indexPathForItem:oldIndex inSection:oldSection];
+}
+
 - (void)reloadData
 {
   [self _ensureNotCompleted];

Tests/ASCollectionModernDataSourceTests.m

@@ -18,7 +18,9 @@
 #import "ASTestCase.h"
 
 @interface ASCollectionModernDataSourceTests : ASTestCase
+@end
 
+@interface ASTestCellNode : ASCellNode
 @end
 
 @implementation ASCollectionModernDataSourceTests {
@@ -61,6 +63,7 @@
 
 - (void)tearDown
 {
+  [collectionNode waitUntilAllUpdatesAreCommitted];
   OCMVerifyAll(mockDataSource);
   [super tearDown];
 }
@@ -79,10 +82,11 @@
   // Reload at (0, 0)
   NSIndexPath *reloadedPath = [NSIndexPath indexPathForItem:0 inSection:0];
 
-  sections[reloadedPath.section][reloadedPath.item] = [NSObject new];
+  [self performUpdateReloadingItems:@{ reloadedPath: [NSObject new] }
-  [self performUpdateInvalidatingItems:@[ reloadedPath ] block:^{
+                     reloadMappings:@{ reloadedPath: reloadedPath }
-    [collectionNode reloadItemsAtIndexPaths:@[ reloadedPath ]];
+                     insertingItems:nil
-  }];
+                      deletingItems:nil
+            skippedReloadIndexPaths:nil];
 }
 
 - (void)testInsertingAnItem
@@ -92,10 +96,36 @@
   // Insert at (1, 0)
   NSIndexPath *insertedPath = [NSIndexPath indexPathForItem:0 inSection:1];
 
-  [sections[insertedPath.section] insertObject:[NSObject new] atIndex:insertedPath.item];
+  [self performUpdateReloadingItems:nil
-  [self performUpdateInvalidatingItems:@[ insertedPath ] block:^{
+                     reloadMappings:nil
-    [collectionNode insertItemsAtIndexPaths:@[ insertedPath ]];
+                     insertingItems:@{ insertedPath: [NSObject new] }
-  }];
+                      deletingItems:nil
+            skippedReloadIndexPaths:nil];
+}
+
+- (void)testReloadingAnItemWithACompatibleViewModel
+{
+  [self loadInitialData];
+
+  // Reload and delete together, for good measure.
+  NSIndexPath *reloadedPath = [NSIndexPath indexPathForItem:1 inSection:0];
+  NSIndexPath *deletedPath = [NSIndexPath indexPathForItem:0 inSection:0];
+
+  id viewModel = [NSObject new];
+
+  // Cell node should get -canUpdateToViewModel:
+  id mockCellNode = [collectionNode nodeForItemAtIndexPath:reloadedPath];
+  [mockCellNode setExpectationOrderMatters:YES];
+  OCMExpect([mockCellNode canUpdateToViewModel:viewModel])
+  .andReturn(YES);
+
+  [self performUpdateReloadingItems:@{ reloadedPath: viewModel }
+                     reloadMappings:@{ reloadedPath: [NSIndexPath indexPathForItem:0 inSection:0] }
+                     insertingItems:nil
+                      deletingItems:@[ deletedPath ]
+            skippedReloadIndexPaths:@[ reloadedPath ]];
+  
+  OCMVerifyAll(mockCellNode);
 }
 
 #pragma mark - Helpers
@@ -114,7 +144,8 @@
       // For each item:
       for (NSInteger i = 0; i < items.count; i++) {
         NSIndexPath *indexPath = [NSIndexPath indexPathForItem:i inSection:section];
-        [self expectContentMethodsForItemAtIndexPath:indexPath];
+        [self expectViewModelMethodForItemAtIndexPath:indexPath viewModel:items[i]];
+        [self expectNodeBlockMethodForItemAtIndexPath:indexPath];
       }
     }
   }
@@ -147,16 +178,25 @@
   }
 }
 
-// Expects viewModelForItemAtIndexPath: and nodeBlockForItemAtIndexPath:
+- (void)expectViewModelMethodForItemAtIndexPath:(NSIndexPath *)indexPath viewModel:(id)viewModel
-- (void)expectContentMethodsForItemAtIndexPath:(NSIndexPath *)indexPath
 {
-  id viewModel = sections[indexPath.section][indexPath.item];
   OCMExpect([mockDataSource collectionNode:collectionNode viewModelForItemAtIndexPath:indexPath])
   .andReturn(viewModel);
-  OCMExpect([mockDataSource collectionNode:collectionNode nodeBlockForItemAtIndexPath:indexPath])
-  .andReturn((ASCellNodeBlock)^{ return [ASCellNode new]; });
 }
 
+- (void)expectNodeBlockMethodForItemAtIndexPath:(NSIndexPath *)indexPath
+{
+  OCMExpect([mockDataSource collectionNode:collectionNode nodeBlockForItemAtIndexPath:indexPath])
+  .andReturn((ASCellNodeBlock)^{
+    ASCellNode *node = [ASTestCellNode new];
+    // Generating multiple partial mocks of the same class is not thread-safe.
+    @synchronized (NSNull.null) {
+      return OCMPartialMock(node);
+    }
+  });
+}
+
+/// Asserts that counts match and all view-models are up-to-date between us and collectionNode.
 - (void)assertCollectionNodeContent
 {
   // Assert section count
@@ -182,21 +222,70 @@
  * Updates the collection node, with expectations and assertions about the call-order and the correctness of the
  * new data. You should update the data source _before_ calling this method.
  *
- * invalidatedIndexPaths are the items we expect to get refetched (reloaded/inserted).
+ * skippedReloadIndexPaths are the old index paths for nodes that should use -canUpdateToViewModel: instead of being refetched.
  */
-- (void)performUpdateInvalidatingItems:(NSArray<NSIndexPath *> *)invalidatedIndexPaths block:(void(^)())update
+- (void)performUpdateReloadingItems:(NSDictionary<NSIndexPath *, id> *)reloadedItems
+                     reloadMappings:(NSDictionary<NSIndexPath *, NSIndexPath *> *)reloadMappings
+                     insertingItems:(NSDictionary<NSIndexPath *, id> *)insertedItems
+                      deletingItems:(NSArray<NSIndexPath *> *)deletedItems
+            skippedReloadIndexPaths:(NSArray<NSIndexPath *> *)skippedReloadIndexPaths
 {
-  // When we do an edit, it'll read the new counts
+  [collectionNode performBatchUpdates:^{
-  [self expectDataSourceCountMethods];
+    // First update our data source.
+    [reloadedItems enumerateKeysAndObjectsUsingBlock:^(NSIndexPath * _Nonnull key, id  _Nonnull obj, BOOL * _Nonnull stop) {
+      sections[key.section][key.item] = obj;
+    }];
     
-  // Then it'll load the contents for inserted/reloaded items.
+    // Deletion paths, sorted descending
-  for (NSIndexPath *indexPath in invalidatedIndexPaths) {
+    for (NSIndexPath *indexPath in [deletedItems sortedArrayUsingSelector:@selector(compare:)].reverseObjectEnumerator) {
-    [self expectContentMethodsForItemAtIndexPath:indexPath];
+      [sections[indexPath.section] removeObjectAtIndex:indexPath.item];
     }
     
-  [collectionNode performBatchUpdates:update completion:nil];
+    // Insertion paths, sorted ascending.
+    NSArray *insertionsSortedAcending = [insertedItems.allKeys sortedArrayUsingSelector:@selector(compare:)];
+    for (NSIndexPath *indexPath in insertionsSortedAcending) {
+      [sections[indexPath.section] insertObject:insertedItems[indexPath] atIndex:indexPath.item];
+    }
     
+    // Then update the collection node.
+    [collectionNode reloadItemsAtIndexPaths:reloadedItems.allKeys];
+    [collectionNode deleteItemsAtIndexPaths:deletedItems];
+    [collectionNode insertItemsAtIndexPaths:insertedItems.allKeys];
+    
+    // Before the commit, lay out our expectations.
+    
+    // Expect it to load the new counts.
+    [self expectDataSourceCountMethods];
+    
+    // Combine reloads + inserts and expect them to load content for all of them, in ascending order.
+    NSMutableDictionary<NSIndexPath *, id> *insertsPlusReloads = [NSMutableDictionary dictionary];
+    [insertsPlusReloads addEntriesFromDictionary:insertedItems];
+    [reloadedItems enumerateKeysAndObjectsUsingBlock:^(NSIndexPath * _Nonnull key, id  _Nonnull obj, BOOL * _Nonnull stop) {
+      insertsPlusReloads[reloadMappings[key]] = obj;
+    }];
+    
+    for (NSIndexPath *indexPath in [insertsPlusReloads.allKeys sortedArrayUsingSelector:@selector(compare:)]) {
+      [self expectViewModelMethodForItemAtIndexPath:indexPath viewModel:insertsPlusReloads[indexPath]];
+      NSIndexPath *oldIndexPath = [reloadMappings allKeysForObject:indexPath].firstObject;
+      BOOL isSkippedReload = oldIndexPath && [skippedReloadIndexPaths containsObject:oldIndexPath];
+      if (!isSkippedReload) {
+        [self expectNodeBlockMethodForItemAtIndexPath:indexPath];
+      }
+    }
+  } completion:nil];
+
+  // Assert that the counts and view models are all correct now.
   [self assertCollectionNodeContent];
 }
 
 @end
+
+@implementation ASTestCellNode
+
+- (BOOL)canUpdateToViewModel:(id)viewModel
+{
+  // Our tests default to NO for migrating view models. We use OCMExpect to return YES when we specifically want to.
+  return NO;
+}
+
+@end

Tests/ASDisplayNode+OCMock.m

@@ -0,0 +1,31 @@
+//
+//  ASDisplayNode+OCMock.m
+//  Texture
+//
+//  Copyright (c) 2017-present, Pinterest, Inc.  All rights reserved.
+//  Licensed under the Apache License, Version 2.0 (the "License");
+//  you may not use this file except in compliance with the License.
+//  You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+
+#import <AsyncDisplayKit/AsyncDisplayKit.h>
+
+/**
+ * For some reason, when creating partial mocks of nodes, OCMock fails to find
+ * these class methods that it swizzled!
+ */
+@implementation ASDisplayNode (OCMock)
+
++ (Class)ocmock_replaced_viewClass
+{
+  return [_ASDisplayView class];
+}
+
++ (Class)ocmock_replaced_layerClass
+{
+  return [_ASDisplayLayer class];
+}
+
+@end