Optimize VirtualizedList sticky header lookup

Author: tarikfp

packages/react-native/Libraries/Lists/__tests__/VirtualizedList-stickyHeaders-benchmark-itest.js

@@ -0,0 +1,85 @@
+/**
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ *
+ * @fantom_mode dev
+ * @flow strict-local
+ * @format
+ */
+
+import '@react-native/fantom/src/setUpDefaultReactNativeEnvironment';
+
+import * as Fantom from '@react-native/fantom';
+import VirtualizedList from '@react-native/virtualized-lists/Lists/VirtualizedList';
+
+const VIEWPORT_SIZE = 100;
+const ROW_COUNTS = [100000, 250000, 500000, 750000, 1000000];
+
+type StickyHeaderCase = {
+  itemCount: number,
+  name: string,
+  stickyHeaderIndices?: ReadonlyArray<number>,
+};
+
+type BenchmarkData = {
+  length: number,
+};
+
+const benchmarkCases: Array<StickyHeaderCase> = [];
+
+for (let i = 0; i < ROW_COUNTS.length; i++) {
+  const itemCount = ROW_COUNTS[i];
+  const label = itemCount === 1000000 ? '1m' : `${itemCount / 1000}k`;
+
+  benchmarkCases.push(
+    {
+      itemCount,
+      name: `${label} rows without sticky headers`,
+    },
+    {
+      itemCount,
+      name: `${label} rows with empty sticky headers`,
+      stickyHeaderIndices: [],
+    },
+    {
+      itemCount,
+      name: `${label} rows with one sticky header at the top`,
+      stickyHeaderIndices: [0],
+    },
+  );
+}
+
+function createProps(
+  itemCount: number,
+  stickyHeaderIndices?: ReadonlyArray<number>,
+) {
+  return {
+    data: {length: itemCount},
+    getItem: (_data: BenchmarkData, index: number) => index,
+    getItemCount: (data: BenchmarkData) => data.length,
+    initialScrollIndex: 1,
+    stickyHeaderIndices,
+  };
+}
+
+Fantom.unstable_benchmark
+  .suite('VirtualizedList sticky headers', {
+    disableOptimizedBuildCheck: true,
+    minIterations: 100,
+  })
+  .test.each(
+    benchmarkCases,
+    benchmarkCase => `create render mask for ${benchmarkCase.name}`,
+    benchmarkCase => {
+      // $FlowExpectedError[prop-missing] Benchmark exercises an internal helper.
+      VirtualizedList._createRenderMask(
+        createProps(benchmarkCase.itemCount, benchmarkCase.stickyHeaderIndices),
+        {
+          first: benchmarkCase.itemCount - VIEWPORT_SIZE,
+          last: benchmarkCase.itemCount - 1,
+        },
+      );
+    },
+  );

packages/virtualized-lists/Lists/VirtualizedList.js

@@ -535,16 +535,18 @@ class VirtualizedList extends StateSafePureComponent<
         renderMask.addCells(initialRegion);
       }
 
-      // The layout coordinates of sticker headers may be off-screen while the
+      // The layout coordinates of sticky headers may be off-screen while the
       // actual header is on-screen. Keep the most recent before the viewport
       // rendered, even if its layout coordinates are not in viewport.
-      const stickyIndicesSet = new Set(props.stickyHeaderIndices);
-      VirtualizedList._ensureClosestStickyHeader(
-        props,
-        stickyIndicesSet,
-        renderMask,
-        cellsAroundViewport.first,
-      );
+      const stickyHeaderIndices = props.stickyHeaderIndices;
+      if (stickyHeaderIndices != null && stickyHeaderIndices.length > 0) {
+        VirtualizedList._ensureClosestStickyHeader(
+          props,
+          stickyHeaderIndices,
+          renderMask,
+          cellsAroundViewport.first,
+        );
+      }
     }
 
     return renderMask;
