[ASRangeController] Optimize calls into UICollectionViewLayout via union rect technique. (#3102)

Details in https://github.com/facebook/AsyncDisplayKit/issues/3082

Source/Details/ASAbstractLayoutController.h

@@ -31,6 +31,8 @@ ASDISPLAYNODE_EXTERN_C_END
 
 - (NSSet *)indexPathsForScrolling:(ASScrollDirection)scrollDirection rangeMode:(ASLayoutRangeMode)rangeMode rangeType:(ASLayoutRangeType)rangeType __unavailable;
 
+- (void)allIndexPathsForScrolling:(ASScrollDirection)scrollDirection rangeMode:(ASLayoutRangeMode)rangeMode displaySet:(NSSet * _Nullable * _Nullable)displaySet preloadSet:(NSSet * _Nullable * _Nullable)preloadSet __unavailable;
+
 @end
 
 NS_ASSUME_NONNULL_END

Source/Details/ASCollectionViewLayoutController.mm

@@ -53,6 +53,50 @@ typedef struct ASRangeGeometry ASRangeGeometry;
   return [self indexPathsForItemsWithinRangeBounds:rangeBounds];
 }
 
+- (void)allIndexPathsForScrolling:(ASScrollDirection)scrollDirection rangeMode:(ASLayoutRangeMode)rangeMode displaySet:(NSSet **)displaySet preloadSet:(NSSet **)preloadSet
+{
+  if (displaySet == NULL || preloadSet == NULL) {
+    return;
+  }
+  
+  ASRangeTuningParameters displayParams = [self tuningParametersForRangeMode:rangeMode rangeType:ASLayoutRangeTypeDisplay];
+  ASRangeTuningParameters preloadParams = [self tuningParametersForRangeMode:rangeMode rangeType:ASLayoutRangeTypePreload];
+  CGRect displayBounds = [self rangeBoundsWithScrollDirection:scrollDirection rangeTuningParameters:displayParams];
+  CGRect preloadBounds = [self rangeBoundsWithScrollDirection:scrollDirection rangeTuningParameters:preloadParams];
+  
+  CGRect unionBounds = CGRectUnion(displayBounds, preloadBounds);
+  NSArray *layoutAttributes = [_collectionViewLayout layoutAttributesForElementsInRect:unionBounds];
+
+  NSMutableSet *display = [NSMutableSet setWithCapacity:layoutAttributes.count];
+  NSMutableSet *preload = [NSMutableSet setWithCapacity:layoutAttributes.count];
+
+  for (UICollectionViewLayoutAttributes *la in layoutAttributes) {
+    // Manually filter out elements that don't intersect the range bounds.
+    // See comment in indexPathsForItemsWithinRangeBounds:
+    // This is re-implemented here so that the iteration over layoutAttributes can be done once to check both ranges.
+    CGRect frame = la.frame;
+    BOOL intersectsDisplay = CGRectIntersectsRect(displayBounds, frame);
+    BOOL intersectsPreload = CGRectIntersectsRect(preloadBounds, frame);
+    if (intersectsDisplay == NO && intersectsPreload == NO && CATransform3DIsIdentity(la.transform3D) == YES) {
+      // Questionable why the element would be included here, but it doesn't belong.
+      continue;
+    }
+    
+    // Avoid excessive retains and releases, as well as property calls. We know the indexPath is kept alive by la.
+    __unsafe_unretained NSIndexPath *indexPath = la.indexPath;
+    if (intersectsDisplay) {
+      [display addObject:indexPath];
+    }
+    if (intersectsPreload) {
+      [preload addObject:indexPath];
+    }
+  }
+
+  *displaySet = display;
+  *preloadSet = preload;
+  return;
+}
+
 - (NSSet *)indexPathsForItemsWithinRangeBounds:(CGRect)rangeBounds
 {
   NSArray *layoutAttributes = [_collectionViewLayout layoutAttributesForElementsInRect:rangeBounds];

Source/Details/ASLayoutController.h

@@ -36,6 +36,8 @@ ASDISPLAYNODE_EXTERN_C_END
 
 - (NSSet<NSIndexPath *> *)indexPathsForScrolling:(ASScrollDirection)scrollDirection rangeMode:(ASLayoutRangeMode)rangeMode rangeType:(ASLayoutRangeType)rangeType;
 
+- (void)allIndexPathsForScrolling:(ASScrollDirection)scrollDirection rangeMode:(ASLayoutRangeMode)rangeMode displaySet:(NSSet * _Nullable * _Nullable)displaySet preloadSet:(NSSet * _Nullable * _Nullable)preloadSet;
+
 @optional
 
 - (void)setVisibleNodeIndexPaths:(NSArray<NSIndexPath *> *)indexPaths;

Source/Details/ASLayoutRangeType.h

@@ -67,3 +67,6 @@ typedef NS_ENUM(NSInteger, ASLayoutRangeType) {
 };
 
 static NSInteger const ASLayoutRangeTypeCount = 2;
+
+#define ASLayoutRangeTypeRender ASLayoutRangeTypeDisplay
+#define ASLayoutRangeTypeFetchData ASLayoutRangeTypePreload

Source/Details/ASRangeController.mm

@@ -238,14 +238,6 @@ static UIApplicationState __ApplicationState = UIApplicationStateActive;
     [_layoutController setVisibleNodeIndexPaths:visibleNodePaths];
   }
   
