Address comments from @levi

AsyncDisplayKit/ASCollectionView.mm

@@ -95,6 +95,7 @@ static NSString * const kCellReuseIdentifier = @"_ASCollectionViewCell";
   ASRangeController *_rangeController;
   ASCollectionViewLayoutController *_layoutController;
   id<ASCollectionViewLayoutInspecting> _defaultLayoutInspector;
+  id<ASCollectionViewLayoutInspecting> _layoutInspector;
   NSMutableSet *_cellsForVisibilityUpdates;
   id<ASCollectionViewLayoutFacilitatorProtocol> _layoutFacilitator;
   
@@ -156,6 +157,11 @@ static NSString * const kCellReuseIdentifier = @"_ASCollectionViewCell";
     unsigned int asyncDataSourceNumberOfSectionsInCollectionView:1;
     unsigned int asyncDataSourceCollectionViewConstrainedSizeForNodeAtIndexPath:1;
   } _asyncDataSourceFlags;
+  
+  struct {
+    unsigned int layoutInspectorDidChangeCollectionViewDataSource:1;
+    unsigned int layoutInspectorDidChangeCollectionViewDelegate:1;
+  } _layoutInspectorFlags;
 }
 
 // Used only when ASCollectionView is created directly rather than through ASCollectionNode.
@@ -245,21 +251,11 @@ static NSString * const kCellReuseIdentifier = @"_ASCollectionViewCell";
   // and should not trigger a relayout.
   _ignoreMaxSizeChange = CGSizeEqualToSize(_maxSizeForNodesConstrainedSize, CGSizeZero);
   
-  if ([layout asdk_isFlowLayout]) {
-    // Register the default layout inspector delegate for flow layouts only
-    UICollectionViewFlowLayout *layout = (UICollectionViewFlowLayout *)self.collectionViewLayout;
-    ASDisplayNodeAssertNotNil(layout, @"Collection view layout must be a flow layout to use the built-in inspector");
-    _defaultLayoutInspector = [[ASCollectionViewFlowLayoutInspector alloc] initWithCollectionView:self flowLayout:layout];
-  } else {
-    // Custom layouts will need to roll their own ASCollectionViewLayoutInspecting implementation and set a layout
-    // delegate. In the meantime ASDK provides a null layout inspector that does not provide any implementation
-    // and throws an exception for methods that should be implemented in the <ASCollectionViewLayoutInspecting>
-    _defaultLayoutInspector = [[ASCollectionViewDefaultCustomLayoutInspector alloc] initWithCollectionView:self];
-  }
-  _layoutInspector = _defaultLayoutInspector;
-  
   _layoutFacilitator = layoutFacilitator;
   
+  // Trigger creating the layout inspector
+  [self layoutInspector];
+  
   _proxyDelegate = [[ASCollectionViewProxy alloc] initWithTarget:nil interceptor:self];
   super.delegate = (id<UICollectionViewDelegate>)_proxyDelegate;
   
@@ -373,8 +369,8 @@ static NSString * const kCellReuseIdentifier = @"_ASCollectionViewCell";
   }
   
   super.dataSource = (id<UICollectionViewDataSource>)_proxyDataSource;
-    
+  
-  if ([_layoutInspector respondsToSelector:@selector(didChangeCollectionViewDataSource:)]) {
+  if (_layoutInspectorFlags.layoutInspectorDidChangeCollectionViewDataSource) {
     [_layoutInspector didChangeCollectionViewDataSource:asyncDataSource];
   }
 }
@@ -409,11 +405,49 @@ static NSString * const kCellReuseIdentifier = @"_ASCollectionViewCell";
 
   super.delegate = (id<UICollectionViewDelegate>)_proxyDelegate;
   
-  if ([_layoutInspector respondsToSelector:@selector(didChangeCollectionViewDelegate:)]) {
+  if (_layoutInspectorFlags.layoutInspectorDidChangeCollectionViewDelegate) {
     [_layoutInspector didChangeCollectionViewDelegate:asyncDelegate];
   }
 }
 
