From a659ce7bf8945a140c68913dcd8ec2b07bb4c917 Mon Sep 17 00:00:00 2001
From: Levi McCallum <levi.mccallum@nfl.com>
Date: Tue, 29 Sep 2015 15:16:30 -0700
Subject: [PATCH] Cleanup setup of default flow layout inspector

---
 AsyncDisplayKit/ASCollectionView.mm | 13 +++++++++----
 1 file changed, 9 insertions(+), 4 deletions(-)

diff --git a/AsyncDisplayKit/ASCollectionView.mm b/AsyncDisplayKit/ASCollectionView.mm
index 18b39f9ee4..1e4154852f 100644
--- a/AsyncDisplayKit/ASCollectionView.mm
+++ b/AsyncDisplayKit/ASCollectionView.mm
@@ -205,10 +205,6 @@ static BOOL _isInterceptedSelector(SEL sel)
   _dataController.delegate = _rangeController;
   _dataController.dataSource = self;
   
-  _flowLayoutInspector = [[ASCollectionViewFlowLayoutInspector alloc] init];
-  // TODO: Implement a better path of falling-back to a flow layout
-  _flowLayoutInspector.layout = (UICollectionViewFlowLayout *)layout;
-  
   _batchContext = [[ASBatchContext alloc] init];
 
   _leadingScreensForBatching = 1.0;
@@ -228,6 +224,15 @@ static BOOL _isInterceptedSelector(SEL sel)
   // and should not trigger a relayout.
   _ignoreMaxSizeChange = CGSizeEqualToSize(_maxSizeForNodesConstrainedSize, CGSizeZero);
   
+  // Register the default layout inspector delegate for flow layouts, custom layouts
+  // will need to roll their own ASCollectionViewLayoutInspecting implementation.
+  if ([layout asdk_isFlowLayout]) {
+    _flowLayoutInspector = [[ASCollectionViewFlowLayoutInspector alloc] init];
+    _flowLayoutInspector.layout = (UICollectionViewFlowLayout *)layout;
+    _flowLayoutInspector.collectionView = self;
+    _layoutDelegate = _flowLayoutInspector;
+  }
+
   self.backgroundColor = [UIColor whiteColor];
   
   [self registerClass:[_ASCollectionViewCell class] forCellWithReuseIdentifier:@"_ASCollectionViewCell"];