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,10 @@ export class BaseCollection<T> implements ICollection<Node<T>> {
       throw new Error('Cannot remove a node to a frozen collection');
     }
 
+    if (this.keyMap.get(key)?.type === 'item' && this.keyMap.get(key) != null) {
+      this.itemCount--;
+    }
+
     this.keyMap.delete(key);
   }
 

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

@@ -106,11 +106,24 @@ export function useGridListItem<T>(props: AriaGridListItemOptions, state: ListSt
       if (parent) {
         // siblings must exist because our original node exists, same with lastItem
         let siblings = state.collection.getChildren?.(parent.key)!;
-        setSize = getLastItem(siblings)!.index + 1;
+        let lastItemOfParent = getLastItem(siblings)!;
+        // Don't include loaders as part of the item count, to be revisited when we support focusing the loaders
+        // Also assumes only one loader per tree parent row
+        if (lastItemOfParent.type === 'loader') {
+          setSize = lastItemOfParent.index;
+        } else {
+          setSize = lastItemOfParent.index + 1;
+        }
       }
     } else {
-      setSize = ([...state.collection].filter(row => row.level === 0).at(-1)?.index ?? 0) + 1;
+      let lastItemOfParent = ([...state.collection].filter(row => row.level === 0).at(-1))!;
+      if (lastItemOfParent.type === 'loader') {
+        setSize = lastItemOfParent.index;
+      } else {
+        setSize = lastItemOfParent.index + 1;
+      }
     }
+
     treeGridRowProps = {
       'aria-expanded': isExpanded,
       'aria-level': node.level + 1,

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/TableView.tsx

@@ -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;

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 = 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

@@ -274,9 +274,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;
@@ -296,7 +293,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;

packages/@react-stately/select/src/useSelectState.ts

@@ -15,7 +15,7 @@ import {FormValidationState, useFormValidationState} from '@react-stately/form';
 import {OverlayTriggerState, useOverlayTriggerState} from '@react-stately/overlays';
 import {SelectProps} from '@react-types/select';
 import {SingleSelectListState, useSingleSelectListState} from '@react-stately/list';
-import {useMemo, useState} from 'react';
+import {useState} from 'react';
 
 export interface SelectStateOptions<T> extends Omit<SelectProps<T>, 'children'>, CollectionStateBase<T> {}
 
@@ -62,7 +62,6 @@ export function useSelectState<T extends object>(props: SelectStateOptions<T>):
   });
 
   let [isFocused, setFocused] = useState(false);
