[ASCollectionLayout] Fix element lookup for supplementary attributes (#…

…1707)
TextureGroup · Dec 24, 2019 · b59c6d0 · b59c6d0
1 parent 88a640a
commit b59c6d0
Showing 1 changed file with 12 additions and 2 deletions.
diff --git a/Source/Private/ASCollectionLayout.mm b/Source/Private/ASCollectionLayout.mm
@@ -338,7 +338,12 @@ + (void)_measureElementsInRect:(CGRect)rect blockingRect:(CGRect)blockingRect la
   if (NSUInteger count = blockingAttrs.count) {
     ASDispatchApply(count, queue, 0, ^(size_t i) {
       UICollectionViewLayoutAttributes *attrs = blockingAttrs[i];
-      ASCellNode *node = [elements elementForItemAtIndexPath:attrs.indexPath].node;
+      ASCellNode *node;
+      if (attrs.representedElementKind == nil) {
+        node = [elements elementForItemAtIndexPath:attrs.indexPath].node;
+      } else {
+        node = [elements supplementaryElementOfKind:attrs.representedElementKind atIndexPath:attrs.indexPath].node;
+      }
       CGSize expectedSize = attrs.frame.size;
       if (! CGSizeEqualToSize(expectedSize, node.calculatedSize)) {
         [node layoutThatFits:ASCollectionLayoutElementSizeRangeFromSize(expectedSize)];
@@ -353,7 +358,12 @@ + (void)_measureElementsInRect:(CGRect)rect blockingRect:(CGRect)blockingRect la
       __strong ASElementMap *strongElements = weakElements;
       if (strongElements) {
         UICollectionViewLayoutAttributes *attrs = nonBlockingAttrs[i];
-        ASCellNode *node = [elements elementForItemAtIndexPath:attrs.indexPath].node;
+        ASCellNode *node;
+        if (attrs.representedElementKind == nil) {
+          node = [elements elementForItemAtIndexPath:attrs.indexPath].node;
+        } else {
+          node = [elements supplementaryElementOfKind:attrs.representedElementKind atIndexPath:attrs.indexPath].node;
+        }
         CGSize expectedSize = attrs.frame.size;
         if (! CGSizeEqualToSize(expectedSize, node.calculatedSize)) {
           [node layoutThatFits:ASCollectionLayoutElementSizeRangeFromSize(expectedSize)];