feat: add support for tree sections

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

@@ -45,11 +45,11 @@ export function useCachedChildren<T extends object>(props: CachedChildrenOptions
           rendered = children(item);
           // @ts-ignore
           let key = rendered.props.id ?? item.key ?? item.id;
-           
+
           if (key == null) {
             throw new Error('Could not determine key for item');
           }
-           
+
           if (idScope != null) {
             key = idScope + ':' + key;
           }

packages/@react-aria/dnd/src/DropTargetKeyboardNavigation.ts

@@ -59,6 +59,10 @@ function nextDropTarget(
       nextKey = keyboardDelegate.getKeyBelow?.(target.key);
     }
     let nextCollectionKey = collection.getKeyAfter(target.key);
+    let nextCollectionNode = nextCollectionKey && collection.getItem(nextCollectionKey);
+    if (nextCollectionNode && nextCollectionNode.type === 'content') {
+      nextCollectionKey = nextCollectionKey ? collection.getKeyAfter(nextCollectionKey) : null;
+    }
 
     // If the keyboard delegate did not move to the next key in the collection,
     // jump to that key with the same drop position. Otherwise, try the other

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

@@ -11,7 +11,7 @@
  */
 
 import {chain, getActiveElement, getEventTarget, getScrollParent, isFocusWithin, mergeProps, nodeContains, scrollIntoViewport, useSlotId, useSyntheticLinkProps} from '@react-aria/utils';
-import {DOMAttributes, FocusableElement, Key, RefObject, Node as RSNode} from '@react-types/shared';
+import {Collection, DOMAttributes, FocusableElement, Key, RefObject, Node as RSNode} from '@react-types/shared';
 import {focusSafely, getFocusableTreeWalker} from '@react-aria/focus';
 import {getRowId, listMap} from './utils';
 import {HTMLAttributes, KeyboardEvent as ReactKeyboardEvent, useRef} from 'react';
@@ -100,11 +100,11 @@ export function useGridListItem<T>(props: AriaGridListItemOptions, state: ListSt
 
     let isExpanded = hasChildRows ? state.expandedKeys.has(node.key) : undefined;
     let setSize = 1;
-    if (node.level > 0 && node?.parentKey != null) {
+    if (node.level >= 0 && node?.parentKey != null) {
       let parent = state.collection.getItem(node.parentKey);
       if (parent) {
         // siblings must exist because our original node exists
-        let siblings = state.collection.getChildren?.(parent.key)!;
+        let siblings = getDirectChildren(parent, state.collection);
         setSize = [...siblings].filter(row => row.type === 'item').length;
       }
     } else {
@@ -326,3 +326,16 @@ function last(walker: TreeWalker) {
   } while (last);
   return next;
 }
+
+function getDirectChildren<T>(parent: RSNode<T>, collection: Collection<RSNode<T>>) {
+  // We can't assume that we can use firstChildKey because if a person builds a tree using hooks, they would not have access to that property (using type Node vs CollectionNode)
+  // Instead, get all children and start at the first node (rather than just using firstChildKey) and only look at its siblings
+  let children = collection.getChildren?.(parent.key);
+  let node = children && Array.from(children).length > 0 ?  Array.from(children)[0] : null;
+  let siblings: RSNode<T>[] = [];
+  while (node) {
+    siblings.push(node);
+    node = node.nextKey != null ? collection.getItem(node.nextKey) : null;
+  }
+  return siblings;
+}

packages/@react-aria/selection/src/ListKeyboardDelegate.ts

@@ -23,7 +23,8 @@ interface ListKeyboardDelegateOptions<T> {
   direction?: Direction,
   disabledKeys?: Set<Key>,
   disabledBehavior?: DisabledBehavior,
-  layoutDelegate?: LayoutDelegate
+  layoutDelegate?: LayoutDelegate,
+  expandedKeys?: Set<Key>
 }
 
 export class ListKeyboardDelegate<T> implements KeyboardDelegate {
@@ -36,8 +37,9 @@ export class ListKeyboardDelegate<T> implements KeyboardDelegate {
   private orientation?: Orientation;
   private direction?: Direction;
   private layoutDelegate: LayoutDelegate;
+  private expandedKeys?: Set<Key>;
 
-  constructor(collection: Collection<Node<T>>, disabledKeys: Set<Key>, ref: RefObject<HTMLElement | null>, collator?: Intl.Collator);
+  constructor(collection: Collection<Node<T>>, disabledKeys: Set<Key>, ref: RefObject<HTMLElement | null>, collator?: Intl.Collator, expandedKeys?: Set<Key>);
   constructor(options: ListKeyboardDelegateOptions<T>);
   constructor(...args: any[]) {
     if (args.length === 1) {
@@ -51,6 +53,7 @@ export class ListKeyboardDelegate<T> implements KeyboardDelegate {
       this.direction = opts.direction;
       this.layout = opts.layout || 'stack';
       this.layoutDelegate = opts.layoutDelegate || new DOMLayoutDelegate(opts.ref);
+      this.expandedKeys = opts.expandedKeys;
     } else {
       this.collection = args[0];
       this.disabledKeys = args[1];
@@ -60,6 +63,7 @@ export class ListKeyboardDelegate<T> implements KeyboardDelegate {
       this.orientation = 'vertical';
       this.disabledBehavior = 'all';
       this.layoutDelegate = new DOMLayoutDelegate(this.ref);
+      this.expandedKeys = args[4];
     }
 
     // If this is a vertical stack, remove the left/right methods completely
@@ -88,6 +92,48 @@ export class ListKeyboardDelegate<T> implements KeyboardDelegate {
     return null;
   }
 
+  private findNextVisible(key: Key | null): Key | null {
+    let node = key ? this.collection.getItem(key) : null;
+    if (!node) {
+      return null;
+    }
+
+    // If the node's parent is expanded, then we can assume that this is a visible node
+    if (node.parentKey && this.expandedKeys?.has(node.parentKey)) {
+      return node.key;
+    }
+
+    // If the node's parent is not expanded, find the top-most non-expanded node since it's possible for them to be nested
+    let parentNode = node.parentKey ? this.collection.getItem(node.parentKey) : null;
+    // if the the parent node is not a section, and the parent node is not included in expanded keys
+    while (parentNode && parentNode.type !== 'section' && node && node.parentKey && this.expandedKeys && !this.expandedKeys.has(parentNode.key)) {
+      node = this.collection.getItem(node.parentKey);
+      parentNode = node && node.parentKey ? this.collection.getItem(node.parentKey) : null;
+    }
+
+    return node?.key ?? null;
+  }
+
+  private findNextNonDisabledVisible(key: Key | null, getNext: (key: Key) => Key | null) {
+    let nextKey = key;
+    while (nextKey !== null) {
+      let visibleKey = this.findNextVisible(nextKey);
+      // If visibleKey is null, that means there are no visibleKeys (don't feel like this is a real use case though, I would assume that there is always one visible node)
+      if (visibleKey == null) {
+        return null;
+      }
+
+      let node = this.collection.getItem(visibleKey);
+      if (node?.type === 'item' && !this.isDisabled(node)) {
+        return visibleKey;
+      }
+
+      nextKey = getNext(visibleKey);
+    }
+
+    return null;
+  }
+
   getNextKey(key: Key): Key | null {
     let nextKey: Key | null = key;
     nextKey = this.collection.getKeyAfter(nextKey);
@@ -201,6 +247,10 @@ export class ListKeyboardDelegate<T> implements KeyboardDelegate {
 
   getLastKey(): Key | null {
     let key = this.collection.getLastKey();
+    // we only need to check for visible keys if items can be expanded/collapsed
+    if (this.expandedKeys) {
+      return this.findNextNonDisabledVisible(key, key => this.collection.getKeyBefore(key));
+    }
     return this.findNextNonDisabled(key, key => this.collection.getKeyBefore(key));
   }
 

packages/@react-stately/data/src/useTreeData.ts

@@ -504,7 +504,7 @@ function moveItems<T extends object>(
     // decrement the index if the child being removed is in the target parent and before the target index
     // the root node is special, it is null, and will not have a key, however, a parentKey can still point to it
     if ((child.parentKey === toParent
-       || child.parentKey === toParent?.key)
+      || child.parentKey === toParent?.key)
       && keyArray.includes(child.key)
       && (toParent?.children ? toParent.children.indexOf(child) : items.indexOf(child)) < originalToIndex) {
       toIndex--;

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

@@ -253,7 +253,9 @@ export class ListLayout<T, O extends ListLayoutOptions = ListLayoutOptions> exte
 
   protected buildCollection(y: number = this.padding): LayoutNode[] {
     let collection = this.virtualizer!.collection;
-    let collectionNodes = [...collection];
+    // filter out content nodes since we don't want them to affect the height
+    // Tree specific for now, if we add content nodes to other collection items, we might need to reconsider this
+    let collectionNodes = [...collection].filter((node) => node.type !== 'content');
     let loaderNodes = collectionNodes.filter(node => node.type === 'loader');
     let nodes: LayoutNode[] = [];
     let isEmptyOrLoading = collection?.size === 0;
@@ -370,6 +372,11 @@ export class ListLayout<T, O extends ListLayoutOptions = ListLayoutOptions> exte
     let skipped = 0;
     let children: LayoutNode[] = [];
     for (let child of getChildNodes(node, collection)) {
+      // skip if it is a content node, Tree specific for now, if we add content nodes to other collection items, we might need to reconsider this
+      if (child.type === 'content') {
+        continue;
+      }
+
       let rowHeight = (this.rowHeight ?? this.estimatedRowHeight ?? DEFAULT_HEIGHT) + this.gap;
 
       // Skip rows before the valid rectangle unless they are already cached.

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

@@ -159,6 +159,12 @@ function useCollectionRender(
     items: parent ? collection.getChildren!(parent.key) : collection,
     dependencies: [renderDropIndicator],
     children(node) {
+      // Return a empty fragment since we don't want to render the content twice
+      // If we don't skip the content node here, we end up rendering them twice in a Tree since we also render the content node in TreeItem
+      if (node.type === 'content') {
+        return <></>;
+      }
+
       let rendered = node.render!(node);
       if (!renderDropIndicator || node.type !== 'item') {
         return rendered;

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

@@ -653,7 +653,7 @@ export const GridListSection = /*#__PURE__*/ createBranchComponent(SectionNode,
 export interface GridListHeaderProps extends DOMRenderProps<'div', undefined>, DOMProps, GlobalDOMAttributes<HTMLElement> {}
 
 export const GridListHeaderContext = createContext<ContextValue<GridListHeaderProps, HTMLDivElement>>({});
-const GridListHeaderInnerContext = createContext<HTMLAttributes<HTMLElement> | null>(null);
+export const GridListHeaderInnerContext = createContext<HTMLAttributes<HTMLElement> | null>(null);
 
 export const GridListHeader = /*#__PURE__*/ createLeafComponent(HeaderNode, function Header(props: GridListHeaderProps, ref: ForwardedRef<HTMLDivElement>) {
   [props, ref] = useContextProps(props, ref, GridListHeaderContext);