-  let isEmpty = useMemo(() => listState.collection.size === 0 || (listState.collection.size === 1 && listState.collection.getItem(listState.collection.getFirstKey()!)?.type === 'loader'), [listState.collection]);
 
   return {
     ...validationState,
@@ -71,13 +70,13 @@ export function useSelectState<T extends object>(props: SelectStateOptions<T>):
     focusStrategy,
     open(focusStrategy: FocusStrategy | null = null) {
       // Don't open if the collection is empty.
-      if (!isEmpty) {
+      if (listState.collection.size !== 0) {
         setFocusStrategy(focusStrategy);
         triggerState.open();
       }
     },
     toggle(focusStrategy: FocusStrategy | null = null) {
-      if (!isEmpty) {
+      if (listState.collection.size !== 0) {
         setFocusStrategy(focusStrategy);
         triggerState.toggle();
       }

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

@@ -194,8 +194,7 @@ function GridListInner<T extends object>({props, collection, gridListRef: ref}:
   }
 
   let {focusProps, isFocused, isFocusVisible} = useFocusRing();
-  // TODO: What do we think about this check? Ideally we could just query the collection and see if ALL node are loaders and thus have it return that it is empty
-  let isEmpty = state.collection.size === 0 || (state.collection.size === 1 && state.collection.getItem(state.collection.getFirstKey()!)?.type === 'loader');
+  let isEmpty = state.collection.size === 0;
   let renderValues = {
     isDropTarget: isRootDropTarget,
     isEmpty,
@@ -529,6 +528,9 @@ export const UNSTABLE_GridListLoadingSentinel = createLeafComponent('loader', fu
     defaultClassName: 'react-aria-GridListLoadingIndicator',
     values: null
   });
+  // For now don't include aria-posinset and aria-setsize on loader since they aren't keyboard focusable
+  // Arguably shouldn't include them ever since it might be confusing to the user to include the loaders as part of the
+  // item count
 
   return (
     <>
@@ -542,7 +544,6 @@ export const UNSTABLE_GridListLoadingSentinel = createLeafComponent('loader', fu
           {...renderProps}
           {...mergeProps(filterDOMProps(props as any))}
           role="row"
-          aria-rowindex={isVirtualized ? item.index + 1 : undefined}
           ref={ref}>
           <div
             aria-colindex={isVirtualized ? 1 : undefined}

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

@@ -202,7 +202,7 @@ function ListBoxInner<T extends object>({state: inputState, props, listBoxRef}:
   }
 
   let {focusProps, isFocused, isFocusVisible} = useFocusRing();
-  let isEmpty = state.collection.size === 0 || (state.collection.size === 1 && state.collection.getItem(state.collection.getFirstKey()!)?.type === 'loader');
+  let isEmpty = state.collection.size === 0;
   let renderValues = {
     isDropTarget: isRootDropTarget,
     isEmpty,
@@ -478,7 +478,6 @@ export interface ListBoxLoadingSentinelProps extends Omit<LoadMoreSentinelProps,
 
 export const UNSTABLE_ListBoxLoadingSentinel = createLeafComponent('loader', function ListBoxLoadingIndicator<T extends object>(props: ListBoxLoadingSentinelProps, ref: ForwardedRef<HTMLDivElement>, item: Node<T>) {
   let state = useContext(ListStateContext)!;
-  let {isVirtualized} = useContext(CollectionRendererContext);
   let {isLoading, onLoadMore, scrollOffset, ...otherProps} = props;
 
   let sentinelRef = useRef<HTMLDivElement>(null);
@@ -500,13 +499,11 @@ export const UNSTABLE_ListBoxLoadingSentinel = createLeafComponent('loader', fun
   let optionProps = {
     // For Android talkback
     tabIndex: -1
+    // For now don't include aria-posinset and aria-setsize on loader since they aren't keyboard focusable
+    // Arguably shouldn't include them ever since it might be confusing to the user to include the loaders as part of the
+    // item count
   };
 
-  if (isVirtualized) {
-    optionProps['aria-posinset'] = item.index + 1;
-    optionProps['aria-setsize'] = state.collection.size;
-  }
-
   return (
     <>
       {/* 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) */}

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

@@ -110,10 +110,6 @@ class TableCollection<T> extends BaseCollection<T> implements ITableCollection<T
     yield this.body;
   }
 
-  get size() {
-    return this.rows.length;
-  }
-
   getFirstKey() {
     return this.body.firstChildKey;
   }
@@ -931,7 +927,7 @@ export const TableBody = /*#__PURE__*/ createBranchComponent('tablebody', <T ext
   let isDroppable = !!dragAndDropHooks?.useDroppableCollectionState && !dropState?.isDisabled;
   let isRootDropTarget = isDroppable && !!dropState && (dropState.isDropTarget({type: 'root'}) ?? false);
 
-  let isEmpty = collection.size === 0 || (collection.rows.length === 1 && collection.rows[0].type === 'loader');
+  let isEmpty = collection.size === 0;
   let renderValues = {
     isDropTarget: isRootDropTarget,
     isEmpty
@@ -954,7 +950,6 @@ export const TableBody = /*#__PURE__*/ createBranchComponent('tablebody', <T ext
     let rowHeaderProps = {};
     let style = {};
     if (isVirtualized) {
-      rowProps['aria-rowindex'] = collection.headerRows.length + 1;
       rowHeaderProps['aria-colspan'] = numColumns;
       style = {display: 'contents'};
     } else {
@@ -1388,7 +1383,9 @@ export const UNSTABLE_TableLoadingSentinel = createLeafComponent('loader', funct
   let style = {};
 
   if (isVirtualized) {
-    rowProps['aria-rowindex'] = item.index + 1 + state.collection.headerRows.length;
+    // For now don't include aria-rowindex on loader since they aren't keyboard focusable
+    // Arguably shouldn't include them ever since it might be confusing to the user to include the loaders as part of the
+    // row count
     rowHeaderProps['aria-colspan'] = numColumns;
     style = {display: 'contents'};
   } else {