Rejigger Cell Visibility Tracking

Source/ASCollectionView.mm

@@ -105,6 +105,9 @@ @interface ASCollectionView () <ASRangeControllerDataSource, ASRangeControllerDe
 
   NSMutableSet *_registeredSupplementaryKinds;
 
+  // CountedSet because UIKit may display the same element in multiple cells e.g. during animations.
+  NSCountedSet<ASCollectionElement *> *_visibleElements;
+
   CGPoint _deceleratingVelocity;
 
   BOOL _zeroContentInsets;
@@ -294,6 +297,7 @@ - (instancetype)_initWithFrame:(CGRect)frame collectionViewLayout:(UICollectionV
   super.dataSource = (id<UICollectionViewDataSource>)_proxyDataSource;
 
   _registeredSupplementaryKinds = [NSMutableSet set];
+  _visibleElements = [[NSCountedSet alloc] init];
 
   _cellsForVisibilityUpdates = [NSMutableSet set];
   _cellsForLayoutUpdates = [NSMutableSet set];
@@ -997,7 +1001,8 @@ - (UICollectionReusableView *)collectionView:(UICollectionView *)collectionView
   }
 
   UICollectionReusableView *view = nil;
-  ASCellNode *node = [_dataController.visibleMap supplementaryElementOfKind:kind atIndexPath:indexPath].node;
+  ASCollectionElement *element = [_dataController.visibleMap supplementaryElementOfKind:kind atIndexPath:indexPath];
+  ASCellNode *node = element.node;
 
   BOOL shouldDequeueExternally = _asyncDataSourceFlags.interopViewForSupplementaryElement && (_asyncDataSourceFlags.interopAlwaysDequeue || node.shouldUseUIKitCell);
   if (shouldDequeueExternally) {
@@ -1009,7 +1014,7 @@ - (UICollectionReusableView *)collectionView:(UICollectionView *)collectionView
   }
 
   if (_ASCollectionReusableView *reusableView = ASDynamicCast(view, _ASCollectionReusableView)) {
-    reusableView.node = node;
+    reusableView.element = element;
   }
 
   if (node) {
@@ -1022,7 +1027,8 @@ - (UICollectionReusableView *)collectionView:(UICollectionView *)collectionView
 - (UICollectionViewCell *)collectionView:(UICollectionView *)collectionView cellForItemAtIndexPath:(NSIndexPath *)indexPath
 {
   UICollectionViewCell *cell = nil;
-  ASCellNode *node = [self nodeForItemAtIndexPath:indexPath];
+  ASCollectionElement *element = [_dataController.visibleMap elementForItemAtIndexPath:indexPath];
+  ASCellNode *node = element.node;
 
   BOOL shouldDequeueExternally = _asyncDataSourceFlags.interopAlwaysDequeue || (_asyncDataSourceFlags.interop && node.shouldUseUIKitCell);
   if (shouldDequeueExternally) {
@@ -1031,30 +1037,33 @@ - (UICollectionViewCell *)collectionView:(UICollectionView *)collectionView cell
     cell = [self dequeueReusableCellWithReuseIdentifier:kReuseIdentifier forIndexPath:indexPath];
   }
 
-  ASDisplayNodeAssert(node != nil, @"Cell node should exist. indexPath = %@, collectionDataSource = %@", indexPath, self);
+  ASDisplayNodeAssert(element != nil, @"Element should exist. indexPath = %@, collectionDataSource = %@", indexPath, self);
 
   if (_ASCollectionViewCell *asCell = ASDynamicCast(cell, _ASCollectionViewCell)) {
-    asCell.node = node;
+    asCell.element = element;
     [_rangeController configureContentView:cell.contentView forCellNode:node];
   }
 
   return cell;
 }
 
-- (void)collectionView:(UICollectionView *)collectionView willDisplayCell:(_ASCollectionViewCell *)cell forItemAtIndexPath:(NSIndexPath *)indexPath
+- (void)collectionView:(UICollectionView *)collectionView willDisplayCell:(UICollectionViewCell *)rawCell forItemAtIndexPath:(NSIndexPath *)indexPath
 {
   if (_asyncDelegateFlags.interopWillDisplayCell) {
-    [(id <ASCollectionDelegateInterop>)_asyncDelegate collectionView:collectionView willDisplayCell:cell forItemAtIndexPath:indexPath];
+    [(id <ASCollectionDelegateInterop>)_asyncDelegate collectionView:collectionView willDisplayCell:rawCell forItemAtIndexPath:indexPath];
   }
 
   // Since _ASCollectionViewCell is not available for subclassing, this is faster than isKindOfClass:
   // We must exit early here, because only _ASCollectionViewCell implements the -node accessor method.
-  if ([cell class] != [_ASCollectionViewCell class]) {
+  if ([rawCell class] != [_ASCollectionViewCell class]) {
     [_rangeController setNeedsUpdate];
     return;
   }
+  auto cell = (_ASCollectionViewCell *)rawCell;
 
-  ASCellNode *cellNode = [cell node];
+  ASCollectionElement *element = cell.element;
+  [_visibleElements addObject:element];
+  ASCellNode *cellNode = element.node;
   cellNode.scrollView = collectionView;
 
   // Update the selected background view in collectionView:willDisplayCell:forItemAtIndexPath: otherwise it could be to
@@ -1090,21 +1099,24 @@ - (void)collectionView:(UICollectionView *)collectionView willDisplayCell:(_ASCo
   }
 }
 
-- (void)collectionView:(UICollectionView *)collectionView didEndDisplayingCell:(_ASCollectionViewCell *)cell forItemAtIndexPath:(NSIndexPath *)indexPath
+- (void)collectionView:(UICollectionView *)collectionView didEndDisplayingCell:(UICollectionViewCell *)rawCell forItemAtIndexPath:(NSIndexPath *)indexPath
 {
   if (_asyncDelegateFlags.interopDidEndDisplayingCell) {
-    [(id <ASCollectionDelegateInterop>)_asyncDelegate collectionView:collectionView didEndDisplayingCell:cell forItemAtIndexPath:indexPath];
+    [(id <ASCollectionDelegateInterop>)_asyncDelegate collectionView:collectionView didEndDisplayingCell:rawCell forItemAtIndexPath:indexPath];
   }
 
   // Since _ASCollectionViewCell is not available for subclassing, this is faster than isKindOfClass:
   // We must exit early here, because only _ASCollectionViewCell implements the -node accessor method.
-  if ([cell class] != [_ASCollectionViewCell class]) {
+  if ([rawCell class] != [_ASCollectionViewCell class]) {
     [_rangeController setNeedsUpdate];
     return;
   }
+  auto cell = (_ASCollectionViewCell *)rawCell;
 
-  ASCellNode *cellNode = [cell node];
-  ASDisplayNodeAssertNotNil(cellNode, @"Expected node associated with removed cell not to be nil.");
+  ASCollectionElement *element = cell.element;
+  [_visibleElements removeObject:element];
+  ASDisplayNodeAssertNotNil(element, @"Expected element associated with removed cell not to be nil.");
+  ASCellNode *cellNode = element.node;
 
   if (_asyncDelegateFlags.collectionNodeDidEndDisplayingItem) {
     if (ASCollectionNode *collectionNode = self.collectionNode) {
@@ -1125,8 +1137,14 @@ - (void)collectionView:(UICollectionView *)collectionView didEndDisplayingCell:(
   cell.layoutAttributes = nil;
 }
 
-- (void)collectionView:(UICollectionView *)collectionView willDisplaySupplementaryView:(_ASCollectionReusableView *)view forElementKind:(NSString *)elementKind atIndexPath:(NSIndexPath *)indexPath
+- (void)collectionView:(UICollectionView *)collectionView willDisplaySupplementaryView:(UICollectionReusableView *)rawView forElementKind:(NSString *)elementKind atIndexPath:(NSIndexPath *)indexPath
 {
+  if (rawView.class != [_ASCollectionReusableView class]) {
+    return;
+  }
+  auto view = (_ASCollectionReusableView *)rawView;
+
+  [_visibleElements addObject:view.element];
   // This is a safeguard similar to the behavior for cells in -[ASCollectionView collectionView:willDisplayCell:forItemAtIndexPath:]
   // It ensures _ASCollectionReusableView receives layoutAttributes and calls applyLayoutAttributes.
   if (view.layoutAttributes == nil) {
@@ -1141,8 +1159,14 @@ - (void)collectionView:(UICollectionView *)collectionView willDisplaySupplementa
   }
 }
 
-- (void)collectionView:(UICollectionView *)collectionView didEndDisplayingSupplementaryView:(UICollectionReusableView *)view forElementOfKind:(NSString *)elementKind atIndexPath:(NSIndexPath *)indexPath
+- (void)collectionView:(UICollectionView *)collectionView didEndDisplayingSupplementaryView:(UICollectionReusableView *)rawView forElementOfKind:(NSString *)elementKind atIndexPath:(NSIndexPath *)indexPath
 {
+  if (rawView.class != [_ASCollectionReusableView class]) {
+    return;
+  }
+  auto view = (_ASCollectionReusableView *)rawView;
+
+  [_visibleElements removeObject:view.element];
   if (_asyncDelegateFlags.collectionNodeDidEndDisplayingSupplementaryElement) {
     GET_COLLECTIONNODE_OR_RETURN(collectionNode, (void)0);
     ASCellNode *node = [self supplementaryNodeForElementKind:elementKind atIndexPath:indexPath];
@@ -1805,35 +1829,7 @@ - (ASRangeController *)rangeController
 
 - (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];
-    if (element != nil) {
-      [result addObject:element];
-    } else {
-      ASDisplayNodeFailAssert(@"Couldn't find 'visible' item at index path %@ in map %@", indexPath, map);
-    }
-  }
-
-  // Visible supplementary elements
-  for (NSString *kind in map.supplementaryElementKinds) {
-    for (NSIndexPath *indexPath in [self asdk_indexPathsForVisibleSupplementaryElementsOfKind:kind]) {
-      ASCollectionElement *element = [map supplementaryElementOfKind:kind atIndexPath:indexPath];
-      if (element != nil) {
-        [result addObject:element];
-      } else {
-        ASDisplayNodeFailAssert(@"Couldn't find 'visible' supplementary element of kind %@ at index path %@ in map %@", kind, indexPath, map);
-      }
-    }
-  }
-  return result;
+  return _visibleElements.allObjects;
 }
 
 - (ASElementMap *)elementMapForRangeController:(ASRangeController *)rangeController

Source/ASDisplayNode.mm

@@ -404,14 +404,8 @@ - (void)dealloc
   _flags.isDeallocating = YES;
 
   // Synchronous nodes may not be able to call the hierarchy notifications, so only enforce for regular nodes.
-  // TODO: This condition should be an assertion, but a workaround is in place until the root issue is fixed:
-  // https://github.com/TextureGroup/Texture/issues/145
-#if DEBUG
-  if (checkFlag(Synchronous) == NO && ASInterfaceStateIncludesVisible(_interfaceState) == YES) {
-    NSLog(@"Node should always be marked invisible before deallocating. Node: %@", self);
-  }
-#endif
-
+  ASDisplayNodeAssert(checkFlag(Synchronous) || !ASInterfaceStateIncludesVisible(_interfaceState), @"Node should always be marked invisible before deallocating. Node: %@", self);
+
   self.asyncLayer.asyncDelegate = nil;
   _view.asyncdisplaykit_node = nil;
   _layer.asyncdisplaykit_node = nil;

Source/ASTableView.mm

@@ -82,7 +82,8 @@ - (void)didLayoutSubviewsOfTableViewCell:(_ASTableViewCell *)tableViewCell;
 
 @interface _ASTableViewCell : UITableViewCell
 @property (nonatomic, weak) id<_ASTableViewCellDelegate> delegate;
-@property (nonatomic, weak) ASCellNode *node;
+@property (nonatomic, strong, readonly) ASCellNode *node;
+@property (nonatomic, strong) ASCollectionElement *element;
 @end
 
 @implementation _ASTableViewCell
@@ -101,9 +102,15 @@ - (void)didTransitionToState:(UITableViewCellStateMask)state
   [super didTransitionToState:state];
 }
 
-- (void)setNode:(ASCellNode *)node
+- (ASCellNode *)node
 {
-  _node = node;
+  return self.element.node;
+}
+
+- (void)setElement:(ASCollectionElement *)element
+{
+  _element = element;
+  ASCellNode *node = element.node;
 
   if (node) {
     self.backgroundColor = node.backgroundColor;
@@ -126,19 +133,19 @@ - (void)setNode:(ASCellNode *)node
 - (void)setSelected:(BOOL)selected animated:(BOOL)animated
 {
   [super setSelected:selected animated:animated];
-  [_node __setSelectedFromUIKit:selected];
+  [self.node __setSelectedFromUIKit:selected];
 }
 
 - (void)setHighlighted:(BOOL)highlighted animated:(BOOL)animated
 {
   [super setHighlighted:highlighted animated:animated];
-  [_node __setHighlightedFromUIKit:highlighted];
+  [self.node __setHighlightedFromUIKit:highlighted];
 }
 
 - (void)prepareForReuse
 {
-  // Need to clear node pointer before UIKit calls setSelected:NO / setHighlighted:NO on its cells
-  self.node = nil;
+  // Need to clear element before UIKit calls setSelected:NO / setHighlighted:NO on its cells
+  self.element = nil;
   [super prepareForReuse];
 }
 
@@ -186,6 +193,9 @@ @interface ASTableView () <ASRangeControllerDataSource, ASRangeControllerDelegat
   BOOL _isDeallocating;
   NSMutableSet *_cellsForVisibilityUpdates;
 
+  // CountedSet because UIKit may display the same element in multiple cells e.g. during animations.
+  NSCountedSet<ASCollectionElement *> *_visibleElements;
+
   BOOL _remeasuringCellNodes;
   NSMutableSet *_cellsForLayoutUpdates;
 
@@ -317,6 +327,7 @@ - (instancetype)_initWithFrame:(CGRect)frame style:(UITableViewStyle)style dataC
 
   _leadingScreensForBatching = 2.0;
   _batchContext = [[ASBatchContext alloc] init];
+  _visibleElements = [[NSCountedSet alloc] init];
 
   _automaticallyAdjustsContentOffset = NO;
 
@@ -898,11 +909,11 @@ - (UITableViewCell *)tableView:(UITableView *)tableView cellForRowAtIndexPath:(N
   _ASTableViewCell *cell = [self dequeueReusableCellWithIdentifier:kCellReuseIdentifier forIndexPath:indexPath];
   cell.delegate = self;
 
-  ASCellNode *node = [_dataController.visibleMap elementForItemAtIndexPath:indexPath].node;
+  ASCollectionElement *element = [_dataController.visibleMap elementForItemAtIndexPath:indexPath];
+  cell.element = element;
+  ASCellNode *node = element.node;
   if (node) {
     [_rangeController configureContentView:cell.contentView forCellNode:node];
-
-    cell.node = node;
   }
 
   return cell;
@@ -965,7 +976,9 @@ - (void)tableView:(UITableView *)tableView moveRowAtIndexPath:(NSIndexPath *)sou
 
 - (void)tableView:(UITableView *)tableView willDisplayCell:(_ASTableViewCell *)cell forRowAtIndexPath:(NSIndexPath *)indexPath
 {
-  ASCellNode *cellNode = [cell node];
+  ASCollectionElement *element = cell.element;
+  [_visibleElements addObject:element];
+  ASCellNode *cellNode = element.node;
   cellNode.scrollView = tableView;
 
   ASDisplayNodeAssertNotNil(cellNode, @"Expected node associated with cell that will be displayed not to be nil. indexPath: %@", indexPath);
@@ -991,7 +1004,9 @@ - (void)tableView:(UITableView *)tableView willDisplayCell:(_ASTableViewCell *)c
 
 - (void)tableView:(UITableView *)tableView didEndDisplayingCell:(_ASTableViewCell *)cell forRowAtIndexPath:(NSIndexPath *)indexPath
 {
-  ASCellNode *cellNode = [cell node];
+  ASCollectionElement *element = cell.element;
+  [_visibleElements removeObject:element];
+  ASCellNode *cellNode = element.node;
 
   [_rangeController setNeedsUpdate];
 
@@ -1439,36 +1454,12 @@ - (void)_beginBatchFetching
 
 - (ASRangeController *)rangeController
 {
-    return _rangeController;
+  return _rangeController;
 }
 
 - (NSArray<ASCollectionElement *> *)visibleElementsForRangeController:(ASRangeController *)rangeController
 {
-  ASDisplayNodeAssertMainThread();
-
-  CGRect bounds = self.bounds;
-  // Calling indexPathsForVisibleRows will trigger UIKit to call reloadData if it never has, which can result
-  // in incorrect layout if performed at zero size.  We can use the fact that nothing can be visible at zero size to return fast.
-  if (CGRectIsEmpty(bounds)) {
-    return @[];
-  }
-
-  NSArray *visibleIndexPaths = self.indexPathsForVisibleRows;
-
-  // In some cases (grouped-style tables with particular geometry) indexPathsForVisibleRows will return extra index paths.
-  // This is a very serious issue because we rely on the fact that any node that is marked Visible is hosted inside of a cell,
-  // or else we may not mark it invisible before the node is released. See testIssue2252.
-  // Calling indexPathForCell: and cellForRowAtIndexPath: are both pretty expensive – this is the quickest approach we have.
-  // It would be possible to cache this NSPredicate as an ivar, but that would require unsafeifying self and calling @c bounds
-  // for each item. Since the performance cost is pretty small, prefer simplicity.
-  if (self.style == UITableViewStyleGrouped && visibleIndexPaths.count != self.visibleCells.count) {
-    visibleIndexPaths = [visibleIndexPaths filteredArrayUsingPredicate:[NSPredicate predicateWithBlock:^BOOL(NSIndexPath *indexPath, NSDictionary<NSString *,id> * _Nullable bindings) {
-      return CGRectIntersectsRect(bounds, [self rectForRowAtIndexPath:indexPath]);
-    }]];
-  }
-
-  ASElementMap *map = _dataController.visibleMap;
-  return ASArrayByFlatMapping(visibleIndexPaths, NSIndexPath *indexPath, [map elementForItemAtIndexPath:indexPath]);
+  return _visibleElements.allObjects;
 }
 
 - (ASScrollDirection)scrollDirectionForRangeController:(ASRangeController *)rangeController

Source/Details/ASRangeController.mm

@@ -42,7 +42,7 @@ @interface ASRangeController ()
   BOOL _rangeIsValid;
   BOOL _needsRangeUpdate;
   NSSet<NSIndexPath *> *_allPreviousIndexPaths;
-  ASWeakSet<ASCellNode *> *_visibleNodes;
+  NSHashTable<ASCellNode *> *_visibleNodes;
   ASLayoutRangeMode _currentRangeMode;
   BOOL _preserveCurrentRangeMode;
   BOOL _didRegisterForNodeDisplayNotifications;
@@ -191,7 +191,7 @@ - (void)setDataSource:(id<ASRangeControllerDataSource>)dataSource
 // NOTE: There is a minor risk here, if a node is transferred from one range controller
 // to another before the first rc updates and clears the node out of this set. It's a pretty
 // wild scenario that I doubt happens in practice.
-- (void)_setVisibleNodes:(ASWeakSet *)newVisibleNodes
+- (void)_setVisibleNodes:(NSHashTable *)newVisibleNodes
 {
   for (ASCellNode *node in _visibleNodes) {
     if (![newVisibleNodes containsObject:node] && node.isVisible) {
@@ -217,7 +217,7 @@ - (void)_updateVisibleNodeIndexPaths
   // 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];
   NSSet<ASCollectionElement *> *visibleElements = [NSSet setWithArray:[_dataSource visibleElementsForRangeController:self]];
-  ASWeakSet *newVisibleNodes = [[ASWeakSet alloc] init];
+  NSHashTable *newVisibleNodes = [NSHashTable hashTableWithOptions:NSHashTableObjectPointerPersonality];
 
   if (visibleElements.count == 0) { // if we don't have any visibleNodes currently (scrolled before or after content)...
     [self _setVisibleNodes:newVisibleNodes];

Source/Details/_ASCollectionReusableView.h

@@ -18,13 +18,14 @@
 #import <UIKit/UIKit.h>
 #import <AsyncDisplayKit/ASBaseDefines.h>
 
-@class ASCellNode;
+@class ASCellNode, ASCollectionElement;
 
 NS_ASSUME_NONNULL_BEGIN
 
 AS_SUBCLASSING_RESTRICTED
 @interface _ASCollectionReusableView : UICollectionReusableView
-@property (nonatomic, weak) ASCellNode *node;
+@property (nonatomic, strong, readonly, nullable) ASCellNode *node;
+@property (nonatomic, strong, nullable) ASCollectionElement *element;
 @property (nonatomic, strong, nullable) UICollectionViewLayoutAttributes *layoutAttributes;
 @end