Feat: Add useCollectionRef hook to support adjacent child collections

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

@@ -15,20 +15,37 @@ import {BaseNode, Document, ElementNode} from './Document';
 import {CachedChildrenOptions, useCachedChildren} from './useCachedChildren';
 import {createPortal} from 'react-dom';
 import {FocusableContext} from '@react-aria/interactions';
-import {forwardRefType, Node} from '@react-types/shared';
+import {forwardRefType, Node, RefObject} from '@react-types/shared';
 import {Hidden} from './Hidden';
-import React, {createContext, ForwardedRef, forwardRef, JSX, ReactElement, ReactNode, useCallback, useContext, useMemo, useRef, useState} from 'react';
+import React, {createContext, ForwardedRef, forwardRef, JSX, ReactElement, ReactNode, Ref, useCallback, useContext, useMemo, useRef, useState} from 'react';
 import {useIsSSR} from '@react-aria/ssr';
-import {useLayoutEffect} from '@react-aria/utils';
+import {useLayoutEffect, useObjectRef} from '@react-aria/utils';
 import {useSyncExternalStore as useSyncExternalStoreShim} from 'use-sync-external-store/shim/index.js';
 
 const ShallowRenderContext = createContext(false);
 const CollectionDocumentContext = createContext<Document<any, BaseCollection<any>> | null>(null);
 
+export interface CollectionProps<T> extends CachedChildrenOptions<T> {}
+
+export interface CollectionChildren<C extends BaseCollection<object>> {
+  (collection: C): ReactNode
+}
+
+export interface CollectionRenderProps<C extends BaseCollection<object>> {
+  /** A hook that will be called before the collection builder to build the content. */
+  useCollectionContent?: (content: ReactNode) => ReactNode,
+  /** A hook that will be called by the collection builder to render the children. */
+  useCollectionChildren?: (children: CollectionChildren<C>) => CollectionChildren<C>
+  // TODO: Do we also want useCollection() to wrap createCollection()?
+}
+
+interface CollectionRef<C extends BaseCollection<object>, E extends Element> extends RefObject<E | null>, CollectionRenderProps<C> {}
+
 export interface CollectionBuilderProps<C extends BaseCollection<object>> {
   content: ReactNode,
-  children: (collection: C) => ReactNode,
-  createCollection?: () => C
+  children: CollectionChildren<C>,
+  createCollection?: () => C,
+  collectionRef?: CollectionRef<C, Element>
 }
 
 /**
@@ -37,13 +54,14 @@ export interface CollectionBuilderProps<C extends BaseCollection<object>> {
 export function CollectionBuilder<C extends BaseCollection<object>>(props: CollectionBuilderProps<C>): ReactElement {
   // If a document was provided above us, we're already in a hidden tree. Just render the content.
   let doc = useContext(CollectionDocumentContext);
+  let content = props.collectionRef?.useCollectionContent?.(props.content) ?? props.content;
   if (doc) {
     // The React types prior to 18 did not allow returning ReactNode from components
     // even though the actual implementation since React 16 did.
     // We must return ReactElement so that TS does not complain that <CollectionBuilder>
     // is not a valid JSX element with React 16 and 17 types.
     // https://github.com/DefinitelyTyped/DefinitelyTyped/issues/20544
-    return props.content as ReactElement;
+    return content as ReactElement;
   }
 
   // Otherwise, render a hidden copy of the children so that we can build the collection before constructing the state.
@@ -52,14 +70,15 @@ export function CollectionBuilder<C extends BaseCollection<object>>(props: Colle
   // This is fine. CollectionDocumentContext never changes after mounting.
   // eslint-disable-next-line react-hooks/rules-of-hooks
   let {collection, document} = useCollectionDocument(props.createCollection);
+  let children = props.collectionRef?.useCollectionChildren?.(props.children) ?? props.children;
   return (
     <>
       <Hidden>
         <CollectionDocumentContext.Provider value={document}>
-          {props.content}
+          {content}
         </CollectionDocumentContext.Provider>
       </Hidden>
-      <CollectionInner render={props.children} collection={collection} />
+      <CollectionInner render={children} collection={collection} />
     </>
   );
 }
@@ -157,6 +176,12 @@ function useSSRCollectionNode<T extends Element>(Type: string, props: object, re
   return <Type ref={itemRef}>{children}</Type>;
 }
 
+export function useCollectionRef<C extends BaseCollection<object>, E extends Element>(props: CollectionRenderProps<C>, ref: Ref<E>): CollectionRef<C, E> {
+  let refObject = useObjectRef(ref) as CollectionRef<C, E>;
+
+  return Object.assign(refObject, props);
+}
+
 // eslint-disable-next-line @typescript-eslint/no-unused-vars
 export function createLeafComponent<T extends object, P extends object, E extends Element>(type: 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>(type: string, render: (props: P, ref: ForwardedRef<E>, node: Node<T>) => ReactElement | null): (props: P & React.RefAttributes<E>) => ReactElement | null;
@@ -206,8 +231,6 @@ function useCollectionChildren<T extends object>(options: CachedChildrenOptions<
   return useCachedChildren({...options, addIdAndValue: true});
 }
 
-export interface CollectionProps<T> extends CachedChildrenOptions<T> {}
-
 const CollectionContext = createContext<CachedChildrenOptions<unknown> | null>(null);
 
 /** A Collection renders a list of items, automatically managing caching and keys. */

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

