fix(BaseTable): adaptive row virtualization improvements

src/components/BaseHeaderRow/BaseHeaderRow.tsx

@@ -10,6 +10,7 @@ import {b} from '../BaseTable/BaseTable.classname';
 import {ColumnReorderingContext} from '../ColumnReorderingContext';
 
 import {getCanReorderHeader} from './utils/getCanReorderHeader';
+import {shouldSkipHeaderRender} from './utils/shouldSkipHeaderRender';
 
 export interface BaseHeaderRowProps<TData, TValue = unknown>
     extends Omit<React.HTMLAttributes<HTMLTableRowElement>, 'className'> {
@@ -24,6 +25,8 @@ export interface BaseHeaderRowProps<TData, TValue = unknown>
     renderSortIndicator: BaseHeaderCellProps<TData, TValue>['renderSortIndicator'];
     resizeHandleClassName?: BaseHeaderCellProps<TData, TValue>['resizeHandleClassName'];
     sortIndicatorClassName: BaseHeaderCellProps<TData, TValue>['sortIndicatorClassName'];
+    /** @internal */
+    tableRenderVersion?: Readonly<Record<string, unknown>>;
     attributes?:
         | React.HTMLAttributes<HTMLTableRowElement>
         | ((
@@ -33,7 +36,7 @@ export interface BaseHeaderRowProps<TData, TValue = unknown>
     cellAttributes?: BaseHeaderCellProps<TData, TValue>['attributes'];
 }
 
-export const BaseHeaderRow = <TData, TValue = unknown>({
+const BaseHeaderRowComponent = <TData, TValue = unknown>({
     cellClassName,
     className: classNameProp,
     headerGroup,
@@ -43,6 +46,7 @@ export const BaseHeaderRow = <TData, TValue = unknown>({
     renderSortIndicator,
     resizeHandleClassName,
     sortIndicatorClassName,
+    tableRenderVersion: _tableRenderVersion,
     attributes: attributesProp,
     cellAttributes,
     ...restProps
@@ -93,3 +97,12 @@ export const BaseHeaderRow = <TData, TValue = unknown>({
         </tr>
     );
 };
+
+export const BaseHeaderRow = React.memo(BaseHeaderRowComponent, shouldSkipHeaderRender) as (<
+    TData,
+    TValue = unknown,
+>(
+    props: BaseHeaderRowProps<TData, TValue>,
+) => React.ReactElement) & {displayName?: string};
+
+BaseHeaderRow.displayName = 'BaseHeaderRow';

src/components/BaseHeaderRow/utils/shouldSkipHeaderRender.ts

@@ -0,0 +1,26 @@
+import {areShallowEqual} from '../../../utils/areShallowEqual';
+import type {BaseHeaderRowProps} from '../BaseHeaderRow';
+
+export const shouldSkipHeaderRender = <TData, TValue>(
+    previousProps: BaseHeaderRowProps<TData, TValue>,
+    nextProps: BaseHeaderRowProps<TData, TValue>,
+) => {
+    // Stable header references do not reflect table state changes without a render version.
+    if (
+        previousProps.tableRenderVersion === undefined ||
+        nextProps.tableRenderVersion === undefined
+    ) {
+        return false;
+    }
+
+    const nextKeys = Object.keys(nextProps) as (keyof BaseHeaderRowProps<TData, TValue>)[];
+
+    return (
+        Object.keys(previousProps).length === nextKeys.length &&
+        nextKeys.every((key) =>
+            key === 'tableRenderVersion'
+                ? areShallowEqual(previousProps.tableRenderVersion, nextProps.tableRenderVersion)
+                : previousProps[key] === nextProps[key],
+        )
+    );
+};

src/components/BaseRow/hooks/useVirtualizedRow.ts

@@ -26,6 +26,7 @@ export function useVirtualizedRow<TData, TScrollElement extends Element | Window
     const runtime = rowVirtualizer ? getRowVirtualizerRuntime(rowVirtualizer) : undefined;
     const directDomUpdates = Boolean(runtime?.directDomUpdates);
     const directDomUpdatesMode = runtime?.directDomUpdatesMode ?? 'transform';
+    const measuredVirtualKeyRef = React.useRef<VirtualItem['key']>();
     const placeholderNodeRef = React.useRef<HTMLTableRowElement>();
     const placeholderVirtualKeyRef = React.useRef<VirtualItem['key']>();
     const virtualIndex = virtualItem?.index;
@@ -52,14 +53,25 @@ export function useVirtualizedRow<TData, TScrollElement extends Element | Window
                 }
             }
 
+            const previousVirtualKey = measuredVirtualKeyRef.current;
+            if (
+                previousVirtualKey !== undefined &&
+                !Object.is(previousVirtualKey, virtualKey) &&
+                rowVirtualizer?.elementsCache.get(previousVirtualKey) === node
+            ) {
+                rowVirtualizer.elementsCache.delete(previousVirtualKey);
+            }
+
             rowVirtualizer?.measureElement(node);
+            measuredVirtualKeyRef.current = virtualKey;
         },
         [
             directDomUpdates,
             directDomUpdatesMode,
             measurementVersion,
             rowVirtualizer,
             virtualIndex,
+            virtualKey,
             virtualItemPosition,
         ],
     );

src/components/BaseTable/BaseTable.tsx

@@ -203,6 +203,7 @@ export const BaseTable = React.forwardRef(
             bodyRows,
             bodyStyle,
             getRowVirtualizationProps,
+            headerRenderVersion,
             resolvedBodyRef,
             virtualizationCoverage,
         } = useTableVirtualization({
@@ -322,6 +323,7 @@ export const BaseTable = React.forwardRef(
                                     cellClassName={headerCellClassName}
                                     className={headerRowClassName}
                                     headerGroup={headerGroup}
+                                    tableRenderVersion={headerRenderVersion}
                                     parentHeaderGroup={headerGroups[index - 1]}
                                     renderHeaderCellContent={renderHeaderCellContent}
                                     renderResizeHandle={renderResizeHandle}

src/components/BaseTable/hooks/useAdaptiveTableVirtualization.ts

@@ -130,11 +130,41 @@ export function useAdaptiveTableVirtualization<TData, TScrollElement extends Ele
             return;
         }
 
+        const previousRows = runtime?.renderedRows;
         if (runtime) {
             runtime.renderedRows = renderedRows;
         }
         controller?.commit(renderPlan ?? null, renderedRows);
-    }, [controller, preparedData, preparedRequiredIndexes, renderedRows, renderPlan, runtime]);
+        if (directDomUpdates && rowVirtualizer && runtime?.bodyElement) {
+            const previousRowsByKey = new Map(
+                previousRows?.map((record) => [record.virtualKey, record]),
+            );
+            const unmeasuredRealRows = renderedRows.filter((record) => {
+                if (record.deferred || rowVirtualizer.itemSizeCache.has(record.virtualKey)) {
+                    return false;
+                }
+                const previousRow = previousRowsByKey.get(record.virtualKey);
+                return (
+                    !previousRow ||
+                    previousRow.deferred ||
+                    !Object.is(previousRow.rowKey, record.rowKey) ||
+                    previousRow.index !== record.index
+                );
+            });
+            // TanStack skips initial measurement during scrolling. Reconcile new real rows
+            // before their estimated height can overlap the following rows.
+            remeasureRenderedRows(rowVirtualizer, runtime.bodyElement, unmeasuredRealRows);
+        }
+    }, [
+        controller,
+        directDomUpdates,
+        preparedData,
+        preparedRequiredIndexes,
+        renderedRows,
+        renderPlan,
+        rowVirtualizer,
+        runtime,
+    ]);
 
     useIsomorphicLayoutEffect(
         () => () => {

src/components/BaseTable/hooks/useTableVirtualization.ts

@@ -147,6 +147,7 @@ export function useTableVirtualization<
 
     return {
         bodyRows,
+        headerRenderVersion: rowRenderVersion,
         bodyStyle: {
             height: getVirtualBodyHeight(bodyRows.length > 0, rowVirtualizer),
             ...bodyStyle,

src/components/BaseTable/utils/createTableRenderVersion.ts

@@ -6,7 +6,7 @@ export function createTableRenderVersion(
     const version: Record<string, unknown> = {__columnGeometry: columnGeometry};
 
     Object.entries(options).forEach(([key, value]) => {
-        if (key !== 'state') {
+        if (key !== 'state' && key !== 'onStateChange') {
             version[`option:${key}`] = value;
         }
     });

src/components/BaseTable/utils/resolveHorizontalScrollElement.ts

@@ -6,7 +6,7 @@ export function resolveHorizontalScrollElement(
     if (!targetWindow) {
         return null;
     }
-    let ancestor = bodyElement?.parentElement;
+    let ancestor: HTMLElement | null = bodyElement;
     while (ancestor) {
         const {overflowX} = targetWindow.getComputedStyle(ancestor);
         if (

src/hooks/useAdaptiveVirtualizer/utils/createAdaptiveVirtualizerController.ts

@@ -471,9 +471,9 @@ export const createAdaptiveVirtualizerController = (
             return;
         }
 
-        // A native terminal token is exact for the gesture that produced it. A new
-        // active direction invalidates it even before the offset has changed.
-        invalidateNativeSettled();
+        if (direction !== lastMotionDirection) {
+            invalidateNativeSettled();
+        }
         recordDirectionReversal(direction);
         updateScrollVelocity(direction, offset, timestamp);
         updateDirectionalSample(direction, offset, timestamp);
@@ -576,20 +576,35 @@ export const createAdaptiveVirtualizerController = (
             return [];
         }
         const direction = virtualizer?.scrollDirection ?? lastActiveDirection;
-        return prioritize(
+        const range = latestRange;
+        const indexes = prioritize(
             [...paintedDeferred].flatMap(([index, identity]) =>
                 sameIdentity(identity, deferred.get(index)) ? [index] : [],
             ),
-            latestRange,
+            range,
             direction,
         );
+        if (getNativeSettled() || !virtualizer?.isScrolling) {
+            const isVisible = (index: number) =>
+                index >= range.startIndex && index <= range.endIndex;
+            return [...indexes.filter(isVisible), ...indexes.filter((index) => !isVisible(index))];
+        }
+        return indexes;
     };
 
+    const getPendingRealizationIndexes = (snapshot: Snapshot | null) =>
+        snapshot
+            ? snapshot.indexes.filter(
+                  (index) => !snapshot.realIndexSet.has(index) && !deferred.has(index),
+              )
+            : [];
+
     const realizePainted = (
         limit: number,
         reason: 'active_recovery' | 'idle_stable' | 'native_settle',
     ) => {
-        const indexes = getRealizationCandidates().slice(0, limit);
+        const remaining = Math.max(0, limit - getPendingRealizationIndexes(committed).length);
+        const indexes = getRealizationCandidates().slice(0, remaining);
         if (indexes.length === 0) {
             return [];
         }
@@ -682,10 +697,25 @@ export const createAdaptiveVirtualizerController = (
             return;
         }
 
+        const direction = virtualizer.scrollDirection ?? lastActiveDirection;
+        const currentTicket = paintTicket;
+
+        if (
+            currentTicket &&
+            currentTicket.virtualizer === virtualizer &&
+            currentTicket.lifecycle === lifecycle &&
+            currentTicket.direction === direction &&
+            currentTicket.identities.size === identities.size &&
+            [...identities].every(([index, identity]) =>
+                sameIdentity(identity, currentTicket.identities.get(index)),
+            )
+        ) {
+            return;
+        }
         cancelPaintTicket('superseded');
         const ticket: PaintTicket = {
             cancel: null,
-            direction: virtualizer.scrollDirection ?? lastActiveDirection,
+            direction,
             id: ++ticketSequence,
             identities,
             lifecycle,
@@ -895,6 +925,13 @@ export const createAdaptiveVirtualizerController = (
                     state.canPrepareDirectionalCoverage),
         );
         const nextDeferred = new Map(deferred);
+        const pendingRealizations = new Set(getPendingRealizationIndexes(snapshot));
+        const usesBoundedRealization =
+            usesDeferredRecovery || nextDeferred.size > 0 || pendingRealizations.size > 0;
+        const realizationLimit = isScrolling
+            ? MAX_REALIZATION_CHUNK_SIZE
+            : MAX_NATIVE_SETTLED_REALIZATION_CHUNK_SIZE;
+        const remainingRealizations = Math.max(0, realizationLimit - pendingRealizations.size);
         const visibleSet = new Set(state.visible);
         const prioritizedRealCritical = [
             ...prioritize(
@@ -907,10 +944,10 @@ export const createAdaptiveVirtualizerController = (
                 range,
                 state.plan.direction,
             ),
-        ];
+        ].filter((index) => !nextDeferred.has(index) && !pendingRealizations.has(index));
         const realCritical = new Set(
-            usesDeferredRecovery
-                ? prioritizedRealCritical.slice(0, MAX_REALIZATION_CHUNK_SIZE)
+            usesBoundedRealization
+                ? prioritizedRealCritical.slice(0, remainingRealizations)
                 : prioritizedRealCritical,
         );
         const realizedDeferredCritical = prioritizedRealCritical.filter(
@@ -923,7 +960,7 @@ export const createAdaptiveVirtualizerController = (
         }
         for (const index of missingCritical) {
             output.add(index);
-            if (usesDeferredRecovery && !realCritical.has(index)) {
+            if (usesBoundedRealization && !realCritical.has(index)) {
                 nextDeferred.set(index, {
                     rowKey: dataRows?.[index]?.id,
                     virtualKey: getKey(index),
@@ -937,7 +974,7 @@ export const createAdaptiveVirtualizerController = (
         }
 
         let warmMissing: number[] = [];
-        if (!recoveryActive && deferred.size === 0) {
+        if (!recoveryActive && deferred.size === 0 && pendingRealizations.size === 0) {
             const mountLimit = isScrolling ? state.plan.mountChunkSize : MAX_WARM_MOUNT_CHUNK_SIZE;
             warmMissing = prioritize(
                 state.target.filter((index) => !output.has(index)),
@@ -996,7 +1033,7 @@ export const createAdaptiveVirtualizerController = (
 
     const getNormalStaleRemoveLimit = (isScrolling: boolean) => {
         if (isScrolling) {
-            return WARM_UNMOUNT_CHUNK_SIZE;
+            return 0;
         }
         if (idleTrimReady) {
             return IDLE_WARM_UNMOUNT_CHUNK_SIZE;
@@ -1017,13 +1054,19 @@ export const createAdaptiveVirtualizerController = (
                 (left, right) =>
                     distanceToRange(right, range) - distanceToRange(left, range) || right - left,
             );
-        const outputLimit = getAdaptiveOutputLimit(
+        const adaptiveOutputLimit = getAdaptiveOutputLimit(
             range,
             state,
             recovery,
             requiredSet.size,
             snapshot,
         );
+        const outputLimit = recovery.isScrolling
+            ? Math.min(
+                  adaptiveOutputLimit,
+                  Math.max(requiredSet.size, state.target.length + getGuard(range.overscan)),
+              )
+            : adaptiveOutputLimit;
         const forcedExcess = Math.max(0, recovery.output.size - outputLimit);
         const staleRemoveLimit = Math.min(
             stale.length,
@@ -1095,7 +1138,7 @@ export const createAdaptiveVirtualizerController = (
     ) => {
         const needsWarmWork = !coversIndexes(state.target, recovery.output);
         const hasStaleOutput = outputIndexes.some((index) => !recovery.targetSet.has(index));
-        if (needsWarmWork || (recovery.isScrolling && hasStaleOutput)) {
+        if (needsWarmWork || (!recovery.isScrolling && idleTrimReady && hasStaleOutput)) {
             scheduleFrame();
         } else if (!recovery.isScrolling && hasStaleOutput) {
             scheduleIdleTrim();