Add UIDataSourceModelAssociation to ASTableView and ASCollectionView

Source/ASCollectionNode.h

@@ -672,6 +672,29 @@ NS_ASSUME_NONNULL_BEGIN
  */
 - (void)collectionNode:(ASCollectionNode *)collectionNode moveItemAtIndexPath:(NSIndexPath *)sourceIndexPath toIndexPath:(NSIndexPath *)destinationIndexPath;
 
+/**
+ * Generate a unique identifier for an element in a collection. This helps state restoration persist the scroll position
+ *  of a collection view even when the data in that table changes. See the documentation for UIDataSourceModelAssociation for more information.
+ *
+ * @param indexPath The index path of the requested node.
+ *
+ * @param collectionNode The sender.
+ *
+ * @return a unique identifier for the element at the given path. Return nil if the index path does not exist in the collection.
+ */
+- (nullable NSString *)modelIdentifierForElementAtIndexPath:(NSIndexPath *)indexPath inNode:(ASCollectionNode *)collectionNode;
+
+/**
+ * Similar to -collectionView:cellForItemAtIndexPath:. See the documentation for UIDataSourceModelAssociation for more information.
+ *
+ * @param identifier The model identifier of the element, previously generated by a call to modelIdentifierForElementAtIndexPath
+ *
+ * @param collectionNode The sender.
+ *
+ * @return the index path to the current position of the matching element in the collection. Return nil if the element is not found.
+ */
+- (nullable NSIndexPath *)indexPathForElementWithModelIdentifier:(NSString *)identifier inNode:(ASCollectionNode *)collectionNode;
+
 /**
  * Similar to -collectionView:cellForItemAtIndexPath:.
  *

Source/ASCollectionView.mm

@@ -230,6 +230,7 @@ @interface ASCollectionView () <ASRangeControllerDataSource, ASRangeControllerDe
     unsigned int interopAlwaysDequeue:1;
     // Whether this interop data source implements viewForSupplementaryElementOfKind:
     unsigned int interopViewForSupplementaryElement:1;
+    unsigned int modelIdentifierMethods:1; // if both modelIdentifierForElementAtIndexPath and indexPathForElementWithModelIdentifier are implemented
   } _asyncDataSourceFlags;
 
   struct {
@@ -486,6 +487,9 @@ - (void)setAsyncDataSource:(id<ASCollectionDataSource>)asyncDataSource
       _asyncDataSourceFlags.interopViewForSupplementaryElement = [interopDataSource respondsToSelector:@selector(collectionView:viewForSupplementaryElementOfKind:atIndexPath:)];
     }
 
+    _asyncDataSourceFlags.modelIdentifierMethods = [_asyncDataSource respondsToSelector:@selector(modelIdentifierForElementAtIndexPath:inNode:)] && [_asyncDataSource respondsToSelector:@selector(indexPathForElementWithModelIdentifier:inNode:)];
+
+
     ASDisplayNodeAssert(_asyncDataSourceFlags.collectionNodeNumberOfItemsInSection || _asyncDataSourceFlags.collectionViewNumberOfItemsInSection, @"Data source must implement collectionNode:numberOfItemsInSection:");
     ASDisplayNodeAssert(_asyncDataSourceFlags.collectionNodeNodeBlockForItem
                         || _asyncDataSourceFlags.collectionNodeNodeForItem
@@ -805,6 +809,31 @@ - (void)invalidateFlowLayoutDelegateMetrics
   // Subclass hook
 }
 
+- (nullable NSString *)modelIdentifierForElementAtIndexPath:(NSIndexPath *)indexPath inView:(UIView *)view {
+    if (_asyncDataSourceFlags.modelIdentifierMethods) {
+        GET_COLLECTIONNODE_OR_RETURN(collectionNode, nil);
+        NSIndexPath *convertedPath = [self convertIndexPathToCollectionNode:indexPath];
+        if (convertedPath == nil) {
+            return nil;
+        } else {
+            return [_asyncDataSource modelIdentifierForElementAtIndexPath:convertedPath inNode:collectionNode];
+        }
+    } else {
+        return nil;
+    }
+}
+
+- (nullable NSIndexPath *)indexPathForElementWithModelIdentifier:(NSString *)identifier inView:(UIView *)view {
+    if (_asyncDataSourceFlags.modelIdentifierMethods) {
+        GET_COLLECTIONNODE_OR_RETURN(collectionNode, nil);
+        NSIndexPath *result = [_asyncDataSource indexPathForElementWithModelIdentifier:identifier inNode:collectionNode];
+        result = [self convertIndexPathToCollectionNode:result];
+        return result;
+    } else {
+        return nil;
+    }
+}
+
 #pragma mark Internal
 
 - (void)_configureCollectionViewLayout:(nonnull UICollectionViewLayout *)layout

Source/ASTableNode.h

@@ -569,6 +569,29 @@ NS_ASSUME_NONNULL_BEGIN
  */
 - (void)tableViewUnlockDataSource:(ASTableView *)tableView ASDISPLAYNODE_DEPRECATED_MSG("Data source accesses are on the main thread. Method will not be called.");
 
