Feat: Add support for multiple adjacent collections

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

@@ -15,20 +15,38 @@ 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) => typeof content,
+  /** A hook that will be called by the collection builder to render the children. */
+  useCollectionChildren?: (children: CollectionChildren<C>) => typeof children,
+  /** A hook that will be called by the collection builder to retrieve the collection and document. */
+  useCollectionDocument?: (state: CollectionDocumentResult<any, BaseCollection<any>>) => typeof state
+}
+
+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,29 +55,33 @@ 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 ref = props.collectionRef ?? {} as CollectionRef<C, Element>;
+  let content = ref.useCollectionContent ? ref.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.
   // This should always come before the real DOM content so we have built the collection by the time it renders during SSR.
 
   // This is fine. CollectionDocumentContext never changes after mounting.
   // eslint-disable-next-line react-hooks/rules-of-hooks
-  let {collection, document} = useCollectionDocument(props.createCollection);
+  let state = useCollectionDocument(props.createCollection);
+  let {collection, document} = ref.useCollectionDocument ? ref.useCollectionDocument(state) : state;
+  let children = ref.useCollectionChildren ? ref.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} />
     </>
   );
 }
@@ -94,7 +116,7 @@ const useSyncExternalStore = typeof React['useSyncExternalStore'] === 'function'
   ? React['useSyncExternalStore']
   : useSyncExternalStoreFallback;
 
