feat: Tree multiple level loading support (#8299)

* Update tree and listlayout to handle multi loaders

* adapting other stories to new loader api and adding useAsync example story

* add tests

* fix story for correctness, should only need to provide a dependecy at the top most collection

* restoring focus back to the tree if the user was focused on the loader

* fixing estimated loader position if sections exist

taken from #8326

* skip test for now

Author: LFDanLu

packages/@react-spectrum/autocomplete/test/SearchAutocomplete.test.js

@@ -1844,7 +1844,7 @@ describe('SearchAutocomplete', function () {
       expect(() => within(tray).getByText('No results')).toThrow();
     });
 
-    it.skip('user can select options by pressing them', async function () {
+    it('user can select options by pressing them', async function () {
       let {getByRole, getByText, getByTestId} = renderSearchAutocomplete();
       let button = getByRole('button');
 
@@ -1892,7 +1892,7 @@ describe('SearchAutocomplete', function () {
       expect(items[1]).toHaveAttribute('aria-selected', 'true');
     });
 
-    it.skip('user can select options by focusing them and hitting enter', async function () {
+    it('user can select options by focusing them and hitting enter', async function () {
       let {getByRole, getByText, getByTestId} = renderSearchAutocomplete();
       let button = getByRole('button');
 

packages/@react-stately/layout/src/ListLayout.ts

@@ -253,41 +253,48 @@ export class ListLayout<T, O extends ListLayoutOptions = ListLayoutOptions> exte
 
   protected buildCollection(y = this.padding): LayoutNode[] {
     let collection = this.virtualizer!.collection;
-    let skipped = 0;
+    let collectionNodes = [...collection];
+    let loaderNodes = collectionNodes.filter(node => node.type === 'loader');
     let nodes: LayoutNode[] = [];
-    let isEmptyOrLoading = collection?.size === 0 || (collection.size === 1 && collection.getItem(collection.getFirstKey()!)!.type === 'loader');
+
+    let isEmptyOrLoading = collection?.size === 0 || !collectionNodes.some(item => item.type !== 'loader');
     if (isEmptyOrLoading) {
       y = 0;
     }
 
-    for (let node of collection) {
+    for (let node of collectionNodes) {
       let rowHeight = (this.rowHeight ?? this.estimatedRowHeight ?? DEFAULT_HEIGHT) + this.gap;
       // Skip rows before the valid rectangle unless they are already cached.
       if (node.type === 'item' && y + rowHeight < this.requestedRect.y && !this.isValid(node, y)) {
         y += rowHeight;
-        skipped++;
         continue;
       }
 
       let layoutNode = this.buildChild(node, this.padding, y, null);
       y = layoutNode.layoutInfo.rect.maxY + this.gap;
       nodes.push(layoutNode);
-      if (node.type === 'item' && y > this.requestedRect.maxY) {
-        let itemsAfterRect = collection.size - (nodes.length + skipped);
-        let lastNode = collection.getItem(collection.getLastKey()!);
-        if (lastNode?.type === 'loader') {
-          itemsAfterRect--;
-        }
-
-        y += itemsAfterRect * rowHeight;
+      if (node.type === 'loader') {
+        let index = loaderNodes.indexOf(node);
+        loaderNodes.splice(index, 1);
+      }
 
-        // Always add the loader sentinel if present. This assumes the loader is the last option/row
-        // will need to refactor when handling multi section loading
-        if (lastNode?.type === 'loader' && nodes.at(-1)?.layoutInfo.type !== 'loader') {
-          let loader = this.buildChild(lastNode, this.padding, y, null);
+      // Build each loader that exists in the collection that is outside the visible rect so that they are persisted
+      // at the proper estimated location. If the node.type is "section" then we don't do this shortcut since we have to
+      // build the sections to see how tall they are.
+      if ((node.type === 'item' || node.type === 'loader') && y > this.requestedRect.maxY) {
+        let lastProcessedIndex = collectionNodes.indexOf(node);
+        for (let loaderNode of loaderNodes) {
+          let loaderNodeIndex = collectionNodes.indexOf(loaderNode);
+          // Subtract by an additional 1 since we've already added the current item's height to y
+          y += (loaderNodeIndex - lastProcessedIndex - 1) * rowHeight;
+          let loader = this.buildChild(loaderNode, this.padding, y, null);
           nodes.push(loader);
           y = loader.layoutInfo.rect.maxY;
+          lastProcessedIndex = loaderNodeIndex;
         }
+
+        // Account for the rest of the items after the last loader spinner, subtract by 1 since we've processed the current node's height already
+        y += (collectionNodes.length - lastProcessedIndex - 1) * rowHeight;
         break;
       }
     }

packages/@react-stately/tree/src/useTreeState.ts

@@ -65,8 +65,13 @@ export function useTreeState<T extends object>(props: TreeProps<T>): TreeState<T
 
   // Reset focused key if that item is deleted from the collection.
   useEffect(() => {
-    if (selectionState.focusedKey != null && !tree.getItem(selectionState.focusedKey)) {
-      selectionState.setFocusedKey(null);
+    if (selectionState.focusedKey != null) {
+      let focusedItem = tree.getItem(selectionState.focusedKey);
+      // TODO: do we want to have the same logic as useListState/useGridState where it tries to find the nearest row?
+      // We could possibly special case this loader case and have it try to find the item just before it/the parent
+      if (!focusedItem || focusedItem.type === 'loader' && !focusedItem.props.isLoading) {
+        selectionState.setFocusedKey(null);
+      }
     }
   // eslint-disable-next-line react-hooks/exhaustive-deps
   }, [tree, selectionState.focusedKey]);

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

@@ -20,7 +20,7 @@ import {DisabledBehavior, DragPreviewRenderer, Expandable, forwardRefType, Hover
 import {DragAndDropContext, DropIndicatorContext, useDndPersistedKeys, useRenderDropIndicator} from './DragAndDrop';
 import {DragAndDropHooks} from './useDragAndDrop';
 import {DraggableCollectionState, DroppableCollectionState, Collection as ICollection, Node, SelectionBehavior, TreeState, useTreeState} from 'react-stately';
-import {filterDOMProps, useObjectRef} from '@react-aria/utils';
+import {filterDOMProps, inertValue, LoadMoreSentinelProps, UNSTABLE_useLoadMoreSentinel, useObjectRef} from '@react-aria/utils';
 import React, {createContext, ForwardedRef, forwardRef, JSX, ReactNode, useContext, useEffect, useMemo, useRef} from 'react';
 import {useControlledState} from '@react-stately/utils';
 
@@ -699,51 +699,87 @@ export const TreeItem = /*#__PURE__*/ createBranchComponent('item', <T extends o
   );
 });
 
-export interface UNSTABLE_TreeLoadingIndicatorRenderProps {
+export interface UNSTABLE_TreeLoadingSentinelRenderProps extends Pick<TreeItemRenderProps, 'isFocused' | 'isFocusVisible'> {
   /**
    * What level the tree item has within the tree.
    * @selector [data-level]
    */
   level: number
 }
 
-export interface TreeLoaderProps extends RenderProps<UNSTABLE_TreeLoadingIndicatorRenderProps>, StyleRenderProps<UNSTABLE_TreeLoadingIndicatorRenderProps> {}
+export interface TreeLoadingSentinelProps extends Omit<LoadMoreSentinelProps, 'collection'>, RenderProps<UNSTABLE_TreeLoadingSentinelRenderProps> {
+  /**
+   * The load more spinner to render when loading additional items.
+   */
+  children?: ReactNode | ((values: UNSTABLE_TreeLoadingSentinelRenderProps & {defaultChildren: ReactNode | undefined}) => ReactNode),
+  /**
+   * Whether or not the loading spinner should be rendered or not.
+   */
+  isLoading?: boolean
+}
 
-export const UNSTABLE_TreeLoadingIndicator = createLeafComponent('loader', function TreeLoader<T extends object>(props: TreeLoaderProps,  ref: ForwardedRef<HTMLDivElement>, item: Node<T>) {
+export const UNSTABLE_TreeLoadingSentinel = createLeafComponent('loader', function TreeLoadingSentinel<T extends object>(props: TreeLoadingSentinelProps,  ref: ForwardedRef<HTMLDivElement>, item: Node<T>) {
   let state = useContext(TreeStateContext)!;
-  // This loader row is is non-interactable, but we want the same aria props calculated as a typical row
-  // @ts-ignore
-  let {rowProps} = useTreeItem({node: item}, state, ref);
+  let {isLoading, onLoadMore, scrollOffset, ...otherProps} = props;
+  let sentinelRef = useRef(null);
+  let memoedLoadMoreProps = useMemo(() => ({
+    onLoadMore,
+    // TODO: this collection will update anytime a row is expanded/collapsed becaused the flattenedRows will change.
+    // This means onLoadMore will trigger but that might be ok cause the user should have logic to handle multiple loadMore calls
+    collection: state?.collection,
+    sentinelRef,
+    scrollOffset
+  }), [onLoadMore, scrollOffset, state?.collection]);
+  UNSTABLE_useLoadMoreSentinel(memoedLoadMoreProps, sentinelRef);
+
+  ref = useObjectRef<HTMLDivElement>(ref);
+  let {rowProps, gridCellProps, ...states} = useTreeItem({node: item}, state, ref);
   let level = rowProps['aria-level'] || 1;
 
   let ariaProps = {
+    role: 'row',
     'aria-level': rowProps['aria-level'],
     'aria-posinset': rowProps['aria-posinset'],
-    'aria-setsize': rowProps['aria-setsize']
+    'aria-setsize': rowProps['aria-setsize'],
+    tabIndex: rowProps.tabIndex
   };
 
+  let {isFocusVisible, focusProps} = useFocusRing();
+
   let renderProps = useRenderProps({
-    ...props,
+    ...otherProps,
     id: undefined,
     children: item.rendered,
     defaultClassName: 'react-aria-TreeLoader',
     values: {
-      level
+      level,
+      isFocused: states.isFocused,
+      isFocusVisible
     }
   });
 
   return (
     <>
-      <div
-        role="row"
-        ref={ref}
-        {...mergeProps(filterDOMProps(props as any), ariaProps)}
-        {...renderProps}
-        data-level={level}>
-        <div role="gridcell" aria-colindex={1}>
-          {renderProps.children}
-        </div>
+      {/* Alway render the sentinel. For now onus is on the user for styling when using flex + gap (this would introduce a gap even though it doesn't take room) */}
+      {/* @ts-ignore - compatibility with React < 19 */}
+      <div style={{position: 'relative', width: 0, height: 0}} inert={inertValue(true)} >
+        <div data-testid="loadMoreSentinel" ref={sentinelRef} style={{position: 'absolute', height: 1, width: 1}} />
       </div>
+      {isLoading && renderProps.children && (
+        <div
+          ref={ref}
+          {...mergeProps(filterDOMProps(props as any), ariaProps, focusProps)}
+          {...renderProps}
+          data-key={rowProps['data-key']}
+          data-collection={rowProps['data-collection']}
+          data-focused={states.isFocused || undefined}
+          data-focus-visible={isFocusVisible || undefined}
+          data-level={level}>
+          <div {...gridCellProps}>
+            {renderProps.children}
+          </div>
+        </div>
+      )}
     </>
   );
 });
@@ -796,9 +832,10 @@ function flattenTree<T>(collection: TreeCollection<T>, opts: TreeGridCollectionO
         keyMap.set(node.key, node as CollectionNode<T>);
       }
 
-      if (node.level === 0 || (parentKey != null && expandedKeys.has(parentKey) && flattenedRows.find(row => row.key === parentKey))) {
-        // Grab the modified node from the key map so our flattened list and modified key map point to the same nodes
-        flattenedRows.push(keyMap.get(node.key) || node);
+      // 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))) {
+        flattenedRows.push(modifiedNode);
       }
     } else if (node.type !== null) {
       keyMap.set(node.key, node as CollectionNode<T>);
@@ -836,7 +873,7 @@ function TreeDropIndicatorWrapper(props: DropIndicatorProps, ref: ForwardedRef<H
   let level = dropState && props.target.type === 'item' ? (dropState.collection.getItem(props.target.key)?.level || 0) + 1 : 1;
 
   return (
-    <TreeDropIndicatorForwardRef 
+    <TreeDropIndicatorForwardRef
       {...props}
       dropIndicatorProps={dropIndicatorProps}
       isDropTarget={isDropTarget}

packages/react-aria-components/src/index.ts

@@ -75,7 +75,7 @@ export {ToggleButton, ToggleButtonContext} from './ToggleButton';
 export {ToggleButtonGroup, ToggleButtonGroupContext, ToggleGroupStateContext} from './ToggleButtonGroup';
 export {Toolbar, ToolbarContext} from './Toolbar';
 export {TooltipTrigger, Tooltip, TooltipTriggerStateContext, TooltipContext} from './Tooltip';
-export {UNSTABLE_TreeLoadingIndicator, Tree, TreeItem, TreeContext, TreeItemContent, TreeStateContext} from './Tree';
+export {UNSTABLE_TreeLoadingSentinel, Tree, TreeItem, TreeContext, TreeItemContent, TreeStateContext} from './Tree';
 export {useDragAndDrop} from './useDragAndDrop';
 export {DropIndicator, DropIndicatorContext, DragAndDropContext} from './DragAndDrop';
 export {Virtualizer} from './Virtualizer';