Initial Work for Range Controller to Support Supplementary Elements (#3140)

* Initial work supporting supplementaries in range controller

* Rename indexPathForElementIfItem

AsyncDisplayKit.xcodeproj/project.pbxproj

@@ -98,7 +98,7 @@
 		509E68611B3AEDA0009B9150 /* ASAbstractLayoutController.h in Headers */ = {isa = PBXBuildFile; fileRef = 205F0E171B37339C007741D0 /* ASAbstractLayoutController.h */; };
 		509E68621B3AEDA5009B9150 /* ASAbstractLayoutController.mm in Sources */ = {isa = PBXBuildFile; fileRef = 205F0E181B37339C007741D0 /* ASAbstractLayoutController.mm */; };
 		509E68631B3AEDB4009B9150 /* ASCollectionViewLayoutController.h in Headers */ = {isa = PBXBuildFile; fileRef = 205F0E1B1B373A2C007741D0 /* ASCollectionViewLayoutController.h */; };
-		509E68641B3AEDB7009B9150 /* ASCollectionViewLayoutController.mm in Sources */ = {isa = PBXBuildFile; fileRef = 205F0E1C1B373A2C007741D0 /* ASCollectionViewLayoutController.mm */; };
+		509E68641B3AEDB7009B9150 /* ASCollectionViewLayoutController.m in Sources */ = {isa = PBXBuildFile; fileRef = 205F0E1C1B373A2C007741D0 /* ASCollectionViewLayoutController.m */; };
 		509E68651B3AEDC5009B9150 /* CoreGraphics+ASConvenience.h in Headers */ = {isa = PBXBuildFile; fileRef = 205F0E1F1B376416007741D0 /* CoreGraphics+ASConvenience.h */; settings = {ATTRIBUTES = (Public, ); }; };
 		509E68661B3AEDD7009B9150 /* CoreGraphics+ASConvenience.m in Sources */ = {isa = PBXBuildFile; fileRef = 205F0E201B376416007741D0 /* CoreGraphics+ASConvenience.m */; };
 		636EA1A41C7FF4EC00EE152F /* NSArray+Diffing.m in Sources */ = {isa = PBXBuildFile; fileRef = DBC452DA1C5BF64600B16017 /* NSArray+Diffing.m */; };
@@ -490,7 +490,7 @@
 		205F0E171B37339C007741D0 /* ASAbstractLayoutController.h */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.h; path = ASAbstractLayoutController.h; sourceTree = "<group>"; };
 		205F0E181B37339C007741D0 /* ASAbstractLayoutController.mm */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.cpp.objcpp; path = ASAbstractLayoutController.mm; sourceTree = "<group>"; };
 		205F0E1B1B373A2C007741D0 /* ASCollectionViewLayoutController.h */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.h; path = ASCollectionViewLayoutController.h; sourceTree = "<group>"; };
-		205F0E1C1B373A2C007741D0 /* ASCollectionViewLayoutController.mm */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.cpp.objcpp; path = ASCollectionViewLayoutController.mm; sourceTree = "<group>"; };
+		205F0E1C1B373A2C007741D0 /* ASCollectionViewLayoutController.m */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.objc; path = ASCollectionViewLayoutController.m; sourceTree = "<group>"; };
 		205F0E1F1B376416007741D0 /* CoreGraphics+ASConvenience.h */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.h; path = "CoreGraphics+ASConvenience.h"; sourceTree = "<group>"; };
 		205F0E201B376416007741D0 /* CoreGraphics+ASConvenience.m */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.objc; path = "CoreGraphics+ASConvenience.m"; sourceTree = "<group>"; };
 		242995D21B29743C00090100 /* ASBasicImageDownloaderTests.m */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.objc; path = ASBasicImageDownloaderTests.m; sourceTree = "<group>"; };
@@ -1071,7 +1071,7 @@
 				68C215561DE10D330019C4BC /* ASCollectionViewLayoutInspector.h */,
 				68C215571DE10D330019C4BC /* ASCollectionViewLayoutInspector.m */,
 				205F0E1B1B373A2C007741D0 /* ASCollectionViewLayoutController.h */,
