Improve Table Performance, Especially When Using Section Index Bar (#2506)

* Only trigger table view height requery if row-node height changed

* Make ASTableView handle the section index bar better

* Remove outdated comment

* Add a main thread assertion

AsyncDisplayKit/ASTableNode.mm

@@ -388,19 +388,24 @@ ASEnvironmentCollectionTableSetEnvironmentState(_environmentStateLock)
 
 - (void)reloadDataInitiallyIfNeeded
 {
+  ASDisplayNodeAssertMainThread();
   if (!self.dataController.initialReloadDataHasBeenCalled) {
-    [self reloadData];
+    // Note: Just calling reloadData isn't enough here – we need to
+    // ensure that _nodesConstrainedWidth is updated first.
+    [self.view layoutIfNeeded];
   }
 }
 
 - (NSInteger)numberOfRowsInSection:(NSInteger)section
 {
+  ASDisplayNodeAssertMainThread();
   [self reloadDataInitiallyIfNeeded];
   return [self.dataController numberOfRowsInSection:section];
 }
 
 - (NSInteger)numberOfSections
 {
+  ASDisplayNodeAssertMainThread();
   [self reloadDataInitiallyIfNeeded];
   return [self.dataController numberOfSections];
 }

AsyncDisplayKit/ASTableView.mm

@@ -130,6 +130,10 @@ static NSString * const kCellReuseIdentifier = @"_ASTableViewCell";
   BOOL _performingBatchUpdates;
   NSMutableSet *_cellsForVisibilityUpdates;
 
+  // The section index overlay view, if there is one present.
+  // This is useful because we need to measure our row nodes against (width - indexView.width).
+  __weak UIView *_sectionIndexView;
+  
   struct {
     unsigned int scrollViewDidScroll:1;
     unsigned int scrollViewWillBeginDragging:1;
@@ -181,6 +185,7 @@ static NSString * const kCellReuseIdentifier = @"_ASTableViewCell";
     unsigned int tableNodeCanMoveRow:1;
     unsigned int tableViewMoveRow:1;
     unsigned int tableNodeMoveRow:1;
+    unsigned int sectionIndexMethods:1; // if both section index methods are implemented
   } _asyncDataSourceFlags;
 }
 
