diff --git a/packages/react-stately/src/virtualizer/Virtualizer.ts b/packages/react-stately/src/virtualizer/Virtualizer.ts index 3959c1cc766..bf604620ddc 100644 --- a/packages/react-stately/src/virtualizer/Virtualizer.ts +++ b/packages/react-stately/src/virtualizer/Virtualizer.ts @@ -71,6 +71,7 @@ export class Virtualizer { private _isScrolling: boolean; private _invalidationContext: InvalidationContext; private _overscanManager: OverscanManager; + private _persistedKeyAncestors: Set | null; private _scrollAnchor: ScrollAnchorTracker; // Together they classify whether the changed content was at the anchored (e.g. bottom) edge to avoid @@ -92,6 +93,7 @@ export class Virtualizer { this._isScrolling = false; this._invalidationContext = {}; this._overscanManager = new OverscanManager(); + this._persistedKeyAncestors = null; this._scrollAnchor = new ScrollAnchorTracker(); this._hadItemResize = false; this._batchIncludedEdgeContent = false; @@ -104,23 +106,24 @@ export class Virtualizer { return true; } - // If not, check if the key is an ancestor of any of the persisted keys. - for (let k of this.persistedKeys) { - while (k != null) { - let layoutInfo = this.layout.getLayoutInfo(k); - if (!layoutInfo || layoutInfo.parentKey == null) { - break; - } - - k = layoutInfo.parentKey; + if (this.persistedKeys.size === 0) { + return false; + } - if (k === key) { - return true; + if (!this._persistedKeyAncestors) { + let ancestors = new Set(); + for (let persistedKey of this.persistedKeys) { + let layoutInfo = this.layout.getLayoutInfo(persistedKey); + while (layoutInfo?.parentKey != null) { + let parentKey = layoutInfo.parentKey; + ancestors.add(parentKey); + layoutInfo = this.layout.getLayoutInfo(parentKey); } } + this._persistedKeyAncestors = ancestors; } - return false; + return this._persistedKeyAncestors.has(key); } private getParentView(layoutInfo: LayoutInfo): ReusableView | undefined { @@ -177,6 +180,7 @@ export class Virtualizer { } private relayout(context: InvalidationContext = {}) { + this._persistedKeyAncestors = null; // @ts-ignore let anchorInfo = this.layout.UNSTABLE_getScrollAnchorInfo?.(context.layoutOptions) ?? null; @@ -368,6 +372,7 @@ export class Virtualizer { if (opts.persistedKeys && !isSetEqual(opts.persistedKeys, this.persistedKeys)) { mutableThis.persistedKeys = opts.persistedKeys; + this._persistedKeyAncestors = null; needsUpdate = true; } diff --git a/packages/react-stately/test/virtualizer/Virtualizer.test.ts b/packages/react-stately/test/virtualizer/Virtualizer.test.ts new file mode 100644 index 00000000000..3519afaab26 --- /dev/null +++ b/packages/react-stately/test/virtualizer/Virtualizer.test.ts @@ -0,0 +1,97 @@ +/* + * Copyright 2026 Adobe. All rights reserved. + * This file is licensed to you under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. You may obtain a copy + * of the License at http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software distributed under + * the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR REPRESENTATIONS + * OF ANY KIND, either express or implied. See the License for the specific language + * governing permissions and limitations under the License. + */ + +import {Collection, Key} from '@react-types/shared'; +import {Layout} from '../../src/virtualizer/Layout'; +import {LayoutInfo} from '../../src/virtualizer/LayoutInfo'; +import {Rect} from '../../src/virtualizer/Rect'; +import {Size} from '../../src/virtualizer/Size'; +import {Virtualizer} from '../../src/virtualizer/Virtualizer'; + +class TestLayout extends Layout { + infos = new Map(); + getLayoutInfoCalls = 0; + + getVisibleLayoutInfos(_rect: Rect): LayoutInfo[] { + return []; + } + + getLayoutInfo(key: Key) { + this.getLayoutInfoCalls++; + return this.infos.get(key) ?? null; + } + + getContentSize(): Size { + return new Size(); + } +} + +function createCollection(): Collection { + return { + size: 0, + getKeys: () => [], + getItem: () => null, + at: () => null, + getKeyBefore: () => null, + getKeyAfter: () => null, + getFirstKey: () => null, + getLastKey: () => null, + [Symbol.iterator]: function* () {} + }; +} + +describe('Virtualizer', () => { + it('caches persisted key ancestors and rebuilds them when persisted keys change', () => { + let layout = new TestLayout(); + let addInfo = (key: Key, parentKey: Key | null) => { + let info = new LayoutInfo('item', key, new Rect()); + info.parentKey = parentKey; + layout.infos.set(key, info); + }; + addInfo('leaf', 'parent'); + addInfo('parent', 'root'); + addInfo('root', null); + addInfo('otherLeaf', 'otherParent'); + addInfo('otherParent', 'otherRoot'); + addInfo('otherRoot', null); + + let virtualizer = new Virtualizer({ + delegate: { + setVisibleRect() {}, + renderView: () => null, + invalidate() {} + }, + collection: createCollection(), + layout + }); + virtualizer.persistedKeys.add('leaf'); + + expect(virtualizer.isPersistedKey('parent')).toBe(true); + let getLayoutInfoCalls = layout.getLayoutInfoCalls; + expect(virtualizer.isPersistedKey('root')).toBe(true); + expect(virtualizer.isPersistedKey('unrelated')).toBe(false); + expect(layout.getLayoutInfoCalls).toBe(getLayoutInfoCalls); + + virtualizer.render({ + layout, + collection: virtualizer.collection, + persistedKeys: new Set(['otherLeaf']), + visibleRect: new Rect(), + size: new Size(), + invalidationContext: {}, + isScrolling: false + }); + + expect(virtualizer.isPersistedKey('parent')).toBe(false); + expect(virtualizer.isPersistedKey('otherParent')).toBe(true); + }); +});