From bba659d04bd3a450376be9934f45c524e0b7ec77 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ce=CC=81sar=20Este=CC=81banez=20Tasco=CC=81n?=
 <cestebanez@gmail.com>
Date: Mon, 11 Dec 2017 21:25:29 +0100
Subject: [PATCH] Check for nil elements on ASTableView as well.

---
 Source/ASCollectionView.mm |  2 +-
 Source/ASTableView.mm      | 29 ++++++++++++++++++++++++++---
 2 files changed, 27 insertions(+), 4 deletions(-)

diff --git a/Source/ASCollectionView.mm b/Source/ASCollectionView.mm
index 2811f6271..1c884c304 100644
--- a/Source/ASCollectionView.mm
+++ b/Source/ASCollectionView.mm
@@ -1174,7 +1174,7 @@ - (void)collectionView:(UICollectionView *)collectionView willDisplayCell:(UICol
   ASCellNode *cellNode = element.node;
   cellNode.scrollView = collectionView;
 
-  // Update the selected background view in collectionView:willDisplayCell:forItemAtIndexPath: otherwise it could be to
+  // Update the selected background view in collectionView:willDisplayCell:forItemAtIndexPath: otherwise it could be too
   // early e.g. if the selectedBackgroundView was set in didLoad()
   cell.selectedBackgroundView = cellNode.selectedBackgroundView;
   
diff --git a/Source/ASTableView.mm b/Source/ASTableView.mm
index 4d9d01f81..1cc369b07 100644
--- a/Source/ASTableView.mm
+++ b/Source/ASTableView.mm
@@ -130,6 +130,15 @@ - (void)setElement:(ASCollectionElement *)element
   [node __setHighlightedFromUIKit:self.highlighted];
 }
 
+- (BOOL)consumesCellNodeVisibilityEvents
+{
+  ASCellNode *node = self.node;
+  if (node == nil) {
+    return NO;
+  }
+  return ASSubclassOverridesSelector([ASCellNode class], [node class], @selector(cellNodeVisibilityEvent:inScrollView:withCellFrame:));
+}
+
 - (void)setSelected:(BOOL)selected animated:(BOOL)animated
 {
   [super setSelected:selected animated:animated];
@@ -971,7 +980,14 @@ - (void)tableView:(UITableView *)tableView moveRowAtIndexPath:(NSIndexPath *)sou
 - (void)tableView:(UITableView *)tableView willDisplayCell:(_ASTableViewCell *)cell forRowAtIndexPath:(NSIndexPath *)indexPath
 {
   ASCollectionElement *element = cell.element;
-  [_visibleElements addObject:element];
+  if (element) {
+    ASDisplayNodeAssertTrue([_dataController.visibleMap elementForItemAtIndexPath:indexPath] == element);
+    [_visibleElements addObject:element];
+  } else {
+    ASDisplayNodeAssert(NO, @"Unexpected nil element for willDisplayCell: %@, %@, %@", cell, self, indexPath);
+    return;
+  }
+
   ASCellNode *cellNode = element.node;
   cellNode.scrollView = tableView;
 
@@ -991,15 +1007,22 @@ - (void)tableView:(UITableView *)tableView willDisplayCell:(_ASTableViewCell *)c
   
   [_rangeController setNeedsUpdate];
   
-  if (ASSubclassOverridesSelector([ASCellNode class], [cellNode class], @selector(cellNodeVisibilityEvent:inScrollView:withCellFrame:))) {
+  if ([cell consumesCellNodeVisibilityEvents]) {
     [_cellsForVisibilityUpdates addObject:cell];
   }
 }
 
 - (void)tableView:(UITableView *)tableView didEndDisplayingCell:(_ASTableViewCell *)cell forRowAtIndexPath:(NSIndexPath *)indexPath
 {
+  // Retrieve the element from cell instead of visible map because at this point visible map could have been updated and no longer holds the element.
   ASCollectionElement *element = cell.element;
-  [_visibleElements removeObject:element];
+  if (element) {
+    [_visibleElements removeObject:element];
+  } else {
+    ASDisplayNodeAssert(NO, @"Unexpected nil element for didEndDisplayingCell: %@, %@, %@", cell, self, indexPath);
+    return;
+  }
+
   ASCellNode *cellNode = element.node;
 
   [_rangeController setNeedsUpdate];