-				205F0E1C1B373A2C007741D0 /* ASCollectionViewLayoutController.mm */,
+				205F0E1C1B373A2C007741D0 /* ASCollectionViewLayoutController.m */,
 				696F01EA1DD2AF450049FBD5 /* ASEventLog.h */,
 				696F01EB1DD2AF450049FBD5 /* ASEventLog.mm */,
 				4640521B1A3F83C40061C0BA /* ASTableLayoutController.h */,
@@ -1861,7 +1861,7 @@
 				69CB62AE1CB8165900024920 /* _ASDisplayViewAccessiblity.mm in Sources */,
 				CC034A021E5FAF9700626263 /* ASElementMap.m in Sources */,
 				B35061F61B010EFD0018CF92 /* ASCollectionView.mm in Sources */,
-				509E68641B3AEDB7009B9150 /* ASCollectionViewLayoutController.mm in Sources */,
+				509E68641B3AEDB7009B9150 /* ASCollectionViewLayoutController.m in Sources */,
 				B35061F91B010EFD0018CF92 /* ASControlNode.mm in Sources */,
 				8021EC1F1D2B00B100799119 /* UIImage+ASConvenience.m in Sources */,
 				B35062181B010EFD0018CF92 /* ASDataController.mm in Sources */,

Source/ASCollectionView.mm

@@ -1679,13 +1679,51 @@ static NSString * const kReuseIdentifier = @"_ASCollectionReuseIdentifier";
   return _rangeController;
 }
 
-- (NSArray *)visibleNodeIndexPathsForRangeController:(ASRangeController *)rangeController
+/// The UIKit version of this method is only available on iOS >= 9
+- (NSArray<NSIndexPath *> *)asdk_indexPathsForVisibleSupplementaryElementsOfKind:(NSString *)kind
 {
-  ASDisplayNodeAssertMainThread();
+  if (NSFoundationVersionNumber >= NSFoundationVersionNumber_iOS_9_0) {
-  // Calling -indexPathsForVisibleItems will trigger UIKit to call reloadData if it never has, which can result
+    return [self indexPathsForVisibleSupplementaryElementsOfKind:kind];
-  // in incorrect layout if performed at zero size.  We can use the fact that nothing can be visible at zero size to return fast.
+  }
-  BOOL isZeroSized = CGSizeEqualToSize(self.bounds.size, CGSizeZero);
+
-  return isZeroSized ? @[] : [self indexPathsForVisibleItems];
+  // iOS 8 workaround
+  // We cannot use willDisplaySupplementaryView/didEndDisplayingSupplementaryView
+  // because those methods send index paths for _deleted items_ (invalid index paths)
+  [self layoutIfNeeded];
+  NSArray<UICollectionViewLayoutAttributes *> *visibleAttributes = [self.collectionViewLayout layoutAttributesForElementsInRect:self.bounds];
+  NSMutableArray *result = [NSMutableArray array];
+  for (UICollectionViewLayoutAttributes *attributes in visibleAttributes) {
+    if (attributes.representedElementCategory == UICollectionElementCategorySupplementaryView
+        && [attributes.representedElementKind isEqualToString:kind]) {
+      [result addObject:attributes.indexPath];
+    }
+  }
+  return result;
+}
+
+- (NSArray<ASCollectionElement *> *)visibleElementsForRangeController:(ASRangeController *)rangeController
+{
+  if (CGRectIsEmpty(self.bounds)) {
+    return @[];
+  }
+
+  ASElementMap *map = _dataController.visibleMap;
+  NSMutableArray<ASCollectionElement *> *result = [NSMutableArray array];
+
+  // Visible items
+  for (NSIndexPath *indexPath in self.indexPathsForVisibleItems) {
+    ASCollectionElement *element = [map elementForItemAtIndexPath:indexPath];
+    [result addObject:element];
+  }
+
+  // Visible supplementary elements
+  for (NSString *kind in map.supplementaryElementKinds) {
+    for (NSIndexPath *indexPath in [self asdk_indexPathsForVisibleSupplementaryElementsOfKind:kind]) {
+      ASCollectionElement *element = [map supplementaryElementOfKind:kind atIndexPath:indexPath];
+      [result addObject:element];
+    }
+  }
+  return result;
 }
 
 - (ASElementMap *)elementMapForRangeController:(ASRangeController *)rangeController

