[ASCollectionElement] Add checks for nil element, prior to other PRs …

…landing. (#421)

Source/ASCollectionView.mm

@@ -1062,7 +1062,13 @@ - (void)collectionView:(UICollectionView *)collectionView willDisplayCell:(UICol
   auto cell = (_ASCollectionViewCell *)rawCell;
 
   ASCollectionElement *element = cell.element;
-  [_visibleElements addObject:element];
+  if (element) {
+    [_visibleElements addObject:element];
+  } else {
+    ASDisplayNodeAssert(NO, @"Unexpected nil element for willDisplayCell: %@, %@, %@", rawCell, self, indexPath);
+    return;
+  }
+
   ASCellNode *cellNode = element.node;
   cellNode.scrollView = collectionView;
 
@@ -1114,8 +1120,13 @@ - (void)collectionView:(UICollectionView *)collectionView didEndDisplayingCell:(
   auto cell = (_ASCollectionViewCell *)rawCell;
 
   ASCollectionElement *element = cell.element;
-  [_visibleElements removeObject:element];
-  ASDisplayNodeAssertNotNil(element, @"Expected element associated with removed cell not to be nil.");
+  if (element) {
+    [_visibleElements removeObject:element];
+  } else {
+    ASDisplayNodeAssert(NO, @"Unexpected nil element for didEndDisplayingCell: %@, %@, %@", rawCell, self, indexPath);
+    return;
+  }
+
   ASCellNode *cellNode = element.node;
 
   if (_asyncDelegateFlags.collectionNodeDidEndDisplayingItem) {
@@ -1143,8 +1154,14 @@ - (void)collectionView:(UICollectionView *)collectionView willDisplaySupplementa
     return;
   }
   auto view = (_ASCollectionReusableView *)rawView;
-
-  [_visibleElements addObject:view.element];
+
+  if (view.element) {
+    [_visibleElements addObject:view.element];
+  } else {
+    ASDisplayNodeAssert(NO, @"Unexpected nil element for willDisplaySupplementaryView: %@, %@, %@", rawView, self, indexPath);
+    return;
+  }
+
   // This is a safeguard similar to the behavior for cells in -[ASCollectionView collectionView:willDisplayCell:forItemAtIndexPath:]
   // It ensures _ASCollectionReusableView receives layoutAttributes and calls applyLayoutAttributes.
   if (view.layoutAttributes == nil) {
@@ -1166,7 +1183,13 @@ - (void)collectionView:(UICollectionView *)collectionView didEndDisplayingSupple
   }
   auto view = (_ASCollectionReusableView *)rawView;
 
-  [_visibleElements removeObject:view.element];
+  if (view.element) {
+    [_visibleElements removeObject:view.element];
+  } else {
+    ASDisplayNodeAssert(NO, @"Unexpected nil element for didEndDisplayingSupplementaryView: %@, %@, %@", rawView, self, indexPath);
+    return;
+  }
+
   if (_asyncDelegateFlags.collectionNodeDidEndDisplayingSupplementaryElement) {
     GET_COLLECTIONNODE_OR_RETURN(collectionNode, (void)0);
     ASCellNode *node = [self supplementaryNodeForElementKind:elementKind atIndexPath:indexPath];