+- (void)setCollectionViewLayout:(UICollectionViewLayout *)collectionViewLayout
+{
+  [super setCollectionViewLayout:collectionViewLayout];
+  
+  // Trigger recreation of layout inspector with new collection view layout
+  if (_layoutInspector != nil) {
+    _layoutInspector = nil;
+    [self layoutInspector];
+  }
+}
+
+- (id<ASCollectionViewLayoutInspecting>)layoutInspector
+{
+  if (_layoutInspector == nil) {
+    UICollectionViewFlowLayout *layout = (UICollectionViewFlowLayout *)self.collectionViewLayout;
+    if ([layout asdk_isFlowLayout]) {
+      // Register the default layout inspector delegate for flow layouts only
+      _defaultLayoutInspector = [[ASCollectionViewFlowLayoutInspector alloc] initWithCollectionView:self flowLayout:layout];
+    } else {
+      // Register the default layout inspector delegate for custom collection view layouts
+      _defaultLayoutInspector = [[ASCollectionViewLayoutInspector alloc] initWithCollectionView:self];
+    }
+    
+    // Explicitly call the setter to wire up the _layoutInspectorFlags
+    self.layoutInspector = _defaultLayoutInspector;
+  }
+
+  return _layoutInspector;
+}
+
+- (void)setLayoutInspector:(id<ASCollectionViewLayoutInspecting>)layoutInspector
+{
+  _layoutInspector = layoutInspector;
+  
+  _layoutInspectorFlags.layoutInspectorDidChangeCollectionViewDataSource = [_layoutInspector respondsToSelector:@selector(didChangeCollectionViewDataSource:)];
+  _layoutInspectorFlags.layoutInspectorDidChangeCollectionViewDelegate = [_layoutInspector respondsToSelector:@selector(didChangeCollectionViewDelegate:)];
+}
+
 - (void)setTuningParameters:(ASRangeTuningParameters)tuningParameters forRangeType:(ASLayoutRangeType)rangeType
 {
   [_rangeController setTuningParameters:tuningParameters forRangeMode:ASLayoutRangeModeFull rangeType:rangeType];

AsyncDisplayKit/Details/ASCollectionViewFlowLayoutInspector.h

@@ -64,7 +64,7 @@ NS_ASSUME_NONNULL_BEGIN
  * far as possible based on the scrollable direction of the collection view. It throws exceptions for delegate
  * methods that are related to supplementary node's management.
  */
-@interface ASCollectionViewDefaultCustomLayoutInspector : NSObject <ASCollectionViewLayoutInspecting>
+@interface ASCollectionViewLayoutInspector : NSObject <ASCollectionViewLayoutInspecting>
 
 - (instancetype)init NS_UNAVAILABLE;
 - (instancetype)initWithCollectionView:(ASCollectionView *)collectionView NS_DESIGNATED_INITIALIZER;

AsyncDisplayKit/Details/ASCollectionViewFlowLayoutInspector.m

@@ -19,7 +19,7 @@
 
 // Returns a constrained size to let the cells layout itself as far as possible based on the scrollable direction
 // of the collection view
-static ASSizeRange ASDefaultConstrainedSizeForNodeForCollectionView(ASCollectionView *collectionView) {
+static inline ASSizeRange NodeConstrainedSizeWithCollectionView(ASCollectionView *collectionView) {
   CGSize maxSize = collectionView.bounds.size;
   if (ASScrollDirectionContainsHorizontalDirection(collectionView.scrollableDirections)) {
     maxSize.width = FLT_MAX;
@@ -29,9 +29,9 @@ static ASSizeRange ASDefaultConstrainedSizeForNodeForCollectionView(ASCollection
   return ASSizeRangeMake(CGSizeZero, maxSize);
 }
 
-#pragma mark - ASCollectionViewDefaultCustomLayoutInspector
+#pragma mark - ASCollectionViewLayoutInspector
 
-@implementation ASCollectionViewDefaultCustomLayoutInspector {
+@implementation ASCollectionViewLayoutInspector {
   struct {
     unsigned int implementsConstrainedSizeForNodeAtIndexPath:1;
   } _dataSourceFlags;
@@ -65,7 +65,7 @@ static ASSizeRange ASDefaultConstrainedSizeForNodeForCollectionView(ASCollection
     return [collectionView.asyncDataSource collectionView:collectionView constrainedSizeForNodeAtIndexPath:indexPath];
   }
   
-  return ASDefaultConstrainedSizeForNodeForCollectionView(collectionView);
+  return NodeConstrainedSizeWithCollectionView(collectionView);
 }
 
 - (ASSizeRange)collectionView:(ASCollectionView *)collectionView constrainedSizeForSupplementaryNodeOfKind:(NSString *)kind atIndexPath:(NSIndexPath *)indexPath
@@ -152,11 +152,11 @@ static ASSizeRange ASDefaultConstrainedSizeForNodeForCollectionView(ASCollection
   }
   
   CGSize itemSize = _layout.itemSize;
-  if (!CGSizeEqualToSize(itemSize, kDefaultItemSize)) {
+  if (CGSizeEqualToSize(itemSize, kDefaultItemSize) == NO) {
     return ASSizeRangeMake(itemSize, itemSize);
   }
   
-  return ASDefaultConstrainedSizeForNodeForCollectionView(collectionView);
+  return NodeConstrainedSizeWithCollectionView(collectionView);
 }
 
 - (ASSizeRange)collectionView:(ASCollectionView *)collectionView constrainedSizeForSupplementaryNodeOfKind:(NSString *)kind atIndexPath:(NSIndexPath *)indexPath

AsyncDisplayKitTests/ASCollectionViewTests.m

@@ -129,7 +129,7 @@
   UICollectionViewLayout *layout = [[UICollectionViewLayout alloc] init];
   ASCollectionView *collectionView = [[ASCollectionView alloc] initWithFrame:CGRectZero collectionViewLayout:layout];
   XCTAssert(collectionView.layoutInspector != nil, @"should automatically set a layout delegate for flow layouts");
-  XCTAssert([collectionView.layoutInspector isKindOfClass:[ASCollectionViewDefaultCustomLayoutInspector class]], @"should have a flow layout inspector by default");
+  XCTAssert([collectionView.layoutInspector isKindOfClass:[ASCollectionViewLayoutInspector class]], @"should have a default layout inspector by default");
 }
 
 - (void)testThatRegisteringASupplementaryNodeStoresItForIntrospection