Source/ASTableView.mm

@@ -649,18 +649,8 @@ static NSString * const kCellReuseIdentifier = @"_ASTableViewCell";
 
 - (NSArray<ASCellNode *> *)visibleNodes
 {
-  NSArray *indexPaths = [self visibleNodeIndexPathsForRangeController:_rangeController];
+  NSArray<ASCollectionElement *> *elements = [self visibleElementsForRangeController:_rangeController];
-  
+  return ASArrayByFlatMapping(elements, ASCollectionElement *e, e.node);
-  NSMutableArray<ASCellNode *> *visibleNodes = [NSMutableArray array];
-  for (NSIndexPath *indexPath in indexPaths) {
-    ASCellNode *node = [self nodeForRowAtIndexPath:indexPath];
-    if (node) {
-      // It is possible for UITableView to return indexPaths before the node is completed.
-      [visibleNodes addObject:node];
-    }
-  }
-  
-  return visibleNodes;
 }
 
 - (void)beginUpdates
@@ -1358,7 +1348,7 @@ static NSString * const kCellReuseIdentifier = @"_ASTableViewCell";
     return _rangeController;
 }
 
-- (NSArray *)visibleNodeIndexPathsForRangeController:(ASRangeController *)rangeController
+- (NSArray<ASCollectionElement *> *)visibleElementsForRangeController:(ASRangeController *)rangeController
 {
   ASDisplayNodeAssertMainThread();
   
@@ -1383,7 +1373,8 @@ static NSString * const kCellReuseIdentifier = @"_ASTableViewCell";
     }]];
   }
 
-  return visibleIndexPaths;
+  ASElementMap *map = _dataController.visibleMap;
+  return ASArrayByFlatMapping(visibleIndexPaths, NSIndexPath *indexPath, [map elementForItemAtIndexPath:indexPath]);
 }
 
 - (ASScrollDirection)scrollDirectionForRangeController:(ASRangeController *)rangeController

Source/Base/ASBaseDefines.h

@@ -217,3 +217,31 @@
   id __val = x;\
   ((c *) ([__val isKindOfClass:[c class]] ? __val : nil));\
 })
+
+/**
+ * Create a new set by mapping `collection` over `work`, ignoring nil.
+ */
+#define ASSetByFlatMapping(collection, decl, work) ({ \
+  NSMutableSet *s = [NSMutableSet set]; \
+  for (decl in collection) {\
+    id result = work; \
+    if (result != nil) { \
+      [s addObject:result]; \
+    } \
+  } \
+  s; \
+})
+
+/**
+ * Create a new array by mapping `collection` over `work`, ignoring nil.
+ */
+#define ASArrayByFlatMapping(collection, decl, work) ({ \
+  NSMutableArray *a = [NSMutableArray array]; \
+  for (decl in collection) {\
+    id result = work; \
+    if (result != nil) { \
+      [a addObject:result]; \
+    } \
+  } \
+  a; \
+})

Source/Details/ASAbstractLayoutController.mm