+/**
+ * Generate a unique identifier for an element in a table. This helps state restoration persist the scroll position
+ *  of a table view even when the data in that table changes. See the documentation for UIDataSourceModelAssociation for more information.
+ *
+ * @param indexPath The index path of the requested node.
+ *
+ * @param tableNode The sender.
+ *
+ * @return a unique identifier for the element at the given path. Return nil if the index path does not exist in the table.
+ */
+- (nullable NSString *)modelIdentifierForElementAtIndexPath:(NSIndexPath *)indexPath inNode:(ASTableNode *)tableNode;
+
+/**
+ * Similar to -tableView:cellForRowAtIndexPath:. See the documentation for UIDataSourceModelAssociation for more information.
+ *
+ * @param identifier The model identifier of the element, previously generated by a call to modelIdentifierForElementAtIndexPath.
+ *
+ * @param tableNode The sender.
+ *
+ * @return the index path to the current position of the matching element in the table. Return nil if the element is not found.
+ */
+- (nullable NSIndexPath *)indexPathForElementWithModelIdentifier:(NSString *)identifier inNode:(ASTableNode *)tableNode;
+
 @end
 
 /**

Source/ASTableView.mm

@@ -271,6 +271,7 @@ @interface ASTableView () <ASRangeControllerDataSource, ASRangeControllerDelegat
     unsigned int tableViewMoveRow:1;
     unsigned int tableNodeMoveRow:1;
     unsigned int sectionIndexMethods:1; // if both section index methods are implemented
+    unsigned int modelIdentifierMethods:1; // if both modelIdentifierForElementAtIndexPath and indexPathForElementWithModelIdentifier are implemented
   } _asyncDataSourceFlags;
 }
 
@@ -427,6 +428,7 @@ - (void)setAsyncDataSource:(id<ASTableDataSource>)asyncDataSource
     _asyncDataSourceFlags.tableViewCanMoveRow = [_asyncDataSource respondsToSelector:@selector(tableView:canMoveRowAtIndexPath:)];
     _asyncDataSourceFlags.tableViewMoveRow = [_asyncDataSource respondsToSelector:@selector(tableView:moveRowAtIndexPath:toIndexPath:)];
     _asyncDataSourceFlags.sectionIndexMethods = [_asyncDataSource respondsToSelector:@selector(sectionIndexTitlesForTableView:)] && [_asyncDataSource respondsToSelector:@selector(tableView:sectionForSectionIndexTitle:atIndex:)];
+    _asyncDataSourceFlags.modelIdentifierMethods = [_asyncDataSource respondsToSelector:@selector(modelIdentifierForElementAtIndexPath:inNode:)] && [_asyncDataSource respondsToSelector:@selector(indexPathForElementWithModelIdentifier:inNode:)];
 
     ASDisplayNodeAssert(_asyncDataSourceFlags.tableViewNodeBlockForRow
                         || _asyncDataSourceFlags.tableViewNodeForRow
@@ -961,6 +963,31 @@ - (NSInteger)tableView:(UITableView *)tableView numberOfRowsInSection:(NSInteger
   return [_dataController.visibleMap numberOfItemsInSection:section];
 }
 
+- (nullable NSString *)modelIdentifierForElementAtIndexPath:(NSIndexPath *)indexPath inView:(UIView *)view {
+    if (_asyncDataSourceFlags.modelIdentifierMethods) {
+        GET_TABLENODE_OR_RETURN(tableNode, nil);
+        NSIndexPath *convertedPath = [self convertIndexPathToTableNode:indexPath];
+        if (convertedPath == nil) {
+            return nil;
+        } else {
+            return [_asyncDataSource modelIdentifierForElementAtIndexPath:convertedPath inNode:tableNode];
+        }
+    } else {
+        return nil;
+    }
+}
+
+- (nullable NSIndexPath *)indexPathForElementWithModelIdentifier:(NSString *)identifier inView:(UIView *)view {
+    if (_asyncDataSourceFlags.modelIdentifierMethods) {
+        GET_TABLENODE_OR_RETURN(tableNode, nil);
+        NSIndexPath *result = [_asyncDataSource indexPathForElementWithModelIdentifier:identifier inNode:tableNode];
+        result = [self convertIndexPathFromTableNode:result waitingIfNeeded:YES];
+        return result;
+    } else {
+        return nil;
+    }
+}
+
 - (BOOL)tableView:(UITableView *)tableView canMoveRowAtIndexPath:(NSIndexPath *)indexPath
 {
   if (_asyncDataSourceFlags.tableViewCanMoveRow) {

Source/Details/ASDelegateProxy.mm

@@ -12,6 +12,20 @@
 #import <AsyncDisplayKit/ASCollectionNode.h>
 #import <AsyncDisplayKit/ASAssert.h>
 
+// UIKit performs a class check for UIDataSourceModelAssociation protocol conformance rather than an instance check, so
+//  the implementation of conformsToProtocol: below never gets called. We need to declare the two as conforming to the protocol here, then
+//  we need to implement dummy methods to get rid of a compiler warning about not conforming to the protocol.
+@interface ASTableViewProxy () <UIDataSourceModelAssociation>
+@end
+
+@interface ASCollectionViewProxy () <UIDataSourceModelAssociation>
+@end
+
+@interface ASDelegateProxy (UIDataSourceModelAssociationPrivate)
+- (nullable NSString *)_modelIdentifierForElementAtIndexPath:(NSIndexPath *)indexPath inView:(UIView *)view;
+- (nullable NSIndexPath *)_indexPathForElementWithModelIdentifier:(NSString *)identifier inView:(UIView *)view;
+@end
+
 @implementation ASTableViewProxy
 
 - (BOOL)interceptsSelector:(SEL)selector
@@ -54,10 +68,22 @@ - (BOOL)interceptsSelector:(SEL)selector
 
           // used for batch fetching API
           selector == @selector(scrollViewWillEndDragging:withVelocity:targetContentOffset:) ||
-          selector == @selector(scrollViewDidEndDecelerating:)
+          selector == @selector(scrollViewDidEndDecelerating:) ||
+
+          // UIDataSourceModelAssociation
+          selector == @selector(modelIdentifierForElementAtIndexPath:inView:) ||
+          selector == @selector(indexPathForElementWithModelIdentifier:inView:)
           );
 }
 
+- (nullable NSString *)modelIdentifierForElementAtIndexPath:(NSIndexPath *)indexPath inView:(UIView *)view {
+  return [self _modelIdentifierForElementAtIndexPath:indexPath inView:view];
+}
+
+- (nullable NSIndexPath *)indexPathForElementWithModelIdentifier:(NSString *)identifier inView:(UIView *)view {
+  return [self _indexPathForElementWithModelIdentifier:identifier inView:view];
+}
+
 @end
 
 @implementation ASCollectionViewProxy
@@ -110,10 +136,22 @@ - (BOOL)interceptsSelector:(SEL)selector
 
           // intercepted due to not being supported by ASCollectionView (prevent bugs caused by usage)
           selector == @selector(collectionView:canMoveItemAtIndexPath:) ||
-          selector == @selector(collectionView:moveItemAtIndexPath:toIndexPath:)
+          selector == @selector(collectionView:moveItemAtIndexPath:toIndexPath:) ||
+
+          // UIDataSourceModelAssociation
+          selector == @selector(modelIdentifierForElementAtIndexPath:inView:) ||
+          selector == @selector(indexPathForElementWithModelIdentifier:inView:)
           );
 }
 
+- (nullable NSString *)modelIdentifierForElementAtIndexPath:(NSIndexPath *)indexPath inView:(UIView *)view {
+  return [self _modelIdentifierForElementAtIndexPath:indexPath inView:view];
+}
+
+- (nullable NSIndexPath *)indexPathForElementWithModelIdentifier:(NSString *)identifier inView:(UIView *)view {
+  return [self _indexPathForElementWithModelIdentifier:identifier inView:view];
+}
+
 @end
 
 @implementation ASPagerNodeProxy
@@ -220,4 +258,12 @@ - (BOOL)interceptsSelector:(SEL)selector
   return NO;
 }
 
+- (nullable NSString *)_modelIdentifierForElementAtIndexPath:(NSIndexPath *)indexPath inView:(UIView *)view {
+  return [(id)_interceptor modelIdentifierForElementAtIndexPath:indexPath inView:view];
+}
+
+- (nullable NSIndexPath *)_indexPathForElementWithModelIdentifier:(NSString *)identifier inView:(UIView *)view {
+  return [(id)_interceptor indexPathForElementWithModelIdentifier:identifier inView:view];
+}
+
 @end