Relayout table view cell nodes if there is a mismatch between content view size and the node's constrained size

- Above is the generic case. Correctly handling it means relayout when the table view enters or leaves editing mode is solved as well.
- Async data source API removal: In a table view, cell nodes should always fill its content view and table view widths. Thus async data source can no longer provide custom constrained size for cell nodes. This removal allows table view to better handle relayout.
- Some more tests are added to ASTableViewTests to check against use cases handled in this diff.

Author: nguyenhuy

AsyncDisplayKit/ASTableView.h

@@ -266,17 +266,6 @@
 
 @optional
 
-/**
- * Provides the constrained size range for measuring the node at the index path.
- *
- * @param tableView The sender.
- *
- * @param indexPath The index path of the node.
- *
- * @returns A constrained size range for layout the node at this index path.
- */
-- (ASSizeRange)tableView:(ASTableView *)tableView constrainedSizeForNodeAtIndexPath:(NSIndexPath *)indexPath;
-
 /**
  * Indicator to lock the data source for data fetching in async mode.
  * We should not update the data source until the data source has been unlocked. Otherwise, it will incur data inconsistence or exception

AsyncDisplayKit/ASTableView.mm

@@ -16,6 +16,7 @@
 #import "ASDisplayNodeInternal.h"
 #import "ASBatchFetching.h"
 #import "ASInternalHelpers.h"
+#import "ASLayout.h"
 
 //#define LOG(...) NSLog(__VA_ARGS__)
 #define LOG(...)
@@ -104,22 +105,30 @@ - (id)forwardingTargetForSelector:(SEL)aSelector
 #pragma mark -
 #pragma mark ASCellNode<->UITableViewCell bridging.
 
+@class _ASTableViewCell;
+
 @protocol _ASTableViewCellDelegate <NSObject>
-- (void)tableViewCell:(UITableViewCell *)cell atIndexPath:(NSIndexPath *)indexPath didTransitionToState:(UITableViewCellStateMask)state;
+- (void)willLayoutSubviewsOfTableViewCell:(_ASTableViewCell *)tableViewCell;
 @end
 
 @interface _ASTableViewCell : UITableViewCell
 @property (nonatomic, weak) id<_ASTableViewCellDelegate> delegate;
-@property (nonatomic) NSIndexPath *indexPath;
+@property (nonatomic, weak) ASCellNode *node;
 @end
 
 @implementation _ASTableViewCell
 // TODO add assertions to prevent use of view-backed UITableViewCell properties (eg .textLabel)
 
+- (void)layoutSubviews
+{
+  [_delegate willLayoutSubviewsOfTableViewCell:self];
+  [super layoutSubviews];
+}
+
 - (void)didTransitionToState:(UITableViewCellStateMask)state
 {
+  [self setNeedsLayout];
   [super didTransitionToState:state];
-  [_delegate tableViewCell:self atIndexPath:_indexPath didTransitionToState:state];
 }
 
 @end
@@ -146,8 +155,6 @@ @interface ASTableView () <ASRangeControllerDelegate, ASDataControllerSource, _A
   NSIndexPath *_contentOffsetAdjustmentTopVisibleRow;
   CGFloat _contentOffsetAdjustment;
 
-  BOOL _asyncDataSourceImplementsConstrainedSizeForNode;
-
   CGFloat _maxWidthForNodesConstrainedSize;
   BOOL _ignoreMaxWidthChange;
 }
@@ -271,12 +278,10 @@ - (void)setAsyncDataSource:(id<ASTableViewDataSource>)asyncDataSource
     super.dataSource = nil;
     _asyncDataSource = nil;
     _proxyDataSource = nil;
-    _asyncDataSourceImplementsConstrainedSizeForNode = NO;
   } else {
     _asyncDataSource = asyncDataSource;
     _proxyDataSource = [[_ASTableViewProxy alloc] initWithTarget:_asyncDataSource interceptor:self];
     super.dataSource = (id<UITableViewDataSource>)_proxyDataSource;
-    _asyncDataSourceImplementsConstrainedSizeForNode = ([_asyncDataSource respondsToSelector:@selector(tableView:constrainedSizeForNodeAtIndexPath:)] ? 1 : 0);
   }
 }
 
@@ -507,7 +512,7 @@ - (UITableViewCell *)tableView:(UITableView *)tableView cellForRowAtIndexPath:(N
   ASCellNode *node = [_dataController nodeAtIndexPath:indexPath];
   [_rangeController configureContentView:cell.contentView forCellNode:node];
 
-  cell.indexPath = indexPath;
+  cell.node = node;
   cell.backgroundColor = node.backgroundColor;
   cell.selectionStyle = node.selectionStyle;
 
@@ -798,11 +803,6 @@ - (ASCellNode *)dataController:(ASDataController *)dataController nodeAtIndexPat
 
 - (ASSizeRange)dataController:(ASDataController *)dataController constrainedSizeForNodeAtIndexPath:(NSIndexPath *)indexPath
 {
-  if (_asyncDataSourceImplementsConstrainedSizeForNode) {
-    return [_asyncDataSource tableView:self constrainedSizeForNodeAtIndexPath:indexPath];
-  }
-  
-  // Default size range
   return ASSizeRangeMake(CGSizeMake(_maxWidthForNodesConstrainedSize, 0),
                          CGSizeMake(_maxWidthForNodesConstrainedSize, FLT_MAX));
 }
@@ -845,22 +845,31 @@ - (NSUInteger)dataControllerNumberOfSections:(ASDataController *)dataController
 
 #pragma mark - _ASTableViewCellDelegate
 
-- (void)tableViewCell:(UITableViewCell *)cell atIndexPath:(NSIndexPath *)indexPath didTransitionToState:(UITableViewCellStateMask)state
+- (void)willLayoutSubviewsOfTableViewCell:(_ASTableViewCell *)tableViewCell
 {
-  [self beginUpdates];
-  ASCellNode *node = [_dataController nodeAtIndexPath:indexPath];
+  CGFloat contentViewWidth = tableViewCell.contentView.bounds.size.width;
+  ASCellNode *node = tableViewCell.node;
+  ASSizeRange constrainedSize = node.constrainedSizeForCalculatedLayout;
 
-  ASSizeRange constrainedSize = [self dataController:_dataController constrainedSizeForNodeAtIndexPath:indexPath];
-  if (state != UITableViewCellStateDefaultMask) {
-    // Edit control or delete confirmation was shown and size of content view was changed.
-    // The new size should be taken into consideration.
-    constrainedSize.min.width = MIN(cell.contentView.frame.size.width, constrainedSize.min.width);
-    constrainedSize.max.width = MIN(cell.contentView.frame.size.width, constrainedSize.max.width);
+  // Table view cells should always fill its content view width.
+  // Normally the content view width equals to the constrained size width (which equals to the table view width).
+  // If there is a mismatch between these values, for example after the table view entered or left editing mode,
+  // content view width is preferred and used to re-measure the cell node.
+  if (contentViewWidth != constrainedSize.max.width) {
+    constrainedSize.min.width = contentViewWidth;
+    constrainedSize.max.width = contentViewWidth;
+
+    // Re-measurement is done on main to ensure thread affinity. In the worst case, this is as fast as UIKit's implementation.
+    //
+    // Unloaded nodes *could* be re-measured off the main thread, but only with the assumption that content view width
+    // 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 calculatedSize = [[node measureWithSizeRange:constrainedSize] size];
+    node.frame = CGRectMake(0, 0, calculatedSize.width, calculatedSize.height);
+    [self endUpdates];
   }
-  
-  [node measureWithSizeRange:constrainedSize];
-  node.frame = CGRectMake(0, 0, node.calculatedSize.width, node.calculatedSize.height);
-  [self endUpdates];
 }
 
 @end

AsyncDisplayKitTests/ASTableViewTests.m

@@ -9,6 +9,7 @@
 #import <XCTest/XCTest.h>
 
 #import "ASTableView.h"
+#import "ASDisplayNode+Subclasses.h"
 
 #define NumberOfSections 10
 #define NumberOfRowsPerSection 20
@@ -56,9 +57,24 @@ - (void)dealloc
 
 @end
 
+@interface ASTestTextCellNode : ASTextCellNode
+/** Calculated by counting how many times -layoutSpecThatFits: is called on the main thread. */
+@property (atomic) int numberOfLayoutsOnMainThread;
+@end
+
+@implementation ASTestTextCellNode
+
+- (ASLayoutSpec *)layoutSpecThatFits:(ASSizeRange)constrainedSize
+{
+  if ([NSThread isMainThread]) {
+    _numberOfLayoutsOnMainThread++;
+  }
+  return [super layoutSpecThatFits:constrainedSize];
+}
+
+@end
+
 @interface ASTableViewFilledDataSource : NSObject <ASTableViewDataSource, ASTableViewDelegate>