@@ -167,13 +167,13 @@ extern CGRect CGRectExpandToRangeWithScrollableDirections(CGRect rect, ASRangeTu
 
 #pragma mark - Abstract Index Path Range Support
 
-- (NSSet *)indexPathsForScrolling:(ASScrollDirection)scrollDirection rangeMode:(ASLayoutRangeMode)rangeMode rangeType:(ASLayoutRangeType)rangeType
+- (NSSet<ASCollectionElement *> *)elementsForScrolling:(ASScrollDirection)scrollDirection rangeMode:(ASLayoutRangeMode)rangeMode rangeType:(ASLayoutRangeType)rangeType map:(ASElementMap *)map
 {
   ASDisplayNodeAssertNotSupported();
   return nil;
 }
 
-- (void)allIndexPathsForScrolling:(ASScrollDirection)scrollDirection rangeMode:(ASLayoutRangeMode)rangeMode displaySet:(NSSet *__autoreleasing  _Nullable *)displaySet preloadSet:(NSSet *__autoreleasing  _Nullable *)preloadSet
+- (void)allElementsForScrolling:(ASScrollDirection)scrollDirection rangeMode:(ASLayoutRangeMode)rangeMode displaySet:(NSSet<ASCollectionElement *> *__autoreleasing  _Nullable *)displaySet preloadSet:(NSSet<ASCollectionElement *> *__autoreleasing  _Nullable *)preloadSet map:(ASElementMap *)map
 {
   ASDisplayNodeAssertNotSupported();
 }

Source/Details/ASCollectionViewLayoutController.m

@@ -12,6 +12,7 @@
 
 #import <AsyncDisplayKit/ASAssert.h>
 #import <AsyncDisplayKit/ASCollectionView.h>
+#import <AsyncDisplayKit/ASElementMap.h>
 #import <AsyncDisplayKit/CoreGraphics+ASConvenience.h>
 #import <AsyncDisplayKit/UICollectionViewLayout+ASConvenience.h>
 
@@ -46,14 +47,14 @@ typedef struct ASRangeGeometry ASRangeGeometry;
   return self;
 }
 
-- (NSSet *)indexPathsForScrolling:(ASScrollDirection)scrollDirection rangeMode:(ASLayoutRangeMode)rangeMode rangeType:(ASLayoutRangeType)rangeType
+- (NSSet<ASCollectionElement *> *)elementsForScrolling:(ASScrollDirection)scrollDirection rangeMode:(ASLayoutRangeMode)rangeMode rangeType:(ASLayoutRangeType)rangeType map:(ASElementMap *)map
 {
   ASRangeTuningParameters tuningParameters = [self tuningParametersForRangeMode:rangeMode rangeType:rangeType];
   CGRect rangeBounds = [self rangeBoundsWithScrollDirection:scrollDirection rangeTuningParameters:tuningParameters];
-  return [self indexPathsForItemsWithinRangeBounds:rangeBounds];
+  return [self elementsWithinRangeBounds:rangeBounds map:map];
 }
 
-- (void)allIndexPathsForScrolling:(ASScrollDirection)scrollDirection rangeMode:(ASLayoutRangeMode)rangeMode displaySet:(NSSet **)displaySet preloadSet:(NSSet **)preloadSet
+- (void)allElementsForScrolling:(ASScrollDirection)scrollDirection rangeMode:(ASLayoutRangeMode)rangeMode displaySet:(NSSet<ASCollectionElement *> *__autoreleasing  _Nullable *)displaySet preloadSet:(NSSet<ASCollectionElement *> *__autoreleasing  _Nullable *)preloadSet map:(ASElementMap *)map
 {
   if (displaySet == NULL || preloadSet == NULL) {
     return;
@@ -67,12 +68,12 @@ typedef struct ASRangeGeometry ASRangeGeometry;
   CGRect unionBounds = CGRectUnion(displayBounds, preloadBounds);
   NSArray *layoutAttributes = [_collectionViewLayout layoutAttributesForElementsInRect:unionBounds];
 
-  NSMutableSet *display = [NSMutableSet setWithCapacity:layoutAttributes.count];
+  NSMutableSet<ASCollectionElement *> *display = [NSMutableSet setWithCapacity:layoutAttributes.count];
-  NSMutableSet *preload = [NSMutableSet setWithCapacity:layoutAttributes.count];
+  NSMutableSet<ASCollectionElement *> *preload = [NSMutableSet setWithCapacity:layoutAttributes.count];
 
   for (UICollectionViewLayoutAttributes *la in layoutAttributes) {
     // Manually filter out elements that don't intersect the range bounds.
-    // See comment in indexPathsForItemsWithinRangeBounds:
+    // See comment in elementsForItemsWithinRangeBounds:
     // This is re-implemented here so that the iteration over layoutAttributes can be done once to check both ranges.
     CGRect frame = la.frame;
     BOOL intersectsDisplay = CGRectIntersectsRect(displayBounds, frame);
@@ -82,13 +83,13 @@ typedef struct ASRangeGeometry ASRangeGeometry;
       continue;
     }
     
-    // Avoid excessive retains and releases, as well as property calls. We know the indexPath is kept alive by la.
+    // Avoid excessive retains and releases, as well as property calls. We know the element is kept alive by map.
-    __unsafe_unretained NSIndexPath *indexPath = la.indexPath;
+    __unsafe_unretained ASCollectionElement *e = [map elementForLayoutAttributes:la];
     if (intersectsDisplay) {
-      [display addObject:indexPath];
+      [display addObject:e];
     }
     if (intersectsPreload) {
-      [preload addObject:indexPath];
+      [preload addObject:e];
     }
   }
 
@@ -97,14 +98,12 @@ typedef struct ASRangeGeometry ASRangeGeometry;
   return;
 }
 