-  NSSet<NSIndexPath *> *visibleIndexPaths = [NSSet setWithArray:visibleNodePaths];
-  NSSet<NSIndexPath *> *displayIndexPaths = nil;
-  NSSet<NSIndexPath *> *preloadIndexPaths = nil;
-  
-  // Prioritize the order in which we visit each.  Visible nodes should be updated first so they are enqueued on
-  // the network or display queues before preloading (offscreen) nodes are enqueued.
-  NSMutableOrderedSet<NSIndexPath *> *allIndexPaths = [[NSMutableOrderedSet alloc] initWithSet:visibleIndexPaths];
-  
   ASInterfaceState selfInterfaceState = [self interfaceState];
   ASLayoutRangeMode rangeMode = _currentRangeMode;
   // If the range mode is explicitly set via updateCurrentRangeWithMode: it will last in that mode until the
@@ -255,28 +247,49 @@ static UIApplicationState __ApplicationState = UIApplicationStateActive;
   }
 
   ASRangeTuningParameters parametersPreload = [_layoutController tuningParametersForRangeMode:rangeMode
-                                                                                      rangeType:ASLayoutRangeTypePreload];
-  if (ASRangeTuningParametersEqualToRangeTuningParameters(parametersPreload, ASRangeTuningParametersZero)) {
-    preloadIndexPaths = visibleIndexPaths;
-  } else {
-    preloadIndexPaths = [_layoutController indexPathsForScrolling:scrollDirection
-                                                          rangeMode:rangeMode
-                                                          rangeType:ASLayoutRangeTypePreload];
-  }
-  
+                                                                                    rangeType:ASLayoutRangeTypePreload];
   ASRangeTuningParameters parametersDisplay = [_layoutController tuningParametersForRangeMode:rangeMode
                                                                                     rangeType:ASLayoutRangeTypeDisplay];
