[ASDataController] Avoid asking for size ranges of soon-to-be-delete …

…elements during relayouts (TextureGroup#442)

* ASDataController to avoid asking for size ranges of soon-to-be-deleted elements during relayouts

* Remove outdated TODOs

Source/ASCellNode.h

@@ -87,7 +87,6 @@ typedef NS_ENUM(NSUInteger, ASCellNodeVisibilityEvent) {
  *
  * @return The supplementary element kind, or @c nil if this node does not represent a supplementary element.
  */
-//TODO change this to be a generic "kind" or "elementKind" that exposes `nil` for row kind
 @property (atomic, copy, readonly, nullable) NSString *supplementaryElementKind;
 
 /*

Source/Details/ASCollectionElement.h

@@ -26,7 +26,6 @@ NS_ASSUME_NONNULL_BEGIN
 AS_SUBCLASSING_RESTRICTED
 @interface ASCollectionElement : NSObject
 
-//TODO change this to be a generic "kind" or "elementKind" that exposes `nil` for row kind
 @property (nonatomic, readonly, copy, nullable) NSString *supplementaryElementKind;
 @property (nonatomic, assign) ASSizeRange constrainedSize;
 @property (nonatomic, readonly, weak) id<ASRangeManagingNode> owningNode;

Source/Details/ASDataController.mm

@@ -433,21 +433,16 @@ - (void)invalidateDataSourceItemCounts
   return @[];
 }
 
-- (ASSizeRange)constrainedSizeForElement:(ASCollectionElement *)element inElementMap:(ASElementMap *)map
-{
-  ASDisplayNodeAssertMainThread();
-  NSString *kind = element.supplementaryElementKind ?: ASDataControllerRowNodeKind;
-  NSIndexPath *indexPath = [map indexPathForElement:element];
-  return [self constrainedSizeForNodeOfKind:kind atIndexPath:indexPath];
-}
-
-
+/**
+ * Returns constrained size for the node of the given kind and at the given index path.
+ * NOTE: index path must be in the data-source index space.
+ */
 - (ASSizeRange)constrainedSizeForNodeOfKind:(NSString *)kind atIndexPath:(NSIndexPath *)indexPath
 {
   ASDisplayNodeAssertMainThread();
 
   id<ASDataControllerSource> dataSource = _dataSource;
-  if (dataSource == nil) {
+  if (dataSource == nil || indexPath == nil) {
     return ASSizeRangeZero;
   }
 
@@ -750,15 +745,18 @@ - (void)relayoutNodes:(id<NSFastEnumeration>)nodes nodesSizeChanged:(NSMutableAr
   auto pendingMap = self.pendingMap;
   for (ASCellNode *node in nodes) {
     auto element = node.collectionElement;
+    NSIndexPath *indexPathInPendingMap = [pendingMap indexPathForElement:element];
     // Ensure the element is present in both maps or skip it. If it's not in the visible map,
     // then we can't check the presented size. If it's not in the pending map, we can't get the constrained size.
     // This will only happen if the element has been deleted, so the specifics of this behavior aren't important.
-    if ([visibleMap indexPathForElement:element] == nil || [pendingMap indexPathForElement:element] == nil) {
+    if (indexPathInPendingMap == nil || [visibleMap indexPathForElement:element] == nil) {
       continue;
     }
 
-    ASSizeRange constrainedSize = [self constrainedSizeForElement:element inElementMap:pendingMap];
+    NSString *kind = element.supplementaryElementKind ?: ASDataControllerRowNodeKind;
+    ASSizeRange constrainedSize = [self constrainedSizeForNodeOfKind:kind atIndexPath:indexPathInPendingMap];
     [self _layoutNode:node withConstrainedSize:constrainedSize];
+
     BOOL matchesSize = [dataSource dataController:self presentedSizeForElement:element matchesSize:node.frame.size];
     if (! matchesSize) {
       [nodesSizesChanged addObject:node];
@@ -785,15 +783,23 @@ - (void)_relayoutAllNodes
 {
   ASDisplayNodeAssertMainThread();
   for (ASCollectionElement *element in _visibleMap) {
-    ASSizeRange constrainedSize = [self constrainedSizeForElement:element inElementMap:_visibleMap];
-    if (ASSizeRangeHasSignificantArea(constrainedSize)) {
-      element.constrainedSize = constrainedSize;
+    // Ignore this element if it is no longer in the latest data. It is still recognized in the UIKit world but will be deleted soon.
+    NSIndexPath *indexPathInPendingMap = [_pendingMap indexPathForElement:element];
+    if (indexPathInPendingMap == nil) {
+      continue;
+    }
+
+    NSString *kind = element.supplementaryElementKind ?: ASDataControllerRowNodeKind;
+    ASSizeRange newConstrainedSize = [self constrainedSizeForNodeOfKind:kind atIndexPath:indexPathInPendingMap];
+
+    if (ASSizeRangeHasSignificantArea(newConstrainedSize)) {
+      element.constrainedSize = newConstrainedSize;
 
       // Node may not be allocated yet (e.g node virtualization or same size optimization)
       // Call context.nodeIfAllocated here to avoid immature node allocation and layout
       ASCellNode *node = element.nodeIfAllocated;
       if (node) {
-        [self _layoutNode:node withConstrainedSize:constrainedSize];
+        [self _layoutNode:node withConstrainedSize:newConstrainedSize];
       }
     }
   }