-- (NSSet *)indexPathsForItemsWithinRangeBounds:(CGRect)rangeBounds
+- (NSSet<ASCollectionElement *> *)elementsWithinRangeBounds:(CGRect)rangeBounds map:(ASElementMap *)map
 {
   NSArray *layoutAttributes = [_collectionViewLayout layoutAttributesForElementsInRect:rangeBounds];
-  NSMutableSet *indexPathSet = [NSMutableSet setWithCapacity:layoutAttributes.count];
+  NSMutableSet<ASCollectionElement *> *elementSet = [NSMutableSet setWithCapacity:layoutAttributes.count];
   
   for (UICollectionViewLayoutAttributes *la in layoutAttributes) {
-    //ASDisplayNodeAssert(![indexPathSet containsObject:la.indexPath], @"Shouldn't already contain indexPath");
-
     // Manually filter out elements that don't intersect the range bounds.
     // If a layout returns elements outside the requested rect this can be a huge problem.
     // For instance in a paging flow, you may only want to preload 3 pages (one center, one on each side)
@@ -113,10 +112,10 @@ typedef struct ASRangeGeometry ASRangeGeometry;
     if (CATransform3DIsIdentity(la.transform3D) && CGRectIntersectsRect(la.frame, rangeBounds) == NO) {
       continue;
     }
-    [indexPathSet addObject:la.indexPath];
+    [elementSet addObject:[map elementForLayoutAttributes:la]];
   }
 
-  return indexPathSet;
+  return elementSet;
 }
 
 - (CGRect)rangeBoundsWithScrollDirection:(ASScrollDirection)scrollDirection

Source/Details/ASLayoutController.h

@@ -16,7 +16,7 @@
 
 NS_ASSUME_NONNULL_BEGIN
 
-@class ASCellNode;
+@class ASCollectionElement, ASElementMap;
 
 ASDISPLAYNODE_EXTERN_C_BEGIN
 
@@ -34,9 +34,9 @@ ASDISPLAYNODE_EXTERN_C_END
 
 - (ASRangeTuningParameters)tuningParametersForRangeMode:(ASLayoutRangeMode)rangeMode rangeType:(ASLayoutRangeType)rangeType;
 
-- (NSSet<NSIndexPath *> *)indexPathsForScrolling:(ASScrollDirection)scrollDirection rangeMode:(ASLayoutRangeMode)rangeMode rangeType:(ASLayoutRangeType)rangeType;
+- (NSSet<ASCollectionElement *> *)elementsForScrolling:(ASScrollDirection)scrollDirection rangeMode:(ASLayoutRangeMode)rangeMode rangeType:(ASLayoutRangeType)rangeType map:(ASElementMap *)map;
 
-- (void)allIndexPathsForScrolling:(ASScrollDirection)scrollDirection rangeMode:(ASLayoutRangeMode)rangeMode displaySet:(NSSet * _Nullable * _Nullable)displaySet preloadSet:(NSSet * _Nullable * _Nullable)preloadSet;
+- (void)allElementsForScrolling:(ASScrollDirection)scrollDirection rangeMode:(ASLayoutRangeMode)rangeMode displaySet:(NSSet<ASCollectionElement *> * _Nullable * _Nullable)displaySet preloadSet:(NSSet<ASCollectionElement *> * _Nullable * _Nullable)preloadSet map:(ASElementMap *)map;
 
 @optional
 