@@ -575,18 +577,30 @@ class VirtualizedList extends StateSafePureComponent<
 
   static _ensureClosestStickyHeader(
     props: VirtualizedListProps,
-    stickyIndicesSet: Set<number>,
+    stickyHeaderIndices: ReadonlyArray<number>,
     renderMask: CellRenderMask,
     cellIdx: number,
   ) {
     const stickyOffset = props.ListHeaderComponent ? 1 : 0;
+    const targetStickyIndex = cellIdx + stickyOffset;
+    let closestStickyIndex = null;
 
-    for (let itemIdx = cellIdx - 1; itemIdx >= 0; itemIdx--) {
-      if (stickyIndicesSet.has(itemIdx + stickyOffset)) {
-        renderMask.addCells({first: itemIdx, last: itemIdx});
-        break;
+    for (let itemIdx = 0; itemIdx < stickyHeaderIndices.length; itemIdx++) {
+      const stickyIndex = stickyHeaderIndices[itemIdx];
+      if (
+        Number.isInteger(stickyIndex) &&
+        stickyIndex < targetStickyIndex &&
+        stickyIndex >= stickyOffset &&
+        (closestStickyIndex == null || stickyIndex > closestStickyIndex)
+      ) {
+        closestStickyIndex = stickyIndex;
       }
     }
+
+    if (closestStickyIndex != null) {
+      const itemIdx = closestStickyIndex - stickyOffset;
+      renderMask.addCells({first: itemIdx, last: itemIdx});
+    }
   }
 
   _adjustCellsAroundViewport(

packages/virtualized-lists/Lists/__tests__/VirtualizedList-test.js

@@ -988,6 +988,52 @@ describe('VirtualizedList', () => {
     // scrolled-past in layout space.
     expect(component).toMatchSnapshot();
   });
+
+  it('does not add a sticky header to the render mask when no sticky headers are configured', () => {
+    const expectedRegions = [
+      {first: 0, last: 9, isSpacer: true},
+      {first: 10, last: 12, isSpacer: false},
+      {first: 13, last: 19, isSpacer: true},
+    ];
+
+    expect(createRenderMaskForStickyHeaderTest().enumerateRegions()).toEqual(
+      expectedRegions,
+    );
+    expect(
+      createRenderMaskForStickyHeaderTest({
+        stickyHeaderIndices: [],
+      }).enumerateRegions(),
+    ).toEqual(expectedRegions);
+  });
+
+  it('adds the closest sticky header above the viewport from unsorted stickyHeaderIndices', () => {
+    expect(
+      createRenderMaskForStickyHeaderTest({
+        stickyHeaderIndices: [12, 0, 8.5, 7, 7, -1],
+      }).enumerateRegions(),
+    ).toEqual([
+      {first: 0, last: 6, isSpacer: true},
+      {first: 7, last: 7, isSpacer: false},
+      {first: 8, last: 9, isSpacer: true},
+      {first: 10, last: 12, isSpacer: false},
+      {first: 13, last: 19, isSpacer: true},
+    ]);
+  });
+
+  it('accounts for ListHeaderComponent offset when adding the closest sticky header', () => {
+    expect(
+      createRenderMaskForStickyHeaderTest({
+        ListHeaderComponent: () => createElement('Header'),
+        stickyHeaderIndices: [3],
+      }).enumerateRegions(),
+    ).toEqual([
+      {first: 0, last: 1, isSpacer: true},
+      {first: 2, last: 2, isSpacer: false},
+      {first: 3, last: 9, isSpacer: true},
+      {first: 10, last: 12, isSpacer: false},
+      {first: 13, last: 19, isSpacer: true},
+    ]);
+  });
 });
 
 it('unmounts sticky headers moved below viewport', async () => {
@@ -2569,6 +2615,26 @@ function fixedHeightItemLayoutProps(height) {
   };
 }
 
+function createRenderMaskForStickyHeaderTest({
+  ListHeaderComponent,
+  stickyHeaderIndices: stickyHeaderIndicesForTest,
+} = {}) {
+  return VirtualizedList._createRenderMask(
+    {
+      data: {length: 20},
+      getItem: (data, index) => index,
+      getItemCount: data => data.length,
+      initialScrollIndex: 1,
+      ListHeaderComponent,
+      stickyHeaderIndices: stickyHeaderIndicesForTest,
+    },
+    {
+      first: 10,
+      last: 12,
+    },
+  );
+}
+
 let lastViewportLayout;
 let lastContentLayout;