@@ -321,6 +326,7 @@ static NSString * const kCellReuseIdentifier = @"_ASTableViewCell";
     _asyncDataSourceFlags.tableNodeNodeBlockForRow = [_asyncDataSource respondsToSelector:@selector(tableNode:nodeBlockForRowAtIndexPath:)];
     _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:)];
     
     ASDisplayNodeAssert(_asyncDataSourceFlags.tableViewNodeBlockForRow
                         || _asyncDataSourceFlags.tableViewNodeForRow
@@ -559,8 +565,9 @@ static NSString * const kCellReuseIdentifier = @"_ASTableViewCell";
 
 - (void)layoutSubviews
 {
-  if (_nodesConstrainedWidth != self.bounds.size.width) {
+  CGFloat constrainedWidth = self.bounds.size.width - [self sectionIndexWidth];
-    _nodesConstrainedWidth = self.bounds.size.width;
+  if (_nodesConstrainedWidth != constrainedWidth) {
+    _nodesConstrainedWidth = constrainedWidth;
 
     // First width change occurs during initial configuration. An expensive relayout pass is unnecessary at that time
     // and should be avoided, assuming that the initial data loading automatically runs shortly afterward.
@@ -1512,12 +1519,17 @@ static NSString * const kCellReuseIdentifier = @"_ASTableViewCell";
     // is the same for all cells (because there is no easy way to get that individual value before the node being assigned to a _ASTableViewCell).
     // Also, in many cases, some nodes may not need to be re-measured at all, such as when user enters and then immediately leaves editing mode.
     // To avoid premature optimization and making such assumption, as well as to keep ASTableView simple, re-measurement is strictly done on main.
-    [self beginUpdates];
+    CGSize oldSize = node.bounds.size;
     const CGSize calculatedSize = [node layoutThatFits:constrainedSize].size;
-    node.frame = CGRectMake(0, 0, calculatedSize.width, calculatedSize.height);
+    node.frame = { .size = calculatedSize };
+
+    // If the node height changed, trigger a height requery.
+    if (oldSize.height != calculatedSize.height) {
+      [self beginUpdates];
       [self endUpdates];
     }
   }
+}
 
 #pragma mark - ASCellNodeDelegate
 
@@ -1579,6 +1591,31 @@ static NSString * const kCellReuseIdentifier = @"_ASTableViewCell";
 
 #pragma mark - Helper Methods
 
+// Note: This is called every layout, and so it is very performance sensitive.
+- (CGFloat)sectionIndexWidth
+{
+  // If they don't implement the methods, then there's no section index.
+  if (_asyncDataSourceFlags.sectionIndexMethods == NO) {
+    return 0;
+  }
+
+  UIView *indexView = _sectionIndexView;
+  if (indexView.superview == self) {
+    return indexView.frame.size.width;
+  }
+
+  CGRect bounds = self.bounds;
+  for (UIView *view in self.subviews) {
+    CGRect frame = view.frame;
+    // Section index is right-aligned and less than half-width.
+    if (CGRectGetMaxX(frame) == CGRectGetMaxX(bounds) && frame.size.width * 2 < bounds.size.width) {
+      _sectionIndexView = view;
+      return frame.size.width;
+    }
+  }
+  return 0;
+}
+
 /// @note This should be a UIKit index path.
 - (BOOL)isIndexPath:(NSIndexPath *)indexPath immediateSuccessorOfIndexPath:(NSIndexPath *)anchor
 {

AsyncDisplayKit/ASTableViewInternal.h

@@ -59,4 +59,7 @@
  */
 - (NSArray<NSIndexPath *> *)convertIndexPathsToTableNode:(NSArray<NSIndexPath *> *)indexPaths;
 
+/// Returns the width of the section index view on the right-hand side of the table, if one is present.
+- (CGFloat)sectionIndexWidth;
+
 @end

AsyncDisplayKitTests/ASTableViewTests.m

@@ -18,6 +18,7 @@
 #import "ASTableNode.h"
 #import "ASTableView+Undeprecated.h"
 #import <JGMethodSwizzler/JGMethodSwizzler.h>
+#import "ASXCTExtensions.h"
 
 #define NumberOfSections 10
 #define NumberOfRowsPerSection 20
@@ -117,10 +118,20 @@
 @end
 
 @interface ASTableViewFilledDataSource : NSObject <ASTableDataSource, ASTableDelegate>
+@property (nonatomic) BOOL usesSectionIndex;
 @end
 
 @implementation ASTableViewFilledDataSource
 
+- (BOOL)respondsToSelector:(SEL)aSelector
+{
+  if (aSelector == @selector(sectionIndexTitlesForTableView:) || aSelector == @selector(tableView:sectionForSectionIndexTitle:atIndex:)) {
+    return _usesSectionIndex;
+  } else {
+    return [super respondsToSelector:aSelector];
+  }
+}
+
 - (NSInteger)numberOfSectionsInTableView:(UITableView *)tableView
 {
   return NumberOfSections;
@@ -148,6 +159,16 @@
   };
 }
 
+- (nullable NSArray<NSString *> *)sectionIndexTitlesForTableView:(UITableView *)tableView
+{
+  return @[ @"A", @"B", @"C" ];
+}
+
+- (NSInteger)tableView:(UITableView *)tableView sectionForSectionIndexTitle:(NSString *)title atIndex:(NSInteger)index
+{
+  return 0;
+}
+
 @end
 
 @interface ASTableViewFilledDelegate : NSObject <ASTableDelegate>
@@ -614,6 +635,75 @@
   [UITableView deswizzleAllInstanceMethods];
 }
 
+/**
+ * This tests an issue where, if the table is loaded before the first layout pass,
+ * the nodes are first measured with a constrained width of 0 which isn't ideal.
+ */
+- (void)testThatNodeConstrainedSizesAreCorrectIfReloadIsPreempted
+{
+  ASTableNode *node = [[ASTableNode alloc] initWithStyle:UITableViewStylePlain];
+
+  ASTableViewFilledDataSource *dataSource = [ASTableViewFilledDataSource new];
+  CGFloat cellWidth = 320;
+  node.frame = CGRectMake(0, 0, cellWidth, 480);
+
+  node.dataSource = dataSource;
+  node.delegate = dataSource;
+
+  // Trigger data load BEFORE first layout pass, to ensure constrained size is correct.
+  XCTAssertGreaterThan(node.numberOfSections, 0);
+  [node waitUntilAllUpdatesAreCommitted];
+
+  ASSizeRange expectedSizeRange = ASSizeRangeMakeExactSize(CGSizeMake(cellWidth, 0));
+  expectedSizeRange.max.height = CGFLOAT_MAX;
+
+  for (NSInteger i = 0; i < node.numberOfSections; i++) {
+    for (NSInteger j = 0; j < [node numberOfRowsInSection:i]; j++) {
+      NSIndexPath *indexPath = [NSIndexPath indexPathForItem:j inSection:i];
+      ASTestTextCellNode *cellNode = (id)[node nodeForRowAtIndexPath:indexPath];
+      ASXCTAssertEqualSizeRanges(cellNode.constrainedSizeForCalculatedLayout, expectedSizeRange);
+      XCTAssertEqual(cellNode.numberOfLayoutsOnMainThread, 0);
+    }
+  }
+}
+
+- (void)testSectionIndexHandling
+{
+  ASTableNode *node = [[ASTableNode alloc] initWithStyle:UITableViewStylePlain];
+
+  ASTableViewFilledDataSource *dataSource = [ASTableViewFilledDataSource new];
+  dataSource.usesSectionIndex = YES;
+  node.frame = CGRectMake(0, 0, 320, 480);
+
+  node.dataSource = dataSource;
+  node.delegate = dataSource;
+
+  // Trigger data load
+  XCTAssertGreaterThan(node.numberOfSections, 0);
+  XCTAssertGreaterThan([node numberOfRowsInSection:0], 0);
+  [node waitUntilAllUpdatesAreCommitted];
+
+  UITableViewCell *cell = [node.view cellForRowAtIndexPath:[NSIndexPath indexPathForItem:0 inSection:0]];
+  XCTAssertNotNil(cell);
+
+  CGFloat cellWidth = cell.contentView.frame.size.width;
+  XCTAssert(cellWidth > 0 && cellWidth < 320, @"Expected cell width to be about 305. Width: %@", @(cellWidth));
+
+  ASSizeRange expectedSizeRange = ASSizeRangeMakeExactSize(CGSizeMake(cellWidth, 0));
+  expectedSizeRange.max.height = CGFLOAT_MAX;
+  
+  for (NSInteger i = 0; i < node.numberOfSections; i++) {
+    for (NSInteger j = 0; j < [node numberOfRowsInSection:i]; j++) {
+      NSIndexPath *indexPath = [NSIndexPath indexPathForItem:j inSection:i];
+      ASTestTextCellNode *cellNode = (id)[node nodeForRowAtIndexPath:indexPath];
+      ASXCTAssertEqualSizeRanges(cellNode.constrainedSizeForCalculatedLayout, expectedSizeRange);
+      // We will have to accept a relayout on main thread, since the index bar won't show
+      // up until some of the cells are inserted.
+      XCTAssertLessThanOrEqual(cellNode.numberOfLayoutsOnMainThread, 1);
+    }
+  }
+}
+
 @end
 
 @implementation UITableView (Testing)

AsyncDisplayKitTests/ASXCTExtensions.h

@@ -31,3 +31,6 @@
 
 #define ASXCTAssertNotEqualDimensions(r0, r1, ...) \
   _XCTPrimitiveAssertNotEqualObjects(self, NSStringFromASDimension(r0), @#r0, NSStringFromASDimension(r1), @#r1, __VA_ARGS__)
+
+#define ASXCTAssertEqualSizeRanges(r0, r1, ...) \
+  _XCTPrimitiveAssertEqualObjects(self, NSStringFromASSizeRange(r0), @#r0, NSStringFromASSizeRange(r1), @#r1, __VA_ARGS__)