Source/Details/ASRangeController.h

@@ -107,9 +107,9 @@ AS_SUBCLASSING_RESTRICTED
 /**
  * @param rangeController Sender.
  *
- * @return an array of index paths corresponding to the nodes currently visible onscreen (i.e., the visible range).
+ * @return an array of elements corresponding to the data currently visible onscreen (i.e., the visible range).
  */
-- (NSArray<NSIndexPath *> *)visibleNodeIndexPathsForRangeController:(ASRangeController *)rangeController;
+- (NSArray<ASCollectionElement *> *)visibleElementsForRangeController:(ASRangeController *)rangeController;
 
 /**
  * @param rangeController Sender.

Source/Details/ASRangeController.mm

@@ -209,10 +209,10 @@ static UIApplicationState __ApplicationState = UIApplicationStateActive;
 
   // TODO: Consider if we need to use this codepath, or can rely on something more similar to the data & display ranges
   // Example: ... = [_layoutController indexPathsForScrolling:scrollDirection rangeType:ASLayoutRangeTypeVisible];
-  NSArray<NSIndexPath *> *visibleNodePaths = [_dataSource visibleNodeIndexPathsForRangeController:self];
+  NSSet<ASCollectionElement *> *visibleElements = [NSSet setWithArray:[_dataSource visibleElementsForRangeController:self]];
   ASWeakSet *newVisibleNodes = [[ASWeakSet alloc] init];
 
-  if (visibleNodePaths.count == 0) { // if we don't have any visibleNodes currently (scrolled before or after content)...
+  if (visibleElements.count == 0) { // if we don't have any visibleNodes currently (scrolled before or after content)...
     [self _setVisibleNodes:newVisibleNodes];
     return; // don't do anything for this update, but leave _rangeIsValid == NO to make sure we update it later
   }
@@ -248,32 +248,36 @@ static UIApplicationState __ApplicationState = UIApplicationStateActive;
   // Check if both Display and Preload are unique. If they are, we load them with a single fetch from the layout controller for performance.
   BOOL optimizedLoadingOfBothRanges = (equalDisplayPreload == NO && equalDisplayVisible == NO && emptyDisplayRange == NO);
 
-  NSSet<NSIndexPath *> *visibleIndexPaths = [NSSet setWithArray:visibleNodePaths];
+  NSSet<ASCollectionElement *> *displayElements = nil;
-  NSSet<NSIndexPath *> *displayIndexPaths = nil;
+  NSSet<ASCollectionElement *> *preloadElements = nil;
-  NSSet<NSIndexPath *> *preloadIndexPaths = nil;
   
   if (optimizedLoadingOfBothRanges) {
-    [_layoutController allIndexPathsForScrolling:scrollDirection rangeMode:rangeMode displaySet:&displayIndexPaths preloadSet:&preloadIndexPaths];
+    [_layoutController allElementsForScrolling:scrollDirection rangeMode:rangeMode displaySet:&displayElements preloadSet:&preloadElements map:map];
   } else {
     if (emptyDisplayRange == YES) {
-      displayIndexPaths = [NSSet set];
+      displayElements = [NSSet set];
     } if (equalDisplayVisible == YES) {
-      displayIndexPaths = visibleIndexPaths;
+      displayElements = visibleElements;
     } else {
       // Calculating only the Display range means the Preload range is either the same as Display or Visible.
-      displayIndexPaths = [_layoutController indexPathsForScrolling:scrollDirection rangeMode:rangeMode rangeType:ASLayoutRangeTypeDisplay];
+      displayElements = [_layoutController elementsForScrolling:scrollDirection rangeMode:rangeMode rangeType:ASLayoutRangeTypeDisplay map:map];
     }
     
     BOOL equalPreloadVisible = ASRangeTuningParametersEqualToRangeTuningParameters(parametersPreload, ASRangeTuningParametersZero);
     if (equalDisplayPreload == YES) {
-      preloadIndexPaths = displayIndexPaths;
+      preloadElements = displayElements;
     } else if (equalPreloadVisible == YES) {
-      preloadIndexPaths = visibleIndexPaths;
+      preloadElements = visibleElements;
     } else {
-      preloadIndexPaths = [_layoutController indexPathsForScrolling:scrollDirection rangeMode:rangeMode rangeType:ASLayoutRangeTypePreload];
+      preloadElements = [_layoutController elementsForScrolling:scrollDirection rangeMode:rangeMode rangeType:ASLayoutRangeTypePreload map:map];
     }
   }
   
+  // For now we are only interested in items. Filter-map out from element to item-index-path.
+  NSSet<NSIndexPath *> *visibleIndexPaths = ASSetByFlatMapping(visibleElements, ASCollectionElement *element, [map indexPathForElementIfCell:element]);
+  NSSet<NSIndexPath *> *displayIndexPaths = ASSetByFlatMapping(displayElements, ASCollectionElement *element, [map indexPathForElementIfCell:element]);
+  NSSet<NSIndexPath *> *preloadIndexPaths = ASSetByFlatMapping(preloadElements, ASCollectionElement *element, [map indexPathForElementIfCell:element]);
+
   // Prioritize the order in which we visit each.  Visible nodes should be updated first so they are enqueued on
   // the network or display queues before preloading (offscreen) nodes are enqueued.
   NSMutableOrderedSet<NSIndexPath *> *allIndexPaths = [[NSMutableOrderedSet alloc] initWithSet:visibleIndexPaths];

Source/Details/ASTableLayoutController.m

@@ -13,6 +13,7 @@
 #import <UIKit/UIKit.h>
 
 #import <AsyncDisplayKit/ASAssert.h>
+#import <AsyncDisplayKit/ASElementMap.h>
 
 @interface ASTableLayoutController()
 @end
@@ -30,28 +31,24 @@
 
 #pragma mark - ASLayoutController
 
-/**
+- (NSSet<ASCollectionElement *> *)elementsForScrolling:(ASScrollDirection)scrollDirection rangeMode:(ASLayoutRangeMode)rangeMode rangeType:(ASLayoutRangeType)rangeType map:(ASElementMap *)map
- * IndexPath array for the element in the working range.
- */
-
-- (NSSet *)indexPathsForScrolling:(ASScrollDirection)scrollDirection rangeMode:(ASLayoutRangeMode)rangeMode rangeType:(ASLayoutRangeType)rangeType
 {
   CGRect bounds = _tableView.bounds;
 
   ASRangeTuningParameters tuningParameters = [self tuningParametersForRangeMode:rangeMode rangeType:rangeType];
   CGRect rangeBounds = CGRectExpandToRangeWithScrollableDirections(bounds, tuningParameters, ASScrollDirectionVerticalDirections, scrollDirection);
   NSArray *array = [_tableView indexPathsForRowsInRect:rangeBounds];
-  return [NSSet setWithArray:array];
+  return ASSetByFlatMapping(array, NSIndexPath *indexPath, [map elementForItemAtIndexPath:indexPath]);
 }
 
