fix: Speed up tree collection builder flattening when expanding keys

packages/react-aria-components/src/Tree.tsx

@@ -818,6 +818,7 @@ function flattenTree<T>(collection: TreeCollection<T>, opts: TreeGridCollectionO
   let flattenedRows: Node<T>[] = [];
   // Need to count the items here because BaseCollection will return the full item count regardless if items are hidden via collapsed rows
   let itemCount = 0;
+  let parentLookup: Map<Key, boolean> = new Map();
 
   let visitNode = (node: Node<T>) => {
     if (node.type === 'item' || node.type === 'loader') {
@@ -844,12 +845,13 @@ function flattenTree<T>(collection: TreeCollection<T>, opts: TreeGridCollectionO
 
       // Grab the modified node from the key map so our flattened list and modified key map point to the same nodes
       let modifiedNode = keyMap.get(node.key) || node;
-      if (modifiedNode.level === 0 || (modifiedNode.parentKey != null && expandedKeys.has(modifiedNode.parentKey) && flattenedRows.find(row => row.key === modifiedNode.parentKey))) {
+      if (modifiedNode.level === 0 || (modifiedNode.parentKey != null && expandedKeys.has(modifiedNode.parentKey) && parentLookup.get(modifiedNode.parentKey))) {
         if (modifiedNode.type === 'item') {
           itemCount++;
         }
 
         flattenedRows.push(modifiedNode);
+        parentLookup.set(modifiedNode.key, true);
       }
     } else if (node.type !== null) {
       keyMap.set(node.key, node as CollectionNode<T>);

packages/react-aria-components/stories/Tree.stories.tsx

@@ -426,7 +426,7 @@ let defaultExpandedKeys = new Set(['projects', 'project-2', 'project-5', 'report
 
 const TreeExampleDynamicRender = <T extends object>(args: TreeProps<T>): JSX.Element => {
   let treeData = useTreeData<any>({
-    initialItems: rows,
+    initialItems: args.items as any ?? rows,
     getKey: item => item.id,
     getChildren: item => item.childItems
   });
@@ -1196,3 +1196,80 @@ export const TreeWithDragAndDropVirtualized = {
   render: TreeDragAndDropVirtualizedRender,
   name: 'Tree with drag and drop (virtualized)'
 };
+
+
+interface ITreeItem {
+  id: string,
+  name: string,
+  childItems?: Iterable<ITreeItem>
+}
+
+let totalItems = 0;
+let itemKeys = new Set<any>();
+/**
+ * Generates a tree data structure with 10 items per level and 6 levels deep.
+ * @returns Array of tree items with the specified structure.
+ */
+function generateTreeData(): Array<ITreeItem> {
+  /**
+   * Recursively generates tree items for a given level.
+   * @param level - Current depth level (1-6).
+   * @param parentId - Parent item ID for generating unique child IDs.
+   * @returns Array of tree items for this level.
+   */
+  function generateLevel(level: number, parentId: string = ''): Array<ITreeItem> {
+    const items: ITreeItem[] = [];
+    const itemsPerLevel = 7;
+
+    for (let i = 1; i < itemsPerLevel; i++) {
+      const itemId = parentId ? `${parentId}-${i}` : `level-${level}-item-${i}`;
+      const itemName = `Level ${level} Item ${i}`;
+
+      const item: ITreeItem = {
+        id: itemId,
+        name: itemName
+      };
+
+      // Add child items if not at the deepest level (level 6)
+      if (level < 6) {
+        item.childItems = generateLevel(level + 1, itemId);
+      }
+
+      totalItems++;
+      items.push(item);
+      itemKeys.add(itemId);
+    }
+
+    return items;
+  }
+
+  return generateLevel(1);
+}
+
+const treeData = generateTreeData();
+console.log(`Total items: ${totalItems}`);
+
+function HugeVirtualizedTreeRender<T extends object>(args: TreeProps<T>): JSX.Element {
+  let [expandedKeys, setExpandedKeys] = useState(new Set<Key>());
+  let expandAll = () => {
+    setExpandedKeys(itemKeys);
+  };
+
+  return (
+    <>
+      <button onClick={expandAll}>Expand All {totalItems}</button>
+      <Virtualizer layout={ListLayout} layoutOptions={{rowHeight: 30}}>
+        <TreeExampleDynamicRender {...args} expandedKeys={expandedKeys} onExpandedChange={setExpandedKeys} />
+      </Virtualizer>
+    </>
+  );
+}
+
+export const HugeVirtualizedTree: StoryObj<typeof VirtualizedTreeRender> = {
+  ...TreeExampleDynamic,
+  args: {
+    ...TreeExampleDynamic.args,
+    items: treeData
+  },
+  render: (args) => <HugeVirtualizedTreeRender {...args} />
+};