-function useCollectionDocument<T extends object, C extends BaseCollection<T>>(createCollection?: () => C): CollectionDocumentResult<T, C> {
+export function useCollectionDocument<T extends object, C extends BaseCollection<T>>(createCollection?: () => C): CollectionDocumentResult<T, C> {
   // The document instance is mutable, and should never change between renders.
   // useSyncExternalStore is used to subscribe to updates, which vends immutable Collection objects.
   let [document] = useState(() => new Document<T, C>(createCollection?.() || new BaseCollection() as C));
@@ -157,6 +179,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 +234,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 {createHideableComponent, useIsHidden} from './Hidden';
+export {CollectionBuilder, Collection, createLeafComponent, createBranchComponent, useCollectionRef, useCollectionDocument} from './CollectionBuilder';
+export {createHideableComponent, useIsHidden, Hidden} 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,14 @@
-import {Collection, CollectionBuilder, createLeafComponent} from '../src';
+import {Collection, CollectionBuilder, createLeafComponent, useCollectionDocument, useCollectionRef} from '../src';
+import {mergeRefs, useObjectRef} from '@react-aria/utils';
 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 id={item} key={item} textValue={item} />)}</Collection>} collectionRef={collectionRef}>
     {collection => {
       spyCollection.current = collection;
       return null;
@@ -30,4 +31,187 @@ describe('CollectionBuilder', () => {
     expect(spyCollection.current.firstKey).toBe(null);
     expect(spyCollection.current.lastKey).toBe(null);
   });
+
+  it('should support modifying the content via useCollectionContent', () => {
+    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);
+  });
+
+  describe('synchronization', () => {
+    let CollectionAContext = React.createContext(null);
+    let CollectionAStateContext = React.createContext(null);
+    let CollectionBContext = React.createContext(null);
+    let CollectionBStateContext = React.createContext(null);
+
+    let CollectionA = React.forwardRef(({items}, ref) => {
+      let ctx = React.useContext(CollectionAContext);
+      let state = React.useContext(CollectionAStateContext);
+      let mergedRef = useObjectRef(React.useMemo(() => mergeRefs(ref, ctx), [ref, ctx]));
+
+      if (state) {
+        return (
+          <div data-testid="collection-a">
+            {Array.from(state).map(item => <div data-testid={item.key} key={item.key} />)}
+          </div>
+        );
+      }
+
+      return renderItems(items, {}, mergedRef);
+    });
+
+    let CollectionB = React.forwardRef(({items}, ref) => {
+      let ctx = React.useContext(CollectionBContext);
+      let state = React.useContext(CollectionBStateContext);
+      let mergedRef = useObjectRef(React.useMemo(() => mergeRefs(ref, ctx), [ref, ctx]));
+
+      if (state) {
+        return (
+          <div data-testid="collection-b">
+            {Array.from(state).map(item => <div data-testid={item.key} key={item.key} />)}
+          </div>
+        );
+      }
+
+      return renderItems(items, {}, mergedRef);
+    });
+
+    let Synchronized = ({collectionA, collectionB, children, filterFn}) => {
+      let synchronized = React.useMemo(() => filterFn(collectionA, collectionB), [filterFn, collectionA, collectionB]);
+      return (
+        <CollectionAStateContext.Provider value={synchronized.collectionA}>
+          <CollectionBStateContext.Provider value={synchronized.collectionB}>
+            {children}
+          </CollectionBStateContext.Provider>
+        </CollectionAStateContext.Provider>
+      );
+    };
+
+    let Synchronizer = ({spy: Spy, children}) => {
+      let useCollectionContent = React.useCallback(() => null, []);
+      let useCollectionChildren = React.useCallback(() => () => null, []);
+
+      let refA = useCollectionRef({useCollectionContent, useCollectionChildren}, React.useRef(null));
+      let refB = useCollectionRef({useCollectionContent, useCollectionChildren}, React.useRef(null));
+
+      let contentA = <CollectionBContext.Provider value={refB}>{children}</CollectionBContext.Provider>;
+      let contentB = <CollectionAContext.Provider value={refA}>{children}</CollectionAContext.Provider>;
+
+      // One way because changes in the DocumentB will not trigger a rerender of DocumentA
+      let filterFn = React.useCallback((cA, cB) => {
+        let keysA = new Set(cA.getKeys());
+        return {collectionA: cA, collectionB: cB.UNSTABLE_filter(item => keysA.has(item))};
+      }, []);
+
+      return (
+        <CollectionBuilder content={contentA}>
+          {(c1) => (<CollectionBuilder content={contentB}>
+            {(c2) => (<>
+              <Synchronized collectionA={c1} collectionB={c2} children={children} filterFn={filterFn} />
+              <Spy collectionA={c1} collectionB={c2} />
+            </>)}
+          </CollectionBuilder>)}
+        </CollectionBuilder>
+      );
+    };
+
+    let TwoWaySynchronizer = ({spy: Spy, children}) => {
+      let useCollectionContent = React.useCallback(() => null, []);
+      let useCollectionChildren = React.useCallback(() => () => null, []);
+
+      let stateA = useCollectionDocument();
+      let stateB = useCollectionDocument();
+
+      let useCollectionDocumentA = React.useCallback(() => stateA, [stateA]);
+      let useCollectionDocumentB = React.useCallback(() => stateB, [stateB]);
+
+      let refA = useCollectionRef({useCollectionContent, useCollectionChildren}, React.useRef(null));
+      let contentA = <CollectionBContext.Provider value={refA}>{children}</CollectionBContext.Provider>;
+
+      let refB = useCollectionRef({useCollectionContent, useCollectionChildren}, React.useRef(null));
+      let contentB = <CollectionAContext.Provider value={refB}>{children}</CollectionAContext.Provider>;
+
+      let refA2 = useCollectionRef({useCollectionDocument: useCollectionDocumentA}, React.useRef(null));
+      let refB2 = useCollectionRef({useCollectionDocument: useCollectionDocumentB}, React.useRef(null));
+
+      let filterFn = React.useCallback((cA, cB) => {
+        let keysA = new Set(cA.getKeys()), keysB = new Set(cB.getKeys());
+        let collectionA = cA.UNSTABLE_filter(item => keysB.has(item));
+        let collectionB = cB.UNSTABLE_filter(item => keysA.has(item));
+        return {collectionA, collectionB};
+      }, []);
+
+      return (
+        <>
+          <CollectionBuilder content={contentA} collectionRef={refA2} children={useCollectionChildren()} />
+          <CollectionBuilder content={contentB} collectionRef={refB2} children={useCollectionChildren()} />
+          <Synchronized collectionA={stateA.collection} collectionB={stateB.collection} children={children} filterFn={filterFn} />
+          <Spy collectionA={stateA.collection} collectionB={stateB.collection} />
+        </>
+      );
+    };
+
+    it('should support one-way synchronization of multiple collections', () => {
+      let Spy = jest.fn(() => null);
+
+      // Synchronizer will force CollectionB to only be rendered with keys of CollectionA
+      let {queryAllByTestId} = render(
+        <Synchronizer spy={Spy}>
+          <CollectionA items={['a', 'b']} />
+          <CollectionB items={['a', 'c']} />
+        </Synchronizer>
+      );
+
+      expect(Spy).toHaveBeenCalledTimes(2);
+      expect(Spy.mock.calls[1][0].collectionA.getFirstKey()).toBe('a');
+      expect(Spy.mock.calls[1][0].collectionA.getLastKey()).toBe('b');
+      expect(Spy.mock.calls[1][0].collectionB.getFirstKey()).toBe('a');
+      expect(Spy.mock.calls[1][0].collectionB.getLastKey()).toBe('c');
+
+      expect(queryAllByTestId('a')).toHaveLength(2);
+      expect(queryAllByTestId('b')).toHaveLength(1);
+      expect(queryAllByTestId('c')).toHaveLength(0);
+    });
+
+    it('should support two-way synchronization of multiple collections', () => {
+      let Spy = jest.fn(() => null);
+
+      // TwoWaySynchronizer will force both collections to only be rendered with mutually shared keys
+      let {queryAllByTestId} = render(
+        <TwoWaySynchronizer spy={Spy}>
+          <CollectionA items={['a', 'b']} />
+          <CollectionB items={['a', 'c']} />
+        </TwoWaySynchronizer>
+      );
+
+      expect(Spy).toHaveBeenCalledTimes(2);
+      expect(Spy.mock.calls[1][0].collectionA.getFirstKey()).toBe('a');
+      expect(Spy.mock.calls[1][0].collectionA.getLastKey()).toBe('b');
+      expect(Spy.mock.calls[1][0].collectionB.getFirstKey()).toBe('a');
+      expect(Spy.mock.calls[1][0].collectionB.getLastKey()).toBe('c');
+
+      expect(queryAllByTestId('a')).toHaveLength(2);
+      expect(queryAllByTestId('b')).toHaveLength(0);
+      expect(queryAllByTestId('c')).toHaveLength(0);
+    });
+  });
 });

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();