From 479d40464e78c6ac1b39406b2bf6059238b8ba85 Mon Sep 17 00:00:00 2001
From: Max Wang <max.wang.cs@gmail.com>
Date: Fri, 9 Feb 2018 12:16:18 -0800
Subject: [PATCH] [ASTableNode & ASCollectionNode] Keepalive reference for node
 if their view is necessarily alive (has a superview). (#793)

* fix SIMULATE_WEB_RESPONSE not imported #449

* Fix to make rangeMode update in right time

* Keep collection/table node alive if view still in use.
---
 Source/ASCollectionNode.mm     |  7 +++++++
 Source/ASCollectionView.mm     | 21 ++++++++++++++++++++-
 Source/ASTableNode.mm          |  7 +++++++
 Source/ASTableView.mm          | 21 ++++++++++++++++++++-
 Tests/ASCollectionViewTests.mm |  1 +
 5 files changed, 55 insertions(+), 2 deletions(-)

diff --git a/Source/ASCollectionNode.mm b/Source/ASCollectionNode.mm
index 2449377fd..01669ca79 100644
--- a/Source/ASCollectionNode.mm
+++ b/Source/ASCollectionNode.mm
@@ -172,6 +172,13 @@ - (instancetype)initWithFrame:(CGRect)frame collectionViewLayout:(UICollectionVi
   return self;
 }
 
+- (void)dealloc
+{
+  if ([self isNodeLoaded]) {
+    ASDisplayNodeAssert(self.view.superview == nil, @"Node's view should be removed from hierarchy.");
+  }
+}
+
 #pragma mark ASDisplayNode
 
 - (void)didLoad
diff --git a/Source/ASCollectionView.mm b/Source/ASCollectionView.mm
index 05ed5ebf9..10dd6dc9a 100644
--- a/Source/ASCollectionView.mm
+++ b/Source/ASCollectionView.mm
@@ -151,7 +151,12 @@ @interface ASCollectionView () <ASRangeControllerDataSource, ASRangeControllerDe
    * Counter used to keep track of nested batch updates.
    */
   NSInteger _batchUpdateCount;
-  
+
+  /**
+   * Keep a strong reference to node till view is ready to release.
+   */
+  ASCollectionNode *_keepalive_node;
+
   struct {
     unsigned int scrollViewDidScroll:1;
     unsigned int scrollViewWillBeginDragging:1;
@@ -2250,6 +2255,20 @@ - (void)didMoveToWindow
   }
 }
 
+- (void)willMoveToSuperview:(UIView *)newSuperview
+{
+  if (self.superview == nil && newSuperview != nil) {
+    _keepalive_node = self.collectionNode;
+  }
+}
+
+- (void)didMoveToSuperview
+{
+  if (self.superview == nil) {
+    _keepalive_node = nil;
+  }
+}
+
 #pragma mark ASCALayerExtendedDelegate
 
 /**
diff --git a/Source/ASTableNode.mm b/Source/ASTableNode.mm
index d68e4a9d7..09e5eb9bd 100644
--- a/Source/ASTableNode.mm
+++ b/Source/ASTableNode.mm
@@ -104,6 +104,13 @@ - (instancetype)init
   return [self initWithStyle:UITableViewStylePlain];
 }
 
+- (void)dealloc
+{
+  if ([self isNodeLoaded]) {
+    ASDisplayNodeAssert(self.view.superview == nil, @"Node's view should be removed from hierarchy.");
+  }
+}
+
 #pragma mark ASDisplayNode
 
 - (void)didLoad
diff --git a/Source/ASTableView.mm b/Source/ASTableView.mm
index 1b03797ca..1896b37cf 100644
--- a/Source/ASTableView.mm
+++ b/Source/ASTableView.mm
@@ -216,7 +216,12 @@ @interface ASTableView () <ASRangeControllerDataSource, ASRangeControllerDelegat
    * Counter used to keep track of nested batch updates.
    */
   NSInteger _batchUpdateCount;
-  
+
+  /**
+   * Keep a strong reference to node till view is ready to release.
+   */
+  ASTableNode *_keepalive_node;
+
   struct {
     unsigned int scrollViewDidScroll:1;
     unsigned int scrollViewWillBeginDragging:1;
@@ -1913,4 +1918,18 @@ - (void)didMoveToWindow
   }
 }
 
+- (void)willMoveToSuperview:(UIView *)newSuperview
+{
+  if (self.superview == nil && newSuperview != nil) {
+    _keepalive_node = self.tableNode;
+  }
+}
+
+- (void)didMoveToSuperview
+{
+  if (self.superview == nil) {
+    _keepalive_node = nil;
+  }
+}
+
 @end
diff --git a/Tests/ASCollectionViewTests.mm b/Tests/ASCollectionViewTests.mm
index 0391e992d..311d7943f 100644
--- a/Tests/ASCollectionViewTests.mm
+++ b/Tests/ASCollectionViewTests.mm
@@ -1071,6 +1071,7 @@ - (void)testInitialRangeBounds
       for (NSInteger i = 0; i < c; i++) {
         NSIndexPath *ip = [NSIndexPath indexPathForItem:i inSection:s];
         ASCellNode *node = [cn nodeForItemAtIndexPath:ip];
+        [[NSRunLoop mainRunLoop] runUntilDate:[NSDate dateWithTimeIntervalSinceNow:0.001]];
         if (node.inPreloadState) {
           CGRect frame = [cn.view layoutAttributesForItemAtIndexPath:ip].frame;
           r = CGRectUnion(r, frame);