@@ -10,10 +10,10 @@
  * governing permissions and limitations under the License.
  */
 
-export {CollectionBuilder, Collection, createLeafComponent, createBranchComponent} from './CollectionBuilder';
+export {CollectionBuilder, Collection, createLeafComponent, createBranchComponent, useCollectionRef} from './CollectionBuilder';
 export {createHideableComponent, useIsHidden} from './Hidden';
 export {useCachedChildren} from './useCachedChildren';
 export {BaseCollection, CollectionNode} from './BaseCollection';
 
-export type {CollectionBuilderProps, CollectionProps} from './CollectionBuilder';
+export type {CollectionBuilderProps, CollectionProps, CollectionRenderProps} from './CollectionBuilder';
 export type {CachedChildrenOptions} from './useCachedChildren';

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

@@ -1,13 +1,13 @@
-import {Collection, CollectionBuilder, createLeafComponent} from '../src';
+import {Collection, CollectionBuilder, createLeafComponent, useCollectionRef} from '../src';
 import React from 'react';
 import {render} from '@testing-library/react';
 
 const Item = createLeafComponent('item', () => {
   return <div />;
 });
 
-const renderItems = (items, spyCollection) => (
-  <CollectionBuilder content={<Collection>{items.map((item) => <Item key={item} />)}</Collection>}>
+const renderItems = (items, spyCollection, collectionRef) => (
+  <CollectionBuilder content={<Collection>{items.map((item) => <Item key={item} />)}</Collection>} collectionRef={collectionRef}>
     {collection => {
       spyCollection.current = collection;
       return null;
@@ -30,4 +30,29 @@ describe('CollectionBuilder', () => {
     expect(spyCollection.current.firstKey).toBe(null);
     expect(spyCollection.current.lastKey).toBe(null);
   });
+
+  it('should support modifying the content via useCollectionChildren', () => {
+    let spyCollection = {};
+    let ref = {current: null};
+    let TestBench = () => {
+      let collectionRef = useCollectionRef({useCollectionContent: () => false}, ref);
+      return renderItems(['a'], spyCollection, collectionRef);
+    };
+    render(<TestBench />);
+    expect(spyCollection.current.frozen).toBe(true);
+    expect(spyCollection.current.firstKey).toBe(null);
+    expect(spyCollection.current.lastKey).toBe(null);
+  });
+
+  it('should support modifying the rendered children via useCollectionChildren', () => {
+    let spyCollection = {};
+    let ref = {current: null};
+    let TestBench = () => {
+      let collectionRef = useCollectionRef({useCollectionChildren: (children) => (c) => <div ref={ref} children={children(c)} />}, ref);
+      return renderItems([], spyCollection, collectionRef);
+    };
+    render(<TestBench />);
+    expect(spyCollection.current.frozen).toBe(true);
+    expect(ref.current).not.toBe(null);
+  });
 });

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

@@ -20,7 +20,7 @@ export function mergeRefs<T>(...refs: Array<Ref<T> | MutableRefObject<T> | null
     return refs[0];
   }
 
-  return (value: T | null) => {
+  let callbackRef = (value: T | null) => {
     let hasCleanup = false;
 
     const cleanups = refs.map(ref => {
@@ -41,6 +41,8 @@ export function mergeRefs<T>(...refs: Array<Ref<T> | MutableRefObject<T> | null
       };
     }
   };
+
+  return Object.assign(callbackRef, ...refs.filter(Boolean));
 }
 
 function setRef<T>(ref: Ref<T> | MutableRefObject<T> | null | undefined, value: T) {

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

@@ -49,6 +49,7 @@ export function useObjectRef<T>(ref?: ((instance: T | null) => (() => void) | vo
 
   return useMemo(
     () => ({
+      ...ref,
       get current() {
         return objRef.current;
       },
@@ -64,6 +65,6 @@ export function useObjectRef<T>(ref?: ((instance: T | null) => (() => void) | vo
         }
       }
     }),
-    [refEffect]
+    [ref, refEffect]
   );
 }

packages/@react-aria/utils/test/mergeRefs.test.tsx

@@ -32,6 +32,18 @@ describe('mergeRefs', () => {
     expect(ref1.current).toBe(ref2.current);
   });
 
+  it('should support additional properties on the refs', () => {
+    // We mock refs here because they are only mutable in React18+
+    let ref1 = {current: null};
+    let ref2 = {current: null, foo: 'bar'};
+    let ref3 = (() => {}) as any;
+    ref3.baz = 'foo';
+
+    let ref = mergeRefs(ref1, ref2, ref3) as any;
+    expect(ref.foo).toBe('bar');
+    expect(ref.baz).toBe('foo');
+  });
+
   if (parseInt(React.version.split('.')[0], 10) >= 19) {
     it('merge Ref Cleanup', () => {
       const cleanUp = jest.fn();

packages/@react-aria/utils/test/useObjectRef.test.js

@@ -63,6 +63,17 @@ describe('useObjectRef', () => {
     expect(ref).toHaveBeenCalledTimes(1);
   });
 
+  it('should support additional properties on the ref', () => {
+    const TextField = React.forwardRef((props, forwardedRef) => {
+      const ref = useObjectRef(forwardedRef);
+      return <input {...props} ref={ref} />;
+    });
+
+    let ref = {current: null, foo: 'bar'};
+    render(<TextField ref={ref} />);
+    expect(ref.foo).toBe('bar');
+  });
+
   /**
    * This describe would completely fail if `useObjectRef` did not account
    * for order of execution and rendering, especially when other components

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

@@ -146,7 +146,8 @@ export const Tabs = forwardRef(function Tabs(props: TabsProps, ref: DOMRef<HTMLD
           'aria-labelledby': props['aria-labelledby']
         }]
       ]}>
-      <CollectionBuilder content={props.children}>
+      {/* @ts-expect-error */}
+      <CollectionBuilder content={props.children} collectionRef={ref}>
         {collection => (
           <CollapsingTabs
             {...props}

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

@@ -92,7 +92,8 @@ export const TagGroup = /*#__PURE__*/ (forwardRef as forwardRefType)(function Ta
   let {onRemove} = props;
   return (
     <InternalTagGroupContext.Provider value={{onRemove}}>
-      <CollectionBuilder content={<Collection {...props} />}>
+      {/* @ts-expect-error */}
+      <CollectionBuilder content={<Collection {...props} />} collectionRef={ref}>
         {collection => <TagGroupInner props={props} forwardedRef={ref} collection={collection} />}
       </CollectionBuilder>
     </InternalTagGroupContext.Provider>

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

@@ -38,7 +38,7 @@ export const Breadcrumbs = /*#__PURE__*/ (forwardRef as forwardRefType)(function
   let DOMProps = filterDOMProps(props, {global: true});
 
   return (
-    <CollectionBuilder content={<Collection {...props} />}>
+    <CollectionBuilder content={<Collection {...props} />} collectionRef={ref}>
       {collection => (
         <ol
           ref={ref}

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

@@ -87,7 +87,7 @@ export const ComboBox = /*#__PURE__*/ (forwardRef as forwardRefType)(function Co
   ), [children, isDisabled, isInvalid, isRequired, props.items, props.defaultItems]);
 
   return (
-    <CollectionBuilder content={content}>
+    <CollectionBuilder content={content} collectionRef={ref}>
       {collection => <ComboBoxInner props={props} collection={collection} comboBoxRef={ref} />}
     </CollectionBuilder>
   );

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

@@ -90,7 +90,7 @@ export const GridList = /*#__PURE__*/ (forwardRef as forwardRefType)(function Gr
   [props, ref] = useContextProps(props, ref, GridListContext);
 
   return (
-    <CollectionBuilder content={<Collection {...props} />}>
+    <CollectionBuilder content={<Collection {...props} />} collectionRef={ref}>
       {collection => <GridListInner props={props} collection={collection} gridListRef={ref} />}
     </CollectionBuilder>
   );

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

@@ -100,7 +100,7 @@ export const ListBox = /*#__PURE__*/ (forwardRef as forwardRefType)(function Lis
   }
 
   return (
-    <CollectionBuilder content={<Collection {...props} />}>
+    <CollectionBuilder content={<Collection {...props} />} collectionRef={ref}>
       {collection => <StandaloneListBox props={props} listBoxRef={ref} collection={collection} />}
     </CollectionBuilder>
   );

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

@@ -169,7 +169,7 @@ export const Menu = /*#__PURE__*/ (forwardRef as forwardRefType)(function Menu<T
 
   // Delay rendering the actual menu until we have the collection so that auto focus works properly.
   return (
-    <CollectionBuilder content={<Collection {...props} />}>
+    <CollectionBuilder content={<Collection {...props} />} collectionRef={ref}>
       {collection => <MenuInner props={props} collection={collection} menuRef={ref} />}
     </CollectionBuilder>
   );

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

@@ -94,7 +94,7 @@ export const Select = /*#__PURE__*/ (forwardRef as forwardRefType)(function Sele
   ), [children, isDisabled, isInvalid, isRequired]);
 
   return (
-    <CollectionBuilder content={content}>
+    <CollectionBuilder content={content} collectionRef={ref}>
       {collection => <SelectInner props={props} collection={collection} selectRef={ref} />}
     </CollectionBuilder>
   );

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

@@ -348,7 +348,7 @@ export const Table = forwardRef(function Table(props: TableProps, ref: Forwarded
   );
 
   return (
-    <CollectionBuilder content={content} createCollection={() => new TableCollection<any>()}>
+    <CollectionBuilder content={content} createCollection={() => new TableCollection<any>()} collectionRef={ref}>
       {collection => <TableInner props={props} forwardedRef={ref as any} selectionState={selectionState} collection={collection} />}
     </CollectionBuilder>
   );

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

@@ -131,7 +131,7 @@ export const Tabs = /*#__PURE__*/ (forwardRef as forwardRefType)(function Tabs(p
   ), [children, orientation]);
 
   return (
-    <CollectionBuilder content={children}>
+    <CollectionBuilder content={children} collectionRef={ref}>
       {collection => <TabsInner props={props} collection={collection} tabsRef={ref} />}
     </CollectionBuilder>
   );

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

@@ -61,7 +61,7 @@ export const TagListContext = createContext<ContextValue<TagListProps<any>, HTML
 export const TagGroup = /*#__PURE__*/ (forwardRef as forwardRefType)(function TagGroup(props: TagGroupProps, ref: ForwardedRef<HTMLDivElement>) {
   [props, ref] = useContextProps(props, ref, TagGroupContext);
   return (
-    <CollectionBuilder content={props.children}>
+    <CollectionBuilder content={props.children} collectionRef={ref}>
       {collection => <TagGroupInner props={props} forwardedRef={ref} collection={collection} />}
     </CollectionBuilder>
   );

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

@@ -164,7 +164,7 @@ export const Tree = /*#__PURE__*/ (forwardRef as forwardRefType)(function Tree<T
   [props, ref] = useContextProps(props, ref, TreeContext);
 
   return (
-    <CollectionBuilder content={<Collection {...props} />}>
+    <CollectionBuilder content={<Collection {...props} />} collectionRef={ref}>
       {collection => <TreeInner props={props} collection={collection} treeRef={ref} />}
     </CollectionBuilder>
   );