-/** Calculated by counting how many times a constrained size is asked for the first node on main thread. */
-@property (atomic) int numberOfRelayouts;
 @end
 
 @implementation ASTableViewFilledDataSource
@@ -75,22 +91,12 @@ - (NSInteger)tableView:(UITableView *)tableView numberOfRowsInSection:(NSInteger
 
 - (ASCellNode *)tableView:(ASTableView *)tableView nodeForRowAtIndexPath:(NSIndexPath *)indexPath
 {
-  ASTextCellNode *textCellNode = [ASTextCellNode new];
+  ASTestTextCellNode *textCellNode = [ASTestTextCellNode new];
   textCellNode.text = indexPath.description;
 
   return textCellNode;
 }
 
-- (ASSizeRange)tableView:(ASTableView *)tableView constrainedSizeForNodeAtIndexPath:(NSIndexPath *)indexPath
-{
-  if ([NSThread isMainThread] && indexPath.section == 0 && indexPath.row == 0) {
-    _numberOfRelayouts++;
-  }
-  CGFloat maxWidth = tableView.bounds.size.width;
-  return ASSizeRangeMake(CGSizeMake(maxWidth, 0),
-                         CGSizeMake(maxWidth, FLT_MAX));
-}
-
 @end
 
 @interface ASTableViewTests : XCTestCase
@@ -260,7 +266,7 @@ - (void)testRelayoutAllRowsWithZeroSizeInitially
 
 - (void)triggerSizeChangeAndAssertRelayoutAllRowsForTableView:(ASTableView *)tableView newSize:(CGSize)newSize
 {
-  XCTestExpectation *nodesMeasuredUsingNewConstrainedSizeExpectation = [self expectationWithDescription:@"nodesMeasuredUsingNewConstrainedSizeExpectation"];
+  XCTestExpectation *nodesMeasuredUsingNewConstrainedSizeExpectation = [self expectationWithDescription:@"nodesMeasuredUsingNewConstrainedSize"];
 
   [tableView beginUpdates];
 
@@ -270,13 +276,11 @@ - (void)triggerSizeChangeAndAssertRelayoutAllRowsForTableView:(ASTableView *)tab
   [tableView layoutIfNeeded];
 
   [tableView endUpdatesAnimated:NO completion:^(BOOL completed) {
-    int numberOfRelayouts = ((ASTableViewFilledDataSource *)(tableView.asyncDataSource)).numberOfRelayouts;
-    XCTAssertEqual(numberOfRelayouts, 1);
-    
     for (int section = 0; section < NumberOfSections; section++) {
       for (int row = 0; row < NumberOfRowsPerSection; row++) {
         NSIndexPath *indexPath = [NSIndexPath indexPathForRow:row inSection:section];
-        ASCellNode *node = [tableView nodeForRowAtIndexPath:indexPath];
+        ASTestTextCellNode *node = (ASTestTextCellNode *)[tableView nodeForRowAtIndexPath:indexPath];
+        XCTAssertEqual(node.numberOfLayoutsOnMainThread, 1);
         XCTAssertEqual(node.constrainedSizeForCalculatedLayout.max.width, newSize.width);
       }
     }
@@ -290,4 +294,138 @@ - (void)triggerSizeChangeAndAssertRelayoutAllRowsForTableView:(ASTableView *)tab
   }];
 }
 
+- (void)testRelayoutVisibleRowsWhenEditingModeIsChanged
+{
+  CGSize tableViewSize = CGSizeMake(100, 500);
+  ASTestTableView *tableView = [[ASTestTableView alloc] initWithFrame:CGRectMake(0, 0, tableViewSize.width, tableViewSize.height)
+                                                                style:UITableViewStylePlain
+                                                    asyncDataFetching:YES];
+  ASTableViewFilledDataSource *dataSource = [ASTableViewFilledDataSource new];
+  
+  tableView.asyncDelegate = dataSource;
+  tableView.asyncDataSource = dataSource;
+  
+  XCTestExpectation *reloadDataExpectation = [self expectationWithDescription:@"reloadData"];
+  [tableView reloadDataWithCompletion:^{
+    for (int section = 0; section < NumberOfSections; section++) {
+      for (int row = 0; row < NumberOfRowsPerSection; row++) {
+        NSIndexPath *indexPath = [NSIndexPath indexPathForRow:row inSection:section];
+        ASTestTextCellNode *node = (ASTestTextCellNode *)[tableView nodeForRowAtIndexPath:indexPath];
+        XCTAssertEqual(node.numberOfLayoutsOnMainThread, 0);
+        XCTAssertEqual(node.constrainedSizeForCalculatedLayout.max.width, tableViewSize.width);
+      }
+    }
+    [reloadDataExpectation fulfill];
+  }];
+  [self waitForExpectationsWithTimeout:5 handler:^(NSError *error) {
+    if (error) {
+      XCTFail(@"Expectation failed: %@", error);
+    }
+  }];
+  
+  NSArray *visibleNodes = [tableView visibleNodes];
+  XCTAssertGreaterThan(visibleNodes.count, 0);
+  
+  // Cause table view to enter editing mode.
+  // Visibile nodes should be re-measured on main thread with the new (smaller) content view width.
+  // Other nodes are untouched.
+  XCTestExpectation *relayoutAfterEnablingEditingExpectation = [self expectationWithDescription:@"relayoutAfterEnablingEditing"];
+  [tableView beginUpdates];
+  [tableView setEditing:YES];
+  [tableView endUpdatesAnimated:YES completion:^(BOOL completed) {
+    for (int section = 0; section < NumberOfSections; section++) {
+      for (int row = 0; row < NumberOfRowsPerSection; row++) {
+        NSIndexPath *indexPath = [NSIndexPath indexPathForRow:row inSection:section];
+        ASTestTextCellNode *node = (ASTestTextCellNode *)[tableView nodeForRowAtIndexPath:indexPath];
+        if ([visibleNodes containsObject:node]) {
+          XCTAssertEqual(node.numberOfLayoutsOnMainThread, 1);
+          XCTAssertLessThan(node.constrainedSizeForCalculatedLayout.max.width, tableViewSize.width);
+        } else {
+          XCTAssertEqual(node.numberOfLayoutsOnMainThread, 0);
+          XCTAssertEqual(node.constrainedSizeForCalculatedLayout.max.width, tableViewSize.width);
+        }
+      }
+    }
+    [relayoutAfterEnablingEditingExpectation fulfill];
+  }];
+  [self waitForExpectationsWithTimeout:5 handler:^(NSError *error) {
+    if (error) {
+      XCTFail(@"Expectation failed: %@", error);
+    }
+  }];
+
+  // Cause table view to leave editing mode.
+  // Visibile nodes should be re-measured again.
+  // All nodes should have max constrained width equals to the table view width.
+  XCTestExpectation *relayoutAfterDisablingEditingExpectation = [self expectationWithDescription:@"relayoutAfterDisablingEditing"];
+  [tableView beginUpdates];
+  [tableView setEditing:NO];
+  [tableView endUpdatesAnimated:YES completion:^(BOOL completed) {
+    for (int section = 0; section < NumberOfSections; section++) {
+      for (int row = 0; row < NumberOfRowsPerSection; row++) {
+        NSIndexPath *indexPath = [NSIndexPath indexPathForRow:row inSection:section];
+        ASTestTextCellNode *node = (ASTestTextCellNode *)[tableView nodeForRowAtIndexPath:indexPath];
+        BOOL visible = [visibleNodes containsObject:node];
+        XCTAssertEqual(node.numberOfLayoutsOnMainThread, visible ? 2: 0);
+        XCTAssertEqual(node.constrainedSizeForCalculatedLayout.max.width, tableViewSize.width);
+      }
+    }
+    [relayoutAfterDisablingEditingExpectation fulfill];
+  }];
+  [self waitForExpectationsWithTimeout:5 handler:^(NSError *error) {
+    if (error) {
+      XCTFail(@"Expectation failed: %@", error);
+    }
+  }];
+}
+
+- (void)testRelayoutRowsAfterEditingModeIsChangedAndTheyBecomeVisible
+{
+  CGSize tableViewSize = CGSizeMake(100, 500);
+  ASTestTableView *tableView = [[ASTestTableView alloc] initWithFrame:CGRectMake(0, 0, tableViewSize.width, tableViewSize.height)
+                                                                style:UITableViewStylePlain
+                                                    asyncDataFetching:YES];
+  ASTableViewFilledDataSource *dataSource = [ASTableViewFilledDataSource new];
+  
+  tableView.asyncDelegate = dataSource;
+  tableView.asyncDataSource = dataSource;
+  
+  XCTestExpectation *reloadDataExpectation = [self expectationWithDescription:@"reloadData"];
+  [tableView reloadDataWithCompletion:^{
+    for (int section = 0; section < NumberOfSections; section++) {
+      for (int row = 0; row < NumberOfRowsPerSection; row++) {
+        NSIndexPath *indexPath = [NSIndexPath indexPathForRow:row inSection:section];
+        ASTestTextCellNode *node = (ASTestTextCellNode *)[tableView nodeForRowAtIndexPath:indexPath];
+        XCTAssertEqual(node.numberOfLayoutsOnMainThread, 0);
+        XCTAssertEqual(node.constrainedSizeForCalculatedLayout.max.width, tableViewSize.width);
+      }
+    }
+    [reloadDataExpectation fulfill];
+  }];
+  [self waitForExpectationsWithTimeout:5 handler:^(NSError *error) {
+    if (error) {
+      XCTFail(@"Expectation failed: %@", error);
+    }
+  }];
+  
+  // Cause table view to enter editing mode and then scroll to the bottom.
+  // The last node should be re-measured on main thread with the new (smaller) content view width.
+  NSIndexPath *lastRowIndexPath = [NSIndexPath indexPathForRow:(NumberOfRowsPerSection - 1) inSection:(NumberOfSections - 1)];
+  XCTestExpectation *relayoutExpectation = [self expectationWithDescription:@"relayout"];
+  [tableView beginUpdates];
+  [tableView setEditing:YES];
+  [tableView setContentOffset:CGPointMake(0, CGFLOAT_MAX) animated:YES];
+  [tableView endUpdatesAnimated:YES completion:^(BOOL completed) {
+    ASTestTextCellNode *node = (ASTestTextCellNode *)[tableView nodeForRowAtIndexPath:lastRowIndexPath];
+    XCTAssertEqual(node.numberOfLayoutsOnMainThread, 1);
+    XCTAssertLessThan(node.constrainedSizeForCalculatedLayout.max.width, tableViewSize.width);
+    [relayoutExpectation fulfill];
+  }];
+  [self waitForExpectationsWithTimeout:5 handler:^(NSError *error) {
+    if (error) {
+      XCTFail(@"Expectation failed: %@", error);
+    }
+  }];
+}
+
 @end