Use the sticky wrapper's layout for sticky header cells (#59199)

packages/react-native/Libraries/Lists/__tests__/FlatList-itest.js

@@ -584,6 +584,65 @@ describe('<FlatList>', () => {
       expect(viewableKeys).toContain('3');
     });
 
+    it('reports the next sticky header and the items after it', () => {
+      // Two groups, each one 50px sticky header followed by five 80px rows:
+      // header A at 0, rows A1-A5, header B at 450, rows B1-B5.
+      const data = ['A', 'B'].flatMap(group => [
+        {key: `header ${group}`, height: 50},
+        ...Array.from({length: 5}, (_, i) => ({
+          key: `${group}${i + 1}`,
+          height: 80,
+        })),
+      ]);
+      const root = Fantom.createRoot({
+        viewportWidth: 400,
+        viewportHeight: 600,
+      });
+      const onViewableItemsChanged = jest.fn();
+      const flatListRef = createRef<FlatList<{key: string, height: number}>>();
+      const timers = Fantom.installTimerMock();
+      try {
+        Fantom.runTask(() => {
+          root.render(
+            <FlatList
+              ref={flatListRef}
+              data={data}
+              renderItem={({item}) => (
+                <View style={{height: item.height}} collapsable={false} />
+              )}
+              stickyHeaderIndices={[0, 6]}
+              initialNumToRender={data.length}
+              viewabilityConfig={{itemVisiblePercentThreshold: 100}}
+              onViewableItemsChanged={onViewableItemsChanged}
+            />,
+          );
+        });
+
+        const scrollView = nullthrows(
+          nullthrows(flatListRef.current).getNativeScrollRef(),
+        );
+
+        // Header A is pinned. Header B (at 450) and B1-B4 are fully on screen.
+        Fantom.scrollTo(scrollView, {x: 0, y: 254});
+        // Flush the render batches scheduled by cell layout events.
+        timers.runAllTimers();
+        Fantom.runWorkLoop();
+
+        const lastCall =
+          onViewableItemsChanged.mock.calls[
+            onViewableItemsChanged.mock.calls.length - 1
+          ][0];
+        const viewableKeys = lastCall.viewableItems.map(
+          (item: {key: string, ...}) => item.key,
+        );
+        for (const key of ['header B', 'B1', 'B2', 'B3', 'B4']) {
+          expect(viewableKeys).toContain(key);
+        }
+      } finally {
+        timers.uninstall();
+      }
+    });
+
     it('supports viewabilityConfigCallbackPairs', () => {
       const root = Fantom.createRoot({
         viewportWidth: 400,

packages/react-native/Libraries/Lists/__tests__/SectionList-itest.js

@@ -16,7 +16,7 @@ import * as Fantom from '@react-native/fantom';
 import nullthrows from 'nullthrows';
 import * as React from 'react';
 import {createRef} from 'react';
-import {SectionList, Text} from 'react-native';
+import {SectionList, Text, View} from 'react-native';
 
 type Item = {key: string};
 
@@ -150,6 +150,75 @@ describe('<SectionList>', () => {
         </rn-scrollView>,
       );
     });
+
+    it('reports the next sticky section header and the items after it', () => {
+      // Two sections, each one 50px sticky header followed by five 80px rows:
+      // header A at 0, rows A1-A5, header B at 450, rows B1-B5.
+      const sections: Array<SectionBase<Item>> = ['A', 'B'].map(name => ({
+        key: name,
+        data: Array.from({length: 5}, (_, i) => ({key: `${name}${i + 1}`})),
+      }));
+      const root = Fantom.createRoot({
+        viewportWidth: 400,
+        viewportHeight: 600,
+      });
+      const onViewableItemsChanged = jest.fn();
+      const sectionListRef = createRef<SectionList<Item>>();
+      const timers = Fantom.installTimerMock();
+      try {
+        Fantom.runTask(() => {
+          root.render(
+            <SectionList
+              ref={sectionListRef}
+              stickySectionHeadersEnabled={true}
+              sections={sections}
+              renderItem={() => (
+                <View style={{height: 80}} collapsable={false} />
+              )}
+              renderSectionHeader={() => (
+                <View style={{height: 50}} collapsable={false} />
+              )}
+              initialNumToRender={20}
+              viewabilityConfig={{itemVisiblePercentThreshold: 100}}
+              onViewableItemsChanged={onViewableItemsChanged}
+            />,
+          );
+        });
+
+        const scrollView = nullthrows(
+          nullthrows(
+            nullthrows(sectionListRef.current).getScrollResponder(),
+          ).getNativeScrollRef(),
+        );
+
+        // Header A is pinned. Header B (at 450) and B1-B4 are fully on screen.
+        Fantom.scrollTo(scrollView, {x: 0, y: 254});
+        // Flush the render batches scheduled by cell layout events.
+        timers.runAllTimers();
+        Fantom.runWorkLoop();
+
+        const lastCall =
+          onViewableItemsChanged.mock.calls[
+            onViewableItemsChanged.mock.calls.length - 1
+          ][0];
+        const viewableLabels = lastCall.viewableItems.map(
+          (token: {
+            index: ?number,
+            item: Item,
+            section: {key: string, ...},
+            ...
+          }) =>
+            token.index == null
+              ? `header ${token.section.key}`
+              : token.item.key,
+        );
+        for (const label of ['header B', 'B1', 'B2', 'B3', 'B4']) {
+          expect(viewableLabels).toContain(label);
+        }
+      } finally {
+        timers.uninstall();
+      }
+    });
   });
 
   describe('renders section footer', () => {

packages/virtualized-lists/Lists/VirtualizedList.js

@@ -959,8 +959,12 @@ class VirtualizedList extends StateSafePureComponent<
       const key = VirtualizedList._keyExtractor(item, ii, this.props);
 
       this._indicesToKeys.set(ii, key);
-      if (stickyIndicesFromProps?.has(ii + stickyOffset)) {
+      const isStickyHeader =
+        stickyIndicesFromProps?.has(ii + stickyOffset) === true;
+      if (isStickyHeader) {
         stickyHeaderIndices.push(cells.length);
+      } else {
+        this._stickyHeaderKeysLaidOutByWrapper.delete(key);
       }
 
       const shouldListenForLayout =
@@ -988,6 +992,14 @@ class VirtualizedList extends StateSafePureComponent<
           {...(shouldListenForLayout && {
             onCellLayout: this._onCellLayout,
           })}
+          // `ScrollViewStickyHeader` calls this with the layout of its wrapper,
+          // which is relative to the list content.
+          onLayout={
+            shouldListenForLayout && isStickyHeader
+              ? (event: LayoutChangeEvent) =>
+                  this._onStickyHeaderCellLayout(event, key, ii)
+              : undefined
+          }
         />,
       );
       prevCellKey = key;
@@ -1481,6 +1493,8 @@ class VirtualizedList extends StateSafePureComponent<
   _scrollRef: ?React.ElementRef<typeof ScrollView> = null;
   _sentStartForContentLength = 0;
   _sentEndForContentLength = 0;
+  // Sticky header cells whose position comes from their sticky wrapper.
+  _stickyHeaderKeysLaidOutByWrapper: Set<string> = new Set();
   _updateCellsToRenderTimeoutID: ?ReturnType<typeof setTimeout> = null;
   _viewabilityTuples: Array<ViewabilityHelperCallbackTuple> = [];
 
@@ -1543,6 +1557,28 @@ class VirtualizedList extends StateSafePureComponent<
     cellKey: string,
     cellIndex: number,
   ): void => {
+    // A sticky header cell measures itself relative to its sticky wrapper, so
+    // prefer the wrapper's layout once it has been reported.
+    if (this._stickyHeaderKeysLaidOutByWrapper.has(cellKey)) {
+      return;
+    }
+    this._handleCellLayout(e, cellKey, cellIndex);
+  };
+
+  _onStickyHeaderCellLayout = (
+    e: LayoutChangeEvent,
+    cellKey: string,
+    cellIndex: number,
+  ): void => {
+    this._stickyHeaderKeysLaidOutByWrapper.add(cellKey);
+    this._handleCellLayout(e, cellKey, cellIndex);
+  };
+
+  _handleCellLayout(
+    e: LayoutChangeEvent,
+    cellKey: string,
+    cellIndex: number,
+  ): void {
     const layoutHasChanged = this._listMetrics.notifyCellLayout({
       cellIndex,
       cellKey,
@@ -1558,7 +1594,7 @@ class VirtualizedList extends StateSafePureComponent<
     this._computeBlankness();
     this._updateViewableItems(this.props, this.state.cellsAroundViewport);
     this._notifyCrossOrientationChildren();
-  };
+  }
 
   _onCellFocusCapture = (cellKey: string) => {
     this._lastFocusedCellKey = cellKey;
@@ -1572,6 +1608,7 @@ class VirtualizedList extends StateSafePureComponent<
 
   _onCellUnmount = (cellKey: string) => {
     delete this._cellRefs[cellKey];
+    this._stickyHeaderKeysLaidOutByWrapper.delete(cellKey);
     this._listMetrics.notifyCellUnmounted(cellKey);
   };
 

packages/virtualized-lists/Lists/VirtualizedListCellRenderer.js

@@ -40,6 +40,9 @@ export type Props<ItemT> = {
     index: number,
   ) => void,
   onCellFocusCapture?: (cellKey: string) => void,
+  // Not used by CellRenderer. A sticky header wrapper reads it from the cell
+  // element and calls it with the wrapper's layout.
+  onLayout?: (event: LayoutChangeEvent) => void,
   onUnmount: (cellKey: string) => void,
   onUpdateSeparators: (
     cellKeys: Array<?string>,

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

@@ -9,6 +9,7 @@
  */
 
 import VirtualizedList from '../VirtualizedList';
+import CellRenderer from '../VirtualizedListCellRenderer';
 import {format} from 'node:util';
 import * as React from 'react';
 import {createElement, createRef} from 'react';
@@ -1140,6 +1141,107 @@ describe('VirtualizedList', () => {
       {first: 13, last: 19, isSpacer: true},
     ]);
   });
+
+  it('uses the sticky wrapper layout for sticky header cells', async () => {
+    // Sticky headers at 0, 5 and 10. Every cell is 50dip tall.
+    const items = generateItemsStickyEveryN(15, 5);
+
+    let component;
+    await act(() => {
+      component = create(
+        <VirtualizedList initialNumToRender={15} {...baseItemProps(items)} />,
+      );
+    });
+
+    await act(() => {
+      simulateLayout(component, {
+        viewport: {width: 10, height: 100},
+        content: {width: 10, height: 750},
+      });
+      for (let i = 0; i < items.length; i++) {
+        // Sticky header cells measure themselves relative to their wrapper.
+        const y = items[i].sticky === true ? 0 : i * 50;
+        simulateCellLayout(component, items, i, {
+          width: 10,
+          height: 50,
+          x: 0,
+          y,
+        });
+      }
+    });
+
+    // The wrapper reports after the cell's own layout.
+    await act(() => {
+      simulateStickyHeaderWrapperLayout(component, items, 5, {
+        width: 10,
+        height: 50,
+        x: 0,
+        y: 250,
+      });
+    });
+    expect(getCellOffset(component, 5)).toBe(250);
+
+    // The wrapper reports before the cell's own layout.
+    await act(() => {
+      simulateStickyHeaderWrapperLayout(component, items, 10, {
+        width: 10,
+        height: 50,
+        x: 0,
+        y: 500,
+      });
+      simulateCellLayout(component, items, 10, {
+        width: 10,
+        height: 50,
+        x: 0,
+        y: 0,
+      });
+    });
+    expect(getCellOffset(component, 10)).toBe(500);
+
+    // Cells that are not sticky keep using their own layout.
+    expect(getCellOffset(component, 7)).toBe(350);
+  });
+
+  it('uses the cell layout again once a cell is no longer sticky', async () => {
+    const items = generateItemsStickyEveryN(10, 5);
+    const props = {initialNumToRender: 10, ...baseItemProps(items)};
+
+    let component;
+    await act(() => {
+      component = create(<VirtualizedList {...props} />);
+    });
+
+    await act(() => {
+      simulateLayout(component, {
+        viewport: {width: 10, height: 100},
+        content: {width: 10, height: 500},
+      });
+      simulateStickyHeaderWrapperLayout(component, items, 5, {
+        width: 10,
+        height: 50,
+        x: 0,
+        y: 250,
+      });
+    });
+    expect(getCellOffset(component, 5)).toBe(250);
+
+    await act(() => {
+      component.update(<VirtualizedList {...props} stickyHeaderIndices={[]} />);
+    });
+    expect(findCellRenderer(component, items, 5).props.onLayout).toBe(
+      undefined,
+    );
+
+    await act(() => {
+      simulateCellLayout(component, items, 5, {
+        width: 10,
+        height: 50,
+        x: 0,
+        y: 260,
+      });
+    });
+    expect(getCellOffset(component, 5)).toBe(260);
+  });
 });
 
 it('unmounts sticky headers moved below viewport', async () => {
@@ -3141,6 +3243,36 @@ function simulateCellLayout(component, items, itemIndex, dimensions) {
   );
 }
 
+function findCellRenderer(component, items, itemIndex) {
+  const cellKey = VirtualizedList._keyExtractor(
+    items[itemIndex],
+    itemIndex,
+    component.getInstance().props,
+  );
+  return component.root.find(
+    node => node.type === CellRenderer && node.props.cellKey === cellKey,
+  );
+}
+
+// Mimics `ScrollViewStickyHeader`, which calls the cell element's `onLayout`
+// with the layout of its wrapper.
+function simulateStickyHeaderWrapperLayout(
+  component,
+  items,
+  itemIndex,
+  dimensions,
+) {
+  findCellRenderer(component, items, itemIndex).props.onLayout({
+    nativeEvent: {layout: dimensions, zoomScale: 1},
+  });
+}
+
+function getCellOffset(component, itemIndex) {
+  const instance = component.getInstance();
+  return instance.__getListMetrics().getCellMetrics(itemIndex, instance.props)
+    .offset;
+}
+
 function simulateScroll(component, position) {
   component.getInstance()._onScroll({
     nativeEvent: {