perf: cache persisted key ancestors

packages/react-stately/src/virtualizer/Virtualizer.ts

@@ -71,6 +71,7 @@ export class Virtualizer<T extends object, V> {
   private _isScrolling: boolean;
   private _invalidationContext: InvalidationContext;
   private _overscanManager: OverscanManager;
+  private _persistedKeyAncestors: Set<Key> | 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<T extends object, V> {
     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<T extends object, V> {
       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<Key>();
+      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<T, V> | undefined {
@@ -177,6 +180,7 @@ export class Virtualizer<T extends object, V> {
   }
 
   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<T extends object, V> {
 
     if (opts.persistedKeys && !isSetEqual(opts.persistedKeys, this.persistedKeys)) {
       mutableThis.persistedKeys = opts.persistedKeys;
+      this._persistedKeyAncestors = null;
       needsUpdate = true;
     }
 

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<object> {
+  infos = new Map<Key, LayoutInfo>();
+  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<object> {
+  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);
+  });
+});