From dcc54c0d366f04dbccb4eb57aff3ff298bc20ab8 Mon Sep 17 00:00:00 2001
From: Michael Schneider <mischneider1@gmail.com>
Date: Mon, 14 Nov 2016 10:10:47 -0800
Subject: [PATCH] Proper relayout if size changes

---
 AsyncDisplayKit/ASCellNode+Internal.h |  6 ++++++
 AsyncDisplayKit/ASCellNode.mm         | 11 +++++++++--
 AsyncDisplayKit/ASCollectionView.mm   |  5 +++++
 AsyncDisplayKit/ASTableView.mm        |  5 +++++
 4 files changed, 25 insertions(+), 2 deletions(-)

diff --git a/AsyncDisplayKit/ASCellNode+Internal.h b/AsyncDisplayKit/ASCellNode+Internal.h
index 99066aa123..d2ebd4e27c 100644
--- a/AsyncDisplayKit/ASCellNode+Internal.h
+++ b/AsyncDisplayKit/ASCellNode+Internal.h
@@ -29,6 +29,12 @@ NS_ASSUME_NONNULL_BEGIN
  */
 - (void)nodeDidRelayout:(ASCellNode *)node sizeChanged:(BOOL)sizeChanged;
 
+/// Returns the constrained size for the given node
+- (ASSizeRange)constrainedSizeForNode:(ASCellNode *)node;
+
+// TODO: coalesc: Maybe we should push the relayout of nodes to the table / collection node?
+//- (void)relayoutNodeIfNeecessary:(ASCellNode *)node;
+
 /*
  * Methods to be called whenever the selection or highlight state changes
  * on ASCellNode. UIKit internally stores these values to update reusable cells.
diff --git a/AsyncDisplayKit/ASCellNode.mm b/AsyncDisplayKit/ASCellNode.mm
index fee0e53935..27e5cb1481 100644
--- a/AsyncDisplayKit/ASCellNode.mm
+++ b/AsyncDisplayKit/ASCellNode.mm
@@ -120,9 +120,8 @@ static NSMutableSet *__cellClassesForVisibilityNotifications = nil; // See +init
 
 - (void)didInvalidateSize
 {
-  // TODO: coalesc: Ask the UITableView for the proper constrained size it can layout
   CGSize oldSize = self.calculatedSize;
-  CGSize newSize = [self sizeThatFits:CGSizeMake(CGRectGetWidth(self.bounds), CGFLOAT_MAX)];
+  CGSize newSize = [self layoutThatFits:[self constrainedSize]].size;
     
   if (CGSizeEqualToSize(oldSize, newSize) == NO) {
     self.frame = {self.frame.origin, newSize};
@@ -174,6 +173,14 @@ static NSMutableSet *__cellClassesForVisibilityNotifications = nil; // See +init
   }
 }
 
+- (ASSizeRange)constrainedSize
+{
+  if (_interactionDelegate != nil) {
+    return [_interactionDelegate constrainedSizeForNode:self];
+  }
+  return ASSizeRangeMake(CGSizeZero, CGSizeMake(CGRectGetWidth(self.bounds), CGFLOAT_MAX));
+}
+
 - (void)setSelected:(BOOL)selected
 {
   if (_selected != selected) {
diff --git a/AsyncDisplayKit/ASCollectionView.mm b/AsyncDisplayKit/ASCollectionView.mm
index 8f4f3aa37e..81132202d7 100644
--- a/AsyncDisplayKit/ASCollectionView.mm
+++ b/AsyncDisplayKit/ASCollectionView.mm
@@ -1579,6 +1579,11 @@ static NSString * const kCellReuseIdentifier = @"_ASCollectionViewCell";
   }
 }
 
+- (ASSizeRange)constrainedSizeForNode:(ASCellNode *)node
+{
+  return [self.layoutInspector collectionView:self constrainedSizeForNodeAtIndexPath:[self indexPathForNode:node]];
+}
+
 - (void)nodeDidRelayout:(ASCellNode *)node sizeChanged:(BOOL)sizeChanged
 {
   ASDisplayNodeAssertMainThread();
diff --git a/AsyncDisplayKit/ASTableView.mm b/AsyncDisplayKit/ASTableView.mm
index 7495a3952e..32ca49ce81 100644
--- a/AsyncDisplayKit/ASTableView.mm
+++ b/AsyncDisplayKit/ASTableView.mm
@@ -1568,6 +1568,11 @@ static NSString * const kCellReuseIdentifier = @"_ASTableViewCell";
   }
 }
 
+- (ASSizeRange)constrainedSizeForNode:(ASCellNode *)node
+{
+  return [self dataController:self.dataController constrainedSizeForNodeAtIndexPath:[self indexPathForNode:node]];
+}
+
 - (void)nodeDidRelayout:(ASCellNode *)node sizeChanged:(BOOL)sizeChanged
 {
   ASDisplayNodeAssertMainThread();