-- (void)allIndexPathsForScrolling:(ASScrollDirection)scrollDirection rangeMode:(ASLayoutRangeMode)rangeMode displaySet:(NSSet **)displaySet preloadSet:(NSSet **)preloadSet
+- (void)allElementsForScrolling:(ASScrollDirection)scrollDirection rangeMode:(ASLayoutRangeMode)rangeMode displaySet:(NSSet<ASCollectionElement *> *__autoreleasing  _Nullable *)displaySet preloadSet:(NSSet<ASCollectionElement *> *__autoreleasing  _Nullable *)preloadSet map:(ASElementMap *)map
 {
   if (displaySet == NULL || preloadSet == NULL) {
     return;
   }
 
-  *displaySet = [self indexPathsForScrolling:scrollDirection rangeMode:rangeMode rangeType:ASLayoutRangeTypeDisplay];
+  *displaySet = [self elementsForScrolling:scrollDirection rangeMode:rangeMode rangeType:ASLayoutRangeTypeDisplay map:map];
-  *preloadSet = [self indexPathsForScrolling:scrollDirection rangeMode:rangeMode rangeType:ASLayoutRangeTypePreload];
+  *preloadSet = [self elementsForScrolling:scrollDirection rangeMode:rangeMode rangeType:ASLayoutRangeTypePreload map:map];
   return;
 }
 

