feat: Tree multiple level loading support and only count "items" for collection size

packages/@react-aria/collections/src/BaseCollection.ts

@@ -79,9 +79,10 @@ export class BaseCollection<T> implements ICollection<Node<T>> {
   private firstKey: Key | null = null;
   private lastKey: Key | null = null;
   private frozen = false;
+  private itemCount: number = 0;
 
   get size(): number {
-    return this.keyMap.size;
+    return this.itemCount;
   }
 
   getKeys(): IterableIterator<Key> {
@@ -184,6 +185,7 @@ export class BaseCollection<T> implements ICollection<Node<T>> {
     collection.keyMap = new Map(this.keyMap);
     collection.firstKey = this.firstKey;
     collection.lastKey = this.lastKey;
+    collection.itemCount = this.itemCount;
     return collection;
   }
 
@@ -192,6 +194,10 @@ export class BaseCollection<T> implements ICollection<Node<T>> {
       throw new Error('Cannot add a node to a frozen collection');
     }
 
+    if (node.type === 'item' && this.keyMap.get(node.key) == null) {
+      this.itemCount++;
+    }
+
     this.keyMap.set(node.key, node);
   }
 
@@ -200,6 +206,11 @@ export class BaseCollection<T> implements ICollection<Node<T>> {
       throw new Error('Cannot remove a node to a frozen collection');
     }
 
+    let node = this.keyMap.get(key);
+    if (node != null && node.type === 'item') {
+      this.itemCount--;
+    }
+
     this.keyMap.delete(key);
   }
 

packages/@react-aria/gridlist/package.json

@@ -32,7 +32,6 @@
     "@react-aria/interactions": "^3.25.3",
     "@react-aria/selection": "^3.24.3",
     "@react-aria/utils": "^3.29.1",
-    "@react-stately/collections": "^3.12.5",
     "@react-stately/list": "^3.12.3",
     "@react-stately/tree": "^3.9.0",
     "@react-types/shared": "^3.30.0",

packages/@react-aria/gridlist/src/useGridListItem.ts

@@ -13,7 +13,6 @@
 import {chain, getScrollParent, mergeProps, scrollIntoViewport, useSlotId, useSyntheticLinkProps} from '@react-aria/utils';
 import {DOMAttributes, FocusableElement, Key, RefObject, Node as RSNode} from '@react-types/shared';
 import {focusSafely, getFocusableTreeWalker} from '@react-aria/focus';
-import {getLastItem} from '@react-stately/collections';
 import {getRowId, listMap} from './utils';
 import {HTMLAttributes, KeyboardEvent as ReactKeyboardEvent, useRef} from 'react';
 import {isFocusVisible} from '@react-aria/interactions';
