feat: Add CollectionNode to collection renderer

.gitignore

@@ -1,6 +1,7 @@
 .DS_Store
 .cache
 .idea
+.vscode
 .package-lock.json
 .parcel-cache
 build-storybook.log

packages/@react-aria/collections/src/CollectionBuilder.tsx

@@ -17,6 +17,7 @@ import {createPortal} from 'react-dom';
 import {FocusableContext} from '@react-aria/interactions';
 import {forwardRefType, Key, Node} from '@react-types/shared';
 import {Hidden} from './Hidden';
+import {mergeRefs} from '@react-aria/utils';
 import React, {createContext, ForwardedRef, forwardRef, JSX, ReactElement, ReactNode, useCallback, useContext, useMemo, useRef, useState} from 'react';
 import {useIsSSR} from '@react-aria/ssr';
 import {useSyncExternalStore as useSyncExternalStoreShim} from 'use-sync-external-store/shim/index.js';
@@ -130,7 +131,7 @@ function createCollectionNodeClass(type: string): CollectionNodeClass<any> {
   return NodeClass;
 }
 
-function useSSRCollectionNode<T extends Element>(CollectionNodeClass: CollectionNodeClass<T> | string, props: object, ref: ForwardedRef<T>, rendered?: any, children?: ReactNode, render?: (node: Node<any>) => ReactElement) {
+function useSSRCollectionNode<T extends Element>(CollectionNodeClass: CollectionNodeClass<T> | string, props: object, ref: ForwardedRef<T>, rendered?: any, children?: ReactNode, render?: (node: Node<any>, ref?: ForwardedRef<T>) => ReactElement) {
   // To prevent breaking change, if CollectionNodeClass is a string, create a CollectionNodeClass using the string as the type
   if (typeof CollectionNodeClass === 'string') {
     CollectionNodeClass = createCollectionNodeClass(CollectionNodeClass);
@@ -169,7 +170,7 @@ function useSSRCollectionNode<T extends Element>(CollectionNodeClass: Collection
 export function createLeafComponent<T extends object, P extends object, E extends Element>(CollectionNodeClass: CollectionNodeClass<any> | string, render: (props: P, ref: ForwardedRef<E>) => ReactElement | null): (props: P & React.RefAttributes<E>) => ReactElement | null;
 export function createLeafComponent<T extends object, P extends object, E extends Element>(CollectionNodeClass: CollectionNodeClass<any> | string, render: (props: P, ref: ForwardedRef<E>, node: Node<T>) => ReactElement | null): (props: P & React.RefAttributes<E>) => ReactElement | null;
 export function createLeafComponent<P extends object, E extends Element>(CollectionNodeClass: CollectionNodeClass<any> | string, render: (props: P, ref: ForwardedRef<E>, node?: any) => ReactElement | null): (props: P & React.RefAttributes<any>) => ReactElement | null {
-  let Component = ({node}) => render(node.props, node.props.ref, node);
+  let Component = (forwardRef as forwardRefType)(({node}: {node: Node<any>}, ref: ForwardedRef<E>) => render(node.props, mergeRefs(node.props.ref, ref), node));
   let Result = (forwardRef as forwardRefType)((props: P, ref: ForwardedRef<E>) => {
     let focusableProps = useContext(FocusableContext);
     let isShallow = useContext(ShallowRenderContext);
@@ -186,10 +187,10 @@ export function createLeafComponent<P extends object, E extends Element>(Collect
       ref,
       'children' in props ? props.children : null,
       null,
-      node => (
+      (node, ref) => (
         // Forward FocusableContext to real DOM tree so tooltips work.
         <FocusableContext.Provider value={focusableProps}>
-          <Component node={node} />
+          <Component node={node} ref={ref} />
         </FocusableContext.Provider>
       )
     );
@@ -200,10 +201,10 @@ export function createLeafComponent<P extends object, E extends Element>(Collect
 }
 
 export function createBranchComponent<T extends object, P extends {children?: any}, E extends Element>(CollectionNodeClass: CollectionNodeClass<any> | string, render: (props: P, ref: ForwardedRef<E>, node: Node<T>) => ReactElement | null, useChildren: (props: P) => ReactNode = useCollectionChildren): (props: P & React.RefAttributes<E>) => ReactElement | null {
-  let Component = ({node}) => render(node.props, node.props.ref, node);
+  let Component = (forwardRef as forwardRefType)(({node}: {node: Node<any>}, ref: ForwardedRef<E>) => render(node.props, mergeRefs(node.props.ref, ref), node));
   let Result = (forwardRef as forwardRefType)((props: P, ref: ForwardedRef<E>) => {
     let children = useChildren(props);
-    return useSSRCollectionNode(CollectionNodeClass, props, ref, null, children, node => <Component node={node} />) ?? <></>;
+    return useSSRCollectionNode(CollectionNodeClass, props, ref, null, children, (node, ref) => <Component node={node} ref={ref} />) ?? <></>;
   });
   // @ts-ignore
   Result.displayName = render.name;

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

@@ -329,7 +329,7 @@ export class ElementNode<T> extends BaseNode<T> {
     }
   }
 
-  setProps<E extends Element>(obj: {[key: string]: any}, ref: ForwardedRef<E>, CollectionNodeClass: CollectionNodeClass<any>, rendered?: ReactNode, render?: (node: Node<T>) => ReactElement): void {
+  setProps<E extends Element>(obj: {[key: string]: any}, ref: ForwardedRef<E>, CollectionNodeClass: CollectionNodeClass<any>, rendered?: ReactNode, render?: (node: Node<T>, ref?: ForwardedRef<E>) => ReactElement): void {
     let node;
     let {value, textValue, id, ...props} = obj;
     if (this.node == null) {

packages/@react-aria/collections/test/CollectionBuilder.test.js

@@ -1,33 +1,33 @@
 import {Collection, CollectionBuilder, CollectionNode, createBranchComponent, createLeafComponent} from '../src';
-import React from 'react';
+import React, {createRef} from 'react';
 import {render} from '@testing-library/react';
 
 class ItemNode extends CollectionNode {
   static type = 'item';
 }
 
-const Item = createLeafComponent(ItemNode, () => {
-  return <div />;
+const Item = createLeafComponent(ItemNode, (props, ref) => {
+  return <div {...props} ref={ref} />;
 });
 
-const ItemsOld = createLeafComponent('item', () => {
-  return <div />;
+const ItemsOld = createLeafComponent('item', (props, ref) => {
+  return <div {...props} ref={ref} />;
 });
 
-const SectionOld = createBranchComponent('section', () =>  {
-  return <div />;
+const SectionOld = createBranchComponent('section', (props, ref) =>  {
+  return <div {...props} ref={ref} />;
 });
 
-const renderItems = (items, spyCollection) => (
+const renderItems = (items, spyCollection, children = () => null) => (
   <CollectionBuilder content={<Collection>{items.map((item) => <Item key={item} />)}</Collection>}>
     {collection => {
       spyCollection.current = collection;
-      return null;
+      return children(collection);
     }}
   </CollectionBuilder>
 );
 
-const renderItemsOld = (items, spyCollection) => (
+const renderItemsOld = (items, spyCollection, children = () => null) => (
   <CollectionBuilder
     content={
       <Collection>
@@ -40,7 +40,7 @@ const renderItemsOld = (items, spyCollection) => (
     }>
     {collection => {
       spyCollection.current = collection;
-      return null;
+      return children(collection);
     }}
   </CollectionBuilder>
 );
@@ -70,4 +70,12 @@ describe('CollectionBuilder', () => {
     expect(spyCollection.current.keyMap.get('react-aria-2').firstChildKey).toBe('react-aria-1');
     expect(spyCollection.current.keyMap.get('react-aria-1').type).toBe('item');
   });
+
+  it('should support ref attachment to a rendered node', () => {
+    let spyRef = createRef();
+    render(renderItems(['a'], {}, collection => Array.from(collection).map(
+      item => <React.Fragment key={item.key}>{item.render(item, spyRef)}</React.Fragment>
+    )));
+    expect(spyRef.current).toBeEmptyDOMElement();
+  });
 });

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

@@ -18,6 +18,8 @@ import {MutableRefObject, Ref} from 'react';
 export function mergeRefs<T>(...refs: Array<Ref<T> | MutableRefObject<T> | null | undefined>): Ref<T> {
   if (refs.length === 1 && refs[0]) {
     return refs[0];
+  } else if (refs.every(ref => ref == null)) {
+    return null;
   }
 
   return (value: T | null) => {

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

@@ -29,7 +29,7 @@ import {baseColor, focusRing, size, style} from '../style' with { type: 'macro'
 import ChevronIcon from '../ui-icons/Chevron';
 import {Collection, DOMRef, DOMRefValue, GlobalDOMAttributes, LinkDOMProps, Node} from '@react-types/shared';
 import {controlFont, controlSize, getAllowedOverrides, StyleProps} from './style-utils' with {type: 'macro'};
-import {createContext, forwardRef, Fragment, ReactNode, RefObject, useCallback, useContext, useEffect, useMemo, useRef, useState} from 'react';
+import {createContext, forwardRef, ReactNode, RefObject, useCallback, useContext, useEffect, useMemo, useRef, useState} from 'react';
 // @ts-ignore
 import FolderIcon from '../s2wf-icons/S2_Icon_FolderBreadcrumb_20_N.svg';
 import {forwardRefType} from './types';
@@ -385,6 +385,7 @@ let CollapsingCollectionRenderer: CollectionRenderer = {
 };
 
 let useCollectionRender = (collection: Collection<Node<unknown>>) => {
+  let {CollectionNode} = DefaultCollectionRenderer;
   let {containerRef, onAction} = useContext(CollapseContext) ?? {};
   let [visibleItems, setVisibleItems] = useState(collection.size);
   let {size = 'M'} = useContext(InternalBreadcrumbsContext);
@@ -473,13 +474,13 @@ let useCollectionRender = (collection: Collection<Node<unknown>>) => {
       <HiddenBreadcrumbs items={children} size={size} listRef={listRef} />
       {visibleItems < collection.size && collection.size > 2 ? (
         <>
-          {children[0].render?.(children[0])}
+          <CollectionNode node={children[0]} parent={null} collection={collection} />
           <BreadcrumbMenu items={children.slice(1, sliceIndex)} onAction={onAction} />
-          {children.slice(sliceIndex).map(node => <Fragment key={node.key}>{node.render?.(node)}</Fragment >)}
+          {children.slice(sliceIndex).map(node => <CollectionNode node={node} parent={null} collection={collection} key={node.key} />)}
         </>
       ) : (
         <>
-          {children.map(node => <Fragment key={node.key}>{node.render?.(node)}</Fragment>)}
+          {children.map(node => <CollectionNode node={node} parent={null} collection={collection} key={node.key} />)}
         </>
       )}
     </>

packages/@react-types/shared/src/collections.d.ts

@@ -10,9 +10,9 @@
  * governing permissions and limitations under the License.
  */
 
+import {ForwardedRef, ReactElement, ReactNode} from 'react';
 import {Key} from '@react-types/shared';
 import {LinkDOMProps} from './dom';
-import {ReactElement, ReactNode} from 'react';
 
 export interface ItemProps<T> extends LinkDOMProps {
   /** Rendered contents of the item or child items. */
@@ -223,5 +223,5 @@ export interface Node<T> {
   /** @private */
   shouldInvalidate?: (context: any) => boolean,
   /** A function that renders this node to a React Element in the DOM. */
-  render?: (node: Node<any>) => ReactElement
+  render?: (node: Node<any>, ref?: ForwardedRef<Element>) => ReactElement
 }

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

@@ -108,7 +108,24 @@ export const Section = /*#__PURE__*/ createBranchComponent('section', <T extends
   return render(props, ref, section, 'react-aria-Section');
 });
 
-export interface CollectionBranchProps {
+export interface CollectionNodeProps<T = Node<unknown>> extends HTMLAttributes<HTMLElement> {
+  /** The collection of items to render. */
+  collection: ICollection<Node<unknown>>,
+  /** The node of the item to render. */
+  node: T,
+  /** The parent node of the item to render. */
+  parent: T | null,
+  /** The content that should be rendered before the item. */
+  before?: ReactNode[],
+  /** The content that should be rendered after the item. */
+  after?: ReactNode[],
+  /** A function that renders this node to a React Element in the DOM. */
+  render?: (node: Node<any>, ref?: ForwardedRef<Element>) => ReactElement,
+  /** A ref to the rendered element of this node. */
+  nodeRef?: ForwardedRef<Element>
+}
+
+export interface CollectionBranchProps extends HTMLAttributes<HTMLElement> {
   /** The collection of items to render. */
   collection: ICollection<Node<unknown>>,
   /** The parent node of the items to render. */
@@ -128,7 +145,7 @@ export interface CollectionRootProps extends HTMLAttributes<HTMLElement> {
   renderDropIndicator?: (target: ItemDropTarget) => ReactNode
 }
 
-export interface CollectionRenderer {
+export interface CollectionRenderer<T = Node<unknown>> {
   /** Whether this is a virtualized collection. */
   isVirtualized?: boolean,
   /** A delegate object that provides layout information for items in the collection. */
@@ -138,15 +155,31 @@ export interface CollectionRenderer {
   /** A component that renders the root collection items. */
   CollectionRoot: React.ComponentType<CollectionRootProps>,
    /** A component that renders the child collection items. */
-  CollectionBranch: React.ComponentType<CollectionBranchProps>
+  CollectionBranch: React.ComponentType<CollectionBranchProps>,
+  /** A component that renders the collection item. */
+  CollectionNode?: React.ComponentType<CollectionNodeProps<T>>
+}
+
+interface DefaultRenderer extends CollectionRenderer<Node<unknown>> {
+  /** A component that renders the collection item. */
+  CollectionNode: React.ComponentType<CollectionNodeProps<Node<unknown>>>
 }
 
-export const DefaultCollectionRenderer: CollectionRenderer = {
+export const DefaultCollectionRenderer: DefaultRenderer = {
   CollectionRoot({collection, renderDropIndicator}) {
     return useCollectionRender(collection, null, renderDropIndicator);
   },
   CollectionBranch({collection, parent, renderDropIndicator}) {
     return useCollectionRender(collection, parent, renderDropIndicator);
+  },
+  CollectionNode({before = [], node, render = node.render, nodeRef, after = []}) {
+    return (
+      <>
+        {before.map((el, i) => <React.Fragment key={`${node.key}::before::${i}`}>{el}</React.Fragment>)}
+        {render!(node, nodeRef)}
+        {after.map((el, i) => <React.Fragment key={`${node.key}::after::${i}`}>{el}</React.Fragment>)}
+      </>
+    );
   }
 };
 
@@ -155,28 +188,28 @@ function useCollectionRender(
   parent: Node<unknown> | null,
   renderDropIndicator?: (target: ItemDropTarget) => ReactNode
 ) {
+  let {CollectionNode = DefaultCollectionRenderer.CollectionNode} = useContext(CollectionRendererContext);
+
   return useCachedChildren({
     items: parent ? collection.getChildren!(parent.key) : collection,
-    dependencies: [renderDropIndicator],
+    dependencies: [CollectionNode, parent, 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;
+      let pseudoProps: Pick<CollectionNodeProps<Node<unknown>>, 'before' | 'after'> = {};
+
+      if (renderDropIndicator && node.type === 'item') {
+        let beforeIndicator = renderDropIndicator({type: 'item', key: node.key, dropPosition: 'before'});
+        let afterIndicator = renderAfterDropIndicators(collection, node, renderDropIndicator);
+
+        pseudoProps = {before: new Array(beforeIndicator), after: new Array(afterIndicator)};
       }
 
-      return (
-        <>
-          {renderDropIndicator({type: 'item', key: node.key, dropPosition: 'before'})}
-          {rendered}
-          {renderAfterDropIndicators(collection, node, renderDropIndicator)}
-        </>
-      );
+      return <CollectionNode {...pseudoProps} node={node} parent={parent} collection={collection} key={node.key} />;
     }
   });
 }
@@ -217,7 +250,7 @@ export function renderAfterDropIndicators(collection: ICollection<Node<unknown>>
   return afterIndicators;
 }
 
-export const CollectionRendererContext = createContext<CollectionRenderer>(DefaultCollectionRenderer);
+export const CollectionRendererContext = createContext<CollectionRenderer<any>>(DefaultCollectionRenderer);
 
 type PersistedKeysReturnValue = Set<Key> | null;
 export function usePersistedKeys(focusedKey: Key | null): PersistedKeysReturnValue {

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

@@ -629,6 +629,7 @@ class TreeItemNode extends CollectionNode<any> {
  */
 export const TreeItem = /*#__PURE__*/ createBranchComponent(TreeItemNode, <T extends object>(props: TreeItemProps<T>, ref: ForwardedRef<HTMLDivElement>, item: Node<T>) => {
   let state = useContext(TreeStateContext)!;
+  let {CollectionNode = DefaultCollectionRenderer.CollectionNode} = useContext(CollectionRendererContext);
   ref = useObjectRef<HTMLDivElement>(ref);
   let {dragAndDropHooks, dragState, dropState} = useContext(DragAndDropContext)!;
 
@@ -733,18 +734,18 @@ export const TreeItem = /*#__PURE__*/ createBranchComponent(TreeItemNode, <T ext
 
   let children = useCachedChildren({
     items: state.collection.getChildren!(item.key),
-    children: item => {
-      switch (item.type) {
+    children: node => {
+      switch (node.type) {
         case 'content': {
-          return item.render!(item);
+          return <CollectionNode node={node} parent={item} collection={state.collection} key={item.key} />;
         }
         // Skip item since we don't render the nested rows as children of the parent row, the flattened collection
         // will render them each as siblings instead
         case 'loader':
         case 'item':
           return <></>;
         default:
-          throw new Error('Unsupported element type in TreeRow: ' + item.type);
+          throw new Error('Unsupported element type in TreeRow: ' + node.type);
       }
     }
   });