Source/Private/ASElementMap.h

@@ -11,7 +11,7 @@
 
 NS_ASSUME_NONNULL_BEGIN
 
-@class ASCollectionElement, ASSection;
+@class ASCollectionElement, ASSection, UICollectionViewLayoutAttributes;
 @protocol ASSectionContext;
 
 /**
@@ -27,6 +27,11 @@ AS_SUBCLASSING_RESTRICTED
  */
 @property (readonly) NSInteger numberOfSections;
 
+/**
+ * The kinds of supplementary elements present in this map. O(1)
+ */
+@property (copy, readonly) NSArray<NSString *> *supplementaryElementKinds;
+
 /**
  * Returns number of items in the given section. O(1)
  */
@@ -54,6 +59,11 @@ AS_SUBCLASSING_RESTRICTED
  */
 - (nullable NSIndexPath *)indexPathForElement:(ASCollectionElement *)element;
 
+/**
+ * Returns the index path for the given element, if it represents a cell. O(1)
+ */
+- (nullable NSIndexPath *)indexPathForElementIfCell:(ASCollectionElement *)element;
+
 /**
  * Returns the item-element at the given index path. O(1)
  */
@@ -64,6 +74,13 @@ AS_SUBCLASSING_RESTRICTED
  */
 - (nullable ASCollectionElement *)supplementaryElementOfKind:(NSString *)supplementaryElementKind atIndexPath:(NSIndexPath *)indexPath;
 
+/**
+ * Returns the element that corresponds to the given layout attributes, if any.
+ *
+ * NOTE: This method only regards the category, kind, and index path of the attributes object. Elements do not
+ * have any concept of size/position.
+ */
+- (nullable ASCollectionElement *)elementForLayoutAttributes:(UICollectionViewLayoutAttributes *)layoutAttributes;
 
 #pragma mark - Initialization -- Only Useful to ASDataController
 

Source/Private/ASElementMap.m

@@ -73,6 +73,11 @@
   return _sectionsOfItems.count;
 }
 
+- (NSArray<NSString *> *)supplementaryElementKinds
+{
+  return _supplementaryElements.allKeys;
+}
+
 - (NSInteger)numberOfItemsInSection:(NSInteger)section
 {
   return _sectionsOfItems[section].count;
@@ -88,6 +93,15 @@
   return [_elementToIndexPathMap objectForKey:element];
 }
 
+- (nullable NSIndexPath *)indexPathForElementIfCell:(ASCollectionElement *)element
+{
+  if (element.supplementaryElementKind == nil) {
+    return [self indexPathForElement:element];
+  } else {
+    return nil;
+  }
+}
+
 - (nullable ASCollectionElement *)elementForItemAtIndexPath:(NSIndexPath *)indexPath
 {
   return (indexPath != nil) ? ASGetElementInTwoDimensionalArray(_sectionsOfItems, indexPath) : nil;
@@ -98,6 +112,21 @@
   return _supplementaryElements[supplementaryElementKind][indexPath];
 }
 
+- (ASCollectionElement *)elementForLayoutAttributes:(UICollectionViewLayoutAttributes *)layoutAttributes
+{
+  switch (layoutAttributes.representedElementCategory) {
+    case UICollectionElementCategoryCell:
+      // Cell
+      return [self elementForItemAtIndexPath:layoutAttributes.indexPath];
+    case UICollectionElementCategorySupplementaryView:
+      // Supplementary element.
+      return [self supplementaryElementOfKind:layoutAttributes.representedElementKind atIndexPath:layoutAttributes.indexPath];
+    case UICollectionElementCategoryDecorationView:
+      // No support for decoration views.
+      return nil;
+  }
+}
+
 - (NSIndexPath *)convertIndexPath:(NSIndexPath *)indexPath fromMap:(ASElementMap *)map
 {
   id element = [map elementForItemAtIndexPath:indexPath];