@@ -104,13 +103,14 @@ export function useGridListItem<T>(props: AriaGridListItemOptions, state: ListSt
     if (node.level > 0 && node?.parentKey != null) {
       let parent = state.collection.getItem(node.parentKey);
       if (parent) {
-        // siblings must exist because our original node exists, same with lastItem
+        // siblings must exist because our original node exists
         let siblings = state.collection.getChildren?.(parent.key)!;
-        setSize = getLastItem(siblings)!.index + 1;
+        setSize = [...siblings].filter(row => row.type === 'item').length;
       }
     } else {
-      setSize = ([...state.collection].filter(row => row.level === 0).at(-1)?.index ?? 0) + 1;
+      setSize = [...state.collection].filter(row => row.level === 0 && row.type === 'item').length;
     }
+
     treeGridRowProps = {
       'aria-expanded': isExpanded,
       'aria-level': node.level + 1,

packages/@react-aria/utils/src/index.ts

@@ -45,7 +45,7 @@ export {useEffectEvent} from './useEffectEvent';
 export {useDeepMemo} from './useDeepMemo';
 export {useFormReset} from './useFormReset';
 export {useLoadMore} from './useLoadMore';
-export {UNSTABLE_useLoadMoreSentinel} from './useLoadMoreSentinel';
+export {useLoadMoreSentinel} from './useLoadMoreSentinel';
 export {inertValue} from './inertValue';
 export {CLEAR_FOCUS_EVENT, FOCUS_EVENT} from './constants';
 export {isCtrlKeyPressed} from './keyboard';

packages/@react-aria/utils/src/useLoadMoreSentinel.ts

@@ -28,7 +28,7 @@ export interface LoadMoreSentinelProps extends Omit<AsyncLoadable, 'isLoading'>
   scrollOffset?: number
 }
 
-export function UNSTABLE_useLoadMoreSentinel(props: LoadMoreSentinelProps, ref: RefObject<HTMLElement | null>): void {
+export function useLoadMoreSentinel(props: LoadMoreSentinelProps, ref: RefObject<HTMLElement | null>): void {
   let {collection, onLoadMore, scrollOffset = 1} = props;
 
   let sentinelObserver = useRef<IntersectionObserver>(null);

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-spectrum/s2/src/CardView.tsx

@@ -16,9 +16,9 @@ import {
   ContextValue,
   GridLayout,
   GridListItem,
+  GridListLoadMoreItem,
   GridListProps,
   Size,
-  UNSTABLE_GridListLoadingSentinel,
   Virtualizer,
   WaterfallLayout
 } from 'react-aria-components';
@@ -246,7 +246,7 @@ export const CardView = /*#__PURE__*/ (forwardRef as forwardRefType)(function Ca
 
   let renderer;
   let cardLoadingSentinel = (
-    <UNSTABLE_GridListLoadingSentinel
+    <GridListLoadMoreItem
       onLoadMore={onLoadMore} />
   );
 

packages/@react-spectrum/s2/src/ComboBox.tsx

@@ -24,12 +24,12 @@ import {
   ListBox,
   ListBoxItem,
   ListBoxItemProps,
+  ListBoxLoadMoreItem,
   ListBoxProps,
   ListLayout,
   ListStateContext,
   Provider,
   SectionProps,
-  UNSTABLE_ListBoxLoadingSentinel,
   Virtualizer
 } from 'react-aria-components';
 import {AsyncLoadable, HelpTextProps, LoadingState, SpectrumLabelableProps} from '@react-types/shared';
@@ -542,7 +542,7 @@ const ComboboxInner = forwardRef(function ComboboxInner(props: ComboBoxProps<any
 
   let renderer;
   let listBoxLoadingCircle = (
-    <UNSTABLE_ListBoxLoadingSentinel
+    <ListBoxLoadMoreItem
       // Only show the spinner in the list when loading more
       isLoading={loadingState === 'loadingMore'}
       onLoadMore={onLoadMore}
@@ -553,7 +553,7 @@ const ComboboxInner = forwardRef(function ComboboxInner(props: ComboBoxProps<any
         styles={progressCircleStyles({size})}
         // Same loading string as table
         aria-label={stringFormatter.format('table.loadingMore')} />
-    </UNSTABLE_ListBoxLoadingSentinel>
+    </ListBoxLoadMoreItem>
   );
 
   if (typeof children === 'function') {

packages/@react-spectrum/s2/src/Picker.tsx

@@ -23,12 +23,12 @@ import {
   ListBox,
   ListBoxItem,
   ListBoxItemProps,
+  ListBoxLoadMoreItem,
   ListBoxProps,
   ListLayout,
   Provider,
   SectionProps,
   SelectValue,
-  UNSTABLE_ListBoxLoadingSentinel,
   Virtualizer
 } from 'react-aria-components';
 import {AsyncLoadable, FocusableRef, FocusableRefValue, HelpTextProps, LoadingState, PressEvent, RefObject, SpectrumLabelableProps} from '@react-types/shared';
@@ -307,12 +307,12 @@ export const Picker = /*#__PURE__*/ (forwardRef as forwardRefType)(function Pick
   let spinnerId = useSlotId([showButtonSpinner]);
 
   let listBoxLoadingCircle = (
-    <UNSTABLE_ListBoxLoadingSentinel
+    <ListBoxLoadMoreItem
       className={loadingWrapperStyles}
       isLoading={loadingState === 'loadingMore'}
       onLoadMore={onLoadMore}>
       <PickerProgressCircle size={size} aria-label={stringFormatter.format('table.loadingMore')} />
-    </UNSTABLE_ListBoxLoadingSentinel>
+    </ListBoxLoadMoreItem>
   );
 
   if (typeof children === 'function' && items) {

packages/@react-spectrum/s2/src/TableView.tsx

@@ -38,8 +38,8 @@ import {
   RowRenderProps,
   TableBodyRenderProps,
   TableLayout,
+  TableLoadMoreItem,
   TableRenderProps,
-  UNSTABLE_TableLoadingSentinel,
   useSlottedContext,
   useTableOptions,
   Virtualizer
@@ -196,11 +196,10 @@ export class S2TableLayout<T> extends TableLayout<T> {
     if (!header) {
       return [];
     }
-    let {children, layoutInfo} = body;
+    let {layoutInfo} = body;
     // TableLayout's buildCollection always sets the body width to the max width between the header width, but
     // we want the body to be sticky and only as wide as the table so it is always in view if loading/empty
-    // TODO: we may want to adjust RAC layouts to do something simlar? Current users of RAC table will probably run into something similar
-    let isEmptyOrLoading = children?.length === 0 || (children?.length === 1 && children[0].layoutInfo.type === 'loader');
+    let isEmptyOrLoading = this.virtualizer?.collection.size === 0;
     if (isEmptyOrLoading) {
       layoutInfo.rect.width = this.virtualizer!.visibleRect.width - 80;
     }
@@ -219,19 +218,19 @@ export class S2TableLayout<T> extends TableLayout<T> {
     // If performing first load or empty, the body will be sticky so we don't want to apply sticky to the loader, otherwise it will
     // affect the positioning of the empty state renderer
     let collection = this.virtualizer!.collection;
-    let isEmptyOrLoading = collection?.size === 0 || (collection.size === 1 && collection.getItem(collection.getFirstKey()!)!.type === 'loader');
+    let isEmptyOrLoading = collection?.size === 0;
     layoutInfo.isSticky = !isEmptyOrLoading;
     return layoutNode;
   }
 
   // y is the height of the headers
   protected buildBody(y: number): LayoutNode {
     let layoutNode = super.buildBody(y);
-    let {children, layoutInfo} = layoutNode;
+    let {layoutInfo} = layoutNode;
     // Needs overflow for sticky loader
     layoutInfo.allowOverflow = true;
     // If loading or empty, we'll want the body to be sticky and centered
-    let isEmptyOrLoading = children?.length === 0 || (children?.length === 1 && children[0].layoutInfo.type === 'loader');
+    let isEmptyOrLoading = this.virtualizer?.collection.size === 0;
     if (isEmptyOrLoading) {
       layoutInfo.rect = new Rect(40, 40, this.virtualizer!.visibleRect.width - 80, this.virtualizer!.visibleRect.height - 80);
       layoutInfo.isSticky = true;
@@ -390,13 +389,13 @@ export const TableBody = /*#__PURE__*/ (forwardRef as forwardRefType)(function T
   // This is because we don't distinguish between loadingMore and loading in the layout, resulting in a different rect being used to build the body. Perhaps can be considered as a user error
   // if they pass loadingMore without having any other items in the table. Alternatively, could update the layout so it knows the current loading state.
   let loadMoreSpinner = (
-    <UNSTABLE_TableLoadingSentinel isLoading={loadingState === 'loadingMore'} onLoadMore={onLoadMore} className={style({height: 'full', width: 'full'})}>
+    <TableLoadMoreItem isLoading={loadingState === 'loadingMore'} onLoadMore={onLoadMore} className={style({height: 'full', width: 'full'})}>
       <div className={centeredWrapper}>
         <ProgressCircle
           isIndeterminate
           aria-label={stringFormatter.format('table.loadingMore')} />
       </div>
-    </UNSTABLE_TableLoadingSentinel>
+    </TableLoadMoreItem>
   );
 
   // If the user is rendering their rows in dynamic fashion, wrap their render function in Collection so we can inject

packages/@react-spectrum/s2/test/Picker.test.tsx

@@ -103,12 +103,20 @@ describe('Picker', () => {
     let options = selectTester.options();
     for (let [index, option] of options.entries()) {
       expect(option).toHaveAttribute('aria-posinset', `${index + 1}`);
+      expect(option).toHaveAttribute('aria-setsize', `${items.length}`);
     }
 
     tree.rerender(<DynamicPicker items={items} loadingState="loadingMore" />);
     options = selectTester.options();
     for (let [index, option] of options.entries()) {
-      expect(option).toHaveAttribute('aria-posinset', `${index + 1}`);
+      if (index === options.length - 1) {
+        // The last row is the loader here which shouldn't have posinset
+        expect(option).not.toHaveAttribute('aria-posinset');
+        expect(option).not.toHaveAttribute('aria-setsize');
+      } else {
+        expect(option).toHaveAttribute('aria-posinset', `${index + 1}`);
+        expect(option).toHaveAttribute('aria-setsize', `${items.length}`);
+      }
     }
 
     let newItems = [...items, {name: 'Chocolate Mint'}, {name: 'Chocolate Chip Cookie Dough'}];
@@ -117,6 +125,7 @@ describe('Picker', () => {
     options = selectTester.options();
     for (let [index, option] of options.entries()) {
       expect(option).toHaveAttribute('aria-posinset', `${index + 1}`);
+      expect(option).toHaveAttribute('aria-setsize', `${newItems.length}`);
     }
   });
 

packages/@react-spectrum/tree/test/TreeView.test.tsx

@@ -1440,8 +1440,8 @@ describe('Tree', () => {
 
       let row = treeTester.rows[0];
       expect(row).toHaveAttribute('aria-level', '1');
-      expect(row).toHaveAttribute('aria-posinset', '1');
-      expect(row).toHaveAttribute('aria-setsize', '1');
+      expect(row).not.toHaveAttribute('aria-posinset');
+      expect(row).not.toHaveAttribute('aria-setsize');
       let gridCell = treeTester.cells({element: row})[0];
       expect(gridCell).toHaveTextContent('No resultsNo results found.');
 

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

@@ -123,7 +123,7 @@ export class GridLayout<T, O extends GridLayoutOptions = GridLayoutOptions> exte
     let collection = this.virtualizer!.collection;
     // Make sure to set rows to 0 if we performing a first time load or are rendering the empty state so that Virtualizer
     // won't try to render its body
-    let isEmptyOrLoading = collection?.size === 0 || (collection.size === 1 && collection.getItem(collection.getFirstKey()!)!.type === 'loader');
+    let isEmptyOrLoading = collection?.size === 0;
     let rows = isEmptyOrLoading ? 0 : Math.ceil(collection.size / numColumns);
     let iterator = collection[Symbol.iterator]();
     let y = rows > 0 ? minSpace.height : 0;

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

@@ -253,41 +253,47 @@ export class ListLayout<T, O extends ListLayoutOptions = ListLayoutOptions> exte
 
   protected buildCollection(y: number = 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;
     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/layout/src/TableLayout.ts

@@ -277,9 +277,6 @@ export class TableLayout<T, O extends TableLayoutProps = TableLayoutProps> exten
       if (y > this.requestedRect.maxY) {
         let rowsAfterRect = collection.size - (children.length + skipped);
         let lastNode = getLastItem(childNodes);
-        if (lastNode?.type === 'loader') {
-          rowsAfterRect--;
-        }
 
         // Estimate the remaining height for rows that we don't need to layout right now.
         y += rowsAfterRect * rowHeight;
@@ -299,7 +296,7 @@ export class TableLayout<T, O extends TableLayoutProps = TableLayoutProps> exten
     }
 
     // Make sure that the table body gets a height if empty or performing initial load
-    let isEmptyOrLoading = collection?.size === 0 || (collection.size === 1 && collection.getItem(collection.getFirstKey()!)!.type === 'loader');
+    let isEmptyOrLoading = collection?.size === 0;
     if (isEmptyOrLoading) {
       y = this.virtualizer!.visibleRect.maxY;
     } else {

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

@@ -170,7 +170,7 @@ export class WaterfallLayout<T extends object, O extends WaterfallLayoutOptions
     }
 
     // Reset all columns to the maximum for the next section. If loading, set to 0 so virtualizer doesn't render its body since there aren't items to render
-    let isEmptyOrLoading = collection?.size === 0 || (collection.size === 1 && collection.getItem(collection.getFirstKey()!)!.type === 'loader');
+    let isEmptyOrLoading = collection?.size === 0;
     let maxHeight = isEmptyOrLoading ? 0 : Math.max(...columnHeights);
     this.contentSize = new Size(this.virtualizer!.visibleRect.width, maxHeight);
     this.layoutInfos = newLayoutInfos;