-  if (rangeMode == ASLayoutRangeModeLowMemory) {
-    displayIndexPaths = [NSSet set];
-  } else if (ASRangeTuningParametersEqualToRangeTuningParameters(parametersDisplay, ASRangeTuningParametersZero)) {
-    displayIndexPaths = visibleIndexPaths;
-  } else if (ASRangeTuningParametersEqualToRangeTuningParameters(parametersDisplay, parametersPreload)) {
-    displayIndexPaths = preloadIndexPaths;
+
+  // Preload can express the ultra-low-memory state with 0, 0 returned for its tuningParameters above, and will match Visible.
+  // However, in this rangeMode, Display is not supposed to contain *any* paths -- not even the visible bounds. TuningParameters can't express this.
+  BOOL emptyDisplayRange = (rangeMode == ASLayoutRangeModeLowMemory);
+  BOOL equalDisplayPreload = ASRangeTuningParametersEqualToRangeTuningParameters(parametersDisplay, parametersPreload);
+  BOOL equalDisplayVisible = (ASRangeTuningParametersEqualToRangeTuningParameters(parametersDisplay, ASRangeTuningParametersZero)
+                              && emptyDisplayRange == NO);
+
+  // Check if both Display and Preload are unique. If they are, we load them with a single fetch from the layout controller for performance.
+  BOOL optimizedLoadingOfBothRanges = (equalDisplayPreload == NO && equalDisplayVisible == NO && emptyDisplayRange == NO);
+
+  NSSet<NSIndexPath *> *visibleIndexPaths = [NSSet setWithArray:visibleNodePaths];
+  NSSet<NSIndexPath *> *displayIndexPaths = nil;
+  NSSet<NSIndexPath *> *preloadIndexPaths = nil;
+  
+  if (optimizedLoadingOfBothRanges) {
+    [_layoutController allIndexPathsForScrolling:scrollDirection rangeMode:rangeMode displaySet:&displayIndexPaths preloadSet:&preloadIndexPaths];
   } else {
-    displayIndexPaths = [_layoutController indexPathsForScrolling:scrollDirection
-                                                        rangeMode:rangeMode
-                                                        rangeType:ASLayoutRangeTypeDisplay];
+    if (emptyDisplayRange == YES) {
+      displayIndexPaths = [NSSet set];
+    } if (equalDisplayVisible == YES) {
+      displayIndexPaths = visibleIndexPaths;
+    } else {
+      // Calculating only the Display range means the Preload range is either the same as Display or Visible.
+      displayIndexPaths = [_layoutController indexPathsForScrolling:scrollDirection rangeMode:rangeMode rangeType:ASLayoutRangeTypeDisplay];
+    }
+    
+    BOOL equalPreloadVisible = ASRangeTuningParametersEqualToRangeTuningParameters(parametersPreload, ASRangeTuningParametersZero);
+    if (equalDisplayPreload == YES) {
+      preloadIndexPaths = displayIndexPaths;
+    } else if (equalPreloadVisible == YES) {
+      preloadIndexPaths = visibleIndexPaths;
+    } else {
+      preloadIndexPaths = [_layoutController indexPathsForScrolling:scrollDirection rangeMode:rangeMode rangeType:ASLayoutRangeTypePreload];
+    }
   }
+
+  // Prioritize the order in which we visit each.  Visible nodes should be updated first so they are enqueued on
+  // the network or display queues before preloading (offscreen) nodes are enqueued.
+  NSMutableOrderedSet<NSIndexPath *> *allIndexPaths = [[NSMutableOrderedSet alloc] initWithSet:visibleIndexPaths];
   
   // Typically the preloadIndexPaths will be the largest, and be a superset of the others, though it may be disjoint.
   // Because allIndexPaths is an NSMutableOrderedSet, this adds the non-duplicate items /after/ the existing items.

Source/Details/ASTableLayoutController.m

@@ -64,6 +64,17 @@
   return indexPaths;
 }
 
+- (void)allIndexPathsForScrolling:(ASScrollDirection)scrollDirection rangeMode:(ASLayoutRangeMode)rangeMode displaySet:(NSSet **)displaySet preloadSet:(NSSet **)preloadSet
+{
+  if (displaySet == NULL || preloadSet == NULL) {
+    return;
+  }
+
+  *displaySet = [self indexPathsForScrolling:scrollDirection rangeMode:rangeMode rangeType:ASLayoutRangeTypeDisplay];
+  *preloadSet = [self indexPathsForScrolling:scrollDirection rangeMode:rangeMode rangeType:ASLayoutRangeTypePreload];
+  return;
+}
+
 #pragma mark - Utility
 
 - (NSIndexPath *)findIndexPathAtDistance:(CGFloat)distance fromIndexPath:(NSIndexPath *)start