From e2bbde2a4e8ba2fb299c377cb7826cbf5f56d1bc Mon Sep 17 00:00:00 2001
From: Levi McCallum <levi.mccallum@nfl.com>
Date: Mon, 12 Oct 2015 18:19:40 -0700
Subject: [PATCH] Clean up usage of internal flow layout inspector

---
 AsyncDisplayKit/ASCollectionView.mm           | 11 +++-
 .../ASCollectionViewFlowLayoutInspector.h     |  3 +-
 .../ASCollectionViewFlowLayoutInspector.m     | 56 ++++++++++---------
 3 files changed, 42 insertions(+), 28 deletions(-)

diff --git a/AsyncDisplayKit/ASCollectionView.mm b/AsyncDisplayKit/ASCollectionView.mm
index c16715b4eb..36d700a0c6 100644
--- a/AsyncDisplayKit/ASCollectionView.mm
+++ b/AsyncDisplayKit/ASCollectionView.mm
@@ -337,7 +337,16 @@ static BOOL _isInterceptedSelector(SEL sel)
     _asyncDelegateImplementsInsetSection = ([_asyncDelegate respondsToSelector:@selector(collectionView:layout:insetForSectionAtIndex:)] ? 1 : 0);
   }
 
-  _flowLayoutInspector.collectionView = self;
+  [_flowLayoutInspector cacheSelectorsForCollectionView:self];
+}
+
+- (void)setCollectionViewLayout:(UICollectionViewLayout *)collectionViewLayout
+{
+  [super setCollectionViewLayout:collectionViewLayout];
+  if ([collectionViewLayout asdk_isFlowLayout]) {
+    _flowLayoutInspector = nil;
+    _layoutDelegate = [self flowLayoutInspector];
+  }
 }
 
 - (void)setTuningParameters:(ASRangeTuningParameters)tuningParameters forRangeType:(ASLayoutRangeType)rangeType
diff --git a/AsyncDisplayKit/Details/ASCollectionViewFlowLayoutInspector.h b/AsyncDisplayKit/Details/ASCollectionViewFlowLayoutInspector.h
index 94d65e2117..4e4f266cfd 100644
--- a/AsyncDisplayKit/Details/ASCollectionViewFlowLayoutInspector.h
+++ b/AsyncDisplayKit/Details/ASCollectionViewFlowLayoutInspector.h
@@ -34,10 +34,11 @@
 @interface ASCollectionViewFlowLayoutInspector : NSObject <ASCollectionViewLayoutInspecting>
 
 @property (nonatomic, weak) UICollectionViewFlowLayout *layout;
-@property (nonatomic, weak) ASCollectionView *collectionView;
 
 - (instancetype)initWithCollectionView:(ASCollectionView *)collectionView flowLayout:(UICollectionViewFlowLayout *)flowLayout;
 
+- (void)cacheSelectorsForCollectionView:(ASCollectionView *)collectionView;
+
 - (ASSizeRange)collectionView:(ASCollectionView *)collectionView constrainedSizeForSupplementaryNodeOfKind:(NSString *)kind atIndexPath:(NSIndexPath *)indexPath;
 
 - (NSUInteger)collectionView:(ASCollectionView *)collectionView numberOfSectionsForSupplementaryKind:(NSString *)kind;
diff --git a/AsyncDisplayKit/Details/ASCollectionViewFlowLayoutInspector.m b/AsyncDisplayKit/Details/ASCollectionViewFlowLayoutInspector.m
index efdb26bdaf..f9ea2f8293 100644
--- a/AsyncDisplayKit/Details/ASCollectionViewFlowLayoutInspector.m
+++ b/AsyncDisplayKit/Details/ASCollectionViewFlowLayoutInspector.m
@@ -19,26 +19,30 @@
 
 #pragma mark - Accessors
 
-- (instancetype)initWithCollectionView:(ASCollectionView *)collectionView flowLayout:(UICollectionViewFlowLayout *)flowLayout
+- (instancetype)initWithCollectionView:(ASCollectionView *)collectionView flowLayout:(UICollectionViewFlowLayout *)flowLayout;
 {
-    self = [super init];
+  self = [super init];
 
-    if (flowLayout == nil) {
-        return nil;
-    }
-    
-    if (self != nil) {
-        self.collectionView = collectionView;
-        _layout = flowLayout;
-    }
-    return self;
+  if (flowLayout == nil) {
+    return nil;
+  }
+
+  if (self != nil) {
+    [self cacheSelectorForCollectionView:collectionView];
+    _layout = flowLayout;
+  }
+  return self;
 }
 
-- (void)setCollectionView:(ASCollectionView *)collectionView
+- (void)cacheSelectorsForCollectionView:(ASCollectionView *)collectionView
 {
-    _collectionView = collectionView;
-    _delegateImplementsReferenceSizeForHeader = [[self layoutDelegate] respondsToSelector:@selector(collectionView:layout:referenceSizeForHeaderInSection:)];
-    _delegateImplementsReferenceSizeForFooter = [[self layoutDelegate] respondsToSelector:@selector(collectionView:layout:referenceSizeForFooterInSection:)];
+  if (collectionView == nil) {
+    _delegateImplementsReferenceSizeForHeader = nil;
+    _delegateImplementsReferenceSizeForFooter = nil;
+  } else {
+    _delegateImplementsReferenceSizeForHeader = [[self delegateForCollectionView:collectionView] respondsToSelector:@selector(collectionView:layout:referenceSizeForHeaderInSection:)];
+    _delegateImplementsReferenceSizeForFooter = [[self delegateForCollectionView:collectionView] respondsToSelector:@selector(collectionView:layout:referenceSizeForFooterInSection:)];
+  }
 }
 
 #pragma mark - ASCollectionViewLayoutInspecting
@@ -46,11 +50,11 @@
 - (ASSizeRange)collectionView:(ASCollectionView *)collectionView constrainedSizeForSupplementaryNodeOfKind:(NSString *)kind atIndexPath:(NSIndexPath *)indexPath
 {
   CGSize constrainedSize;
-  CGSize supplementarySize = [self sizeForSupplementaryViewOfKind:kind inSection:indexPath.section];
+  CGSize supplementarySize = [self sizeForSupplementaryViewOfKind:kind inSection:indexPath.section collectionView:collectionView];
   if (_layout.scrollDirection == UICollectionViewScrollDirectionVertical) {
-    constrainedSize = CGSizeMake(_collectionView.bounds.size.width, supplementarySize.height);
+    constrainedSize = CGSizeMake(collectionView.bounds.size.width, supplementarySize.height);
   } else {
-    constrainedSize = CGSizeMake(supplementarySize.height, _collectionView.bounds.size.height);
+    constrainedSize = CGSizeMake(supplementarySize.height, collectionView.bounds.size.height);
   }
   return ASSizeRangeMake(CGSizeZero, constrainedSize);
 }
@@ -66,22 +70,22 @@
 
 - (NSUInteger)collectionView:(ASCollectionView *)collectionView supplementaryViewsOfKind:(NSString *)kind inSection:(NSUInteger)section
 {
-  return [self layoutHasSupplementaryViewOfKind:kind inSection:section] ? 1 : 0;
+  return [self layoutHasSupplementaryViewOfKind:kind inSection:section collectionView:collectionView] ? 1 : 0;
 }
 
 #pragma mark - Private helpers
 
-- (CGSize)sizeForSupplementaryViewOfKind:(NSString *)kind inSection:(NSUInteger)section
+- (CGSize)sizeForSupplementaryViewOfKind:(NSString *)kind inSection:(NSUInteger)section collectionView:(ASCollectionView *)collectionView
 {
   if ([kind isEqualToString:UICollectionElementKindSectionHeader]) {
     if (_delegateImplementsReferenceSizeForHeader) {
-      return [[self layoutDelegate] collectionView:_collectionView layout:_layout referenceSizeForHeaderInSection:section];
+      return [[self delegateForCollectionView:collectionView] collectionView:collectionView layout:_layout referenceSizeForHeaderInSection:section];
     } else {
       return [self.layout headerReferenceSize];
     }
   } else if ([kind isEqualToString:UICollectionElementKindSectionFooter]) {
     if (_delegateImplementsReferenceSizeForFooter) {
-      return [[self layoutDelegate] collectionView:_collectionView layout:_layout referenceSizeForFooterInSection:section];
+      return [[self delegateForCollectionView:collectionView] collectionView:collectionView layout:_layout referenceSizeForFooterInSection:section];
     } else {
       return [self.layout footerReferenceSize];
     }
@@ -90,9 +94,9 @@
   }
 }
 
-- (BOOL)layoutHasSupplementaryViewOfKind:(NSString *)kind inSection:(NSUInteger)section
+- (BOOL)layoutHasSupplementaryViewOfKind:(NSString *)kind inSection:(NSUInteger)section collectionView:(ASCollectionView *)collectionView
 {
-  CGSize size = [self sizeForSupplementaryViewOfKind:kind inSection:section];
+  CGSize size = [self sizeForSupplementaryViewOfKind:kind inSection:section collectionView:collectionView];
   if ([self usedLayoutValueForSize:size] > 0) {
     return YES;
   } else {
@@ -109,9 +113,9 @@
   }
 }
 
-- (id<ASCollectionViewDelegateFlowLayout>)layoutDelegate
+- (id<ASCollectionViewDelegateFlowLayout>)delegateForCollectionView:(ASCollectionView *)collectionView
 {
-  return (id<ASCollectionViewDelegateFlowLayout>)_collectionView.asyncDelegate;
+  return (id<ASCollectionViewDelegateFlowLayout>)collectionView.asyncDelegate;
 }
 
 @end