Refactor collection internals in preparation for public API (#6608)

Author: devongovett

packages/@react-stately/grid/src/useGridState.ts

@@ -1,7 +1,7 @@
 import {getChildNodes, getFirstItem, getLastItem} from '@react-stately/collections';
 import {GridCollection, GridNode} from '@react-types/grid';
 import {Key} from '@react-types/shared';
-import {MultipleSelectionStateProps, SelectionManager, useMultipleSelectionState} from '@react-stately/selection';
+import {MultipleSelectionState, MultipleSelectionStateProps, SelectionManager, useMultipleSelectionState} from '@react-stately/selection';
 import {useEffect, useMemo, useRef} from 'react';
 
 export interface GridState<T, C extends GridCollection<T>> {
@@ -17,15 +17,18 @@ export interface GridState<T, C extends GridCollection<T>> {
 export interface GridStateOptions<T, C extends GridCollection<T>> extends MultipleSelectionStateProps {
   collection: C,
   disabledKeys?: Iterable<Key>,
-  focusMode?: 'row' | 'cell'
+  focusMode?: 'row' | 'cell',
+  /** @private - do not use unless you know what you're doing. */
+  UNSAFE_selectionState?: MultipleSelectionState
 }
 
 /**
  * Provides state management for a grid component. Handles row selection and focusing a grid cell's focusable child if applicable.
  */
 export function useGridState<T extends object, C extends GridCollection<T>>(props: GridStateOptions<T, C>): GridState<T, C> {
   let {collection, focusMode} = props;
-  let selectionState = useMultipleSelectionState(props);
+  // eslint-disable-next-line react-hooks/rules-of-hooks
+  let selectionState = props.UNSAFE_selectionState || useMultipleSelectionState(props);
   let disabledKeys = useMemo(() =>
       props.disabledKeys ? new Set(props.disabledKeys) : new Set<Key>()
     , [props.disabledKeys]);

packages/@react-stately/table/src/useTableState.ts

@@ -13,7 +13,7 @@
 import {GridState, useGridState} from '@react-stately/grid';
 import {TableCollection as ITableCollection, TableBodyProps, TableHeaderProps} from '@react-types/table';
 import {Key, Node, SelectionMode, Sortable, SortDescriptor, SortDirection} from '@react-types/shared';
-import {MultipleSelectionStateProps} from '@react-stately/selection';
+import {MultipleSelectionState, MultipleSelectionStateProps} from '@react-stately/selection';
 import {ReactElement, useCallback, useMemo, useState} from 'react';
 import {TableCollection} from './TableCollection';
 import {useCollection} from '@react-stately/collections';
@@ -52,7 +52,9 @@ export interface TableStateProps<T> extends MultipleSelectionStateProps, Sortabl
   /** Whether the row drag button should be displayed.
    * @private
    */
-  showDragButtons?: boolean
+  showDragButtons?: boolean,
+  /** @private - do not use unless you know what you're doing. */
+  UNSAFE_selectionState?: MultipleSelectionState
 }
 
 const OPPOSITE_SORT_DIRECTION = {

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

@@ -10,13 +10,13 @@
  * governing permissions and limitations under the License.
  */
 import {AriaBreadcrumbsProps} from 'react-aria';
-import {Collection, Node} from 'react-stately';
-import {CollectionProps, CollectionRendererContext, createLeafComponent, useCollection} from './Collection';
+import {Collection, CollectionBuilder, CollectionProps, CollectionRendererContext, createLeafComponent} from './Collection';
 import {ContextValue, forwardRefType, RenderProps, SlotProps, StyleProps, useContextProps, useRenderProps, useSlottedContext} from './utils';
 import {filterDOMProps} from '@react-aria/utils';
 import {Key} from '@react-types/shared';
 import {LinkContext} from './Link';
-import React, {createContext, ForwardedRef, forwardRef, ReactNode, RefObject, useContext} from 'react';
+import {Node} from 'react-stately';
+import React, {createContext, ForwardedRef, forwardRef, ReactNode, useContext} from 'react';
 
 export interface BreadcrumbsProps<T> extends Omit<CollectionProps<T>, 'disabledKeys'>, AriaBreadcrumbsProps, StyleProps, SlotProps {
   /** Whether the breadcrumbs are disabled. */
@@ -29,36 +29,23 @@ export const BreadcrumbsContext = createContext<ContextValue<BreadcrumbsProps<an
 
 function Breadcrumbs<T extends object>(props: BreadcrumbsProps<T>, ref: ForwardedRef<HTMLOListElement>) {
   [props, ref] = useContextProps(props, ref, BreadcrumbsContext);
-  let {portal, collection} = useCollection(props);
-
-  // Render the portal first so that we have the collection by the time we render the DOM in SSR
-  return (
-    <>
-      {portal}
-      <BreadcrumbsInner props={props} collection={collection} breadcrumbsRef={ref} />
-    </>
-  );
-}
-
-interface BreadcrumbsInnerProps<T> {
-  props: BreadcrumbsProps<T>,
-  collection: Collection<Node<T>>,
-  breadcrumbsRef: RefObject<HTMLOListElement | null>
-}
-
-function BreadcrumbsInner<T extends object>({props, collection, breadcrumbsRef: ref}: BreadcrumbsInnerProps<T>) {
   let {CollectionRoot} = useContext(CollectionRendererContext);
+
   return (
-    <ol
-      ref={ref}
-      {...filterDOMProps(props, {labelable: true})}
-      slot={props.slot || undefined}
-      style={props.style}
-      className={props.className ?? 'react-aria-Breadcrumbs'}>
-      <BreadcrumbsContext.Provider value={props}>
-        <CollectionRoot collection={collection} />
-      </BreadcrumbsContext.Provider>
-    </ol>
+    <CollectionBuilder content={<Collection {...props} />}>
+      {collection => (
+        <ol
+          ref={ref}
+          {...filterDOMProps(props, {labelable: true})}
+          slot={props.slot || undefined}
+          style={props.style}
+          className={props.className ?? 'react-aria-Breadcrumbs'}>
+          <BreadcrumbsContext.Provider value={props}>
+            <CollectionRoot collection={collection} />
+          </BreadcrumbsContext.Provider>
+        </ol>
+      )}
+    </CollectionBuilder>
   );
 }
 

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

@@ -11,10 +11,10 @@
  */
 import {CollectionBase, DropTargetDelegate, Key, LayoutDelegate} from '@react-types/shared';
 import {createPortal} from 'react-dom';
-import {forwardRefType, StyleProps} from './utils';
+import {forwardRefType, Hidden, StyleProps} from './utils';
 import {Collection as ICollection, Node, SelectionBehavior, SelectionMode, SectionProps as SharedSectionProps} from 'react-stately';
 import {mergeProps, useIsSSR} from 'react-aria';
-import React, {cloneElement, createContext, ForwardedRef, forwardRef, HTMLAttributes, JSX, ReactElement, ReactNode, RefObject, useCallback, useContext, useMemo, useRef} from 'react';
+import React, {cloneElement, createContext, ForwardedRef, forwardRef, HTMLAttributes, JSX, ReactElement, ReactNode, RefObject, useCallback, useContext, useMemo, useRef, useState} from 'react';
 import {useLayoutEffect} from '@react-aria/utils';
 import {useSyncExternalStore as useSyncExternalStoreShim} from 'use-sync-external-store/shim/index.js';
 
@@ -274,7 +274,7 @@ class BaseNode<T> {
  * A mutable element node in the fake DOM tree. It owns an immutable
  * Collection Node which is copied on write.
  */
-export class ElementNode<T> extends BaseNode<T> {
+class ElementNode<T> extends BaseNode<T> {
   nodeType = 8; // COMMENT_NODE (we'd use ELEMENT_NODE but React DevTools will fail to get its dimensions)
   node: NodeValue<T>;
   private _index: number = 0;
@@ -503,7 +503,7 @@ export class BaseCollection<T> implements ICollection<Node<T>> {
  * A mutable Document in the fake DOM. It owns an immutable Collection instance,
  * which is lazily copied on write during updates.
  */
-export class Document<T, C extends BaseCollection<T> = BaseCollection<T>> extends BaseNode<T> {
+class Document<T, C extends BaseCollection<T> = BaseCollection<T>> extends BaseNode<T> {
   nodeType = 11; // DOCUMENT_FRAGMENT_NODE
   ownerDocument = this;
   dirtyNodes: Set<BaseNode<T>> = new Set();
@@ -709,16 +709,41 @@ export function useCollectionChildren<T extends object>(props: CachedChildrenOpt
 }
 
 const ShallowRenderContext = createContext(false);
+const CollectionDocumentContext = createContext<Document<any, BaseCollection<any>> | null>(null);
 
-interface CollectionResult<C> {
-  portal: ReactNode,
-  collection: C
+export interface CollectionBuilderProps<C extends BaseCollection<object>> {
+  content: ReactNode,
+  children: (collection: C) => ReactNode,
+  createCollection?: () => C
 }
 
-export function useCollection<T extends object, C extends BaseCollection<T>>(props: CollectionProps<T>, initialCollection?: C): CollectionResult<C> {
-  let {collection, document} = useCollectionDocument<T, C>(initialCollection);
-  let portal = useCollectionPortal<T, C>(props, document);
-  return {portal, collection};
+export function CollectionBuilder<C extends BaseCollection<object>>(props: CollectionBuilderProps<C>) {
+  // If a document was provided above us, we're already in a hidden tree. Just render the content.
+  let doc = useContext(CollectionDocumentContext);
+  if (doc) {
+    return props.content;
+  }
+
+  // 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);
+  return (
+    <>
+      <Hidden>
+        <CollectionDocumentContext.Provider value={document}>
+          {props.content}
+        </CollectionDocumentContext.Provider>
+      </Hidden>
+      <CollectionInner render={props.children} collection={collection} />
+    </>
+  );
+}
+
+function CollectionInner({collection, render}) {
+  return render(collection);
 }
 
 interface CollectionDocumentResult<T, C extends BaseCollection<T>> {
@@ -747,10 +772,10 @@ const useSyncExternalStore = typeof React['useSyncExternalStore'] === 'function'
   ? React['useSyncExternalStore']
   : useSyncExternalStoreFallback;
 
-export function useCollectionDocument<T extends object, C extends BaseCollection<T>>(initialCollection?: C): CollectionDocumentResult<T, C> {
+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 = useMemo(() => new Document<T, C>(initialCollection || new BaseCollection() as C), [initialCollection]);
+  let [document] = useState(() => new Document<T, C>(createCollection?.() || new BaseCollection() as C));
   let subscribe = useCallback((fn: () => void) => document.subscribe(fn), [document]);
   let getSnapshot = useCallback(() => {
     let collection = document.getCollection();
@@ -779,27 +804,6 @@ export function useCollectionDocument<T extends object, C extends BaseCollection
 }
 
 const SSRContext = createContext<BaseNode<any> | null>(null);
-export const CollectionDocumentContext = createContext<Document<any, BaseCollection<any>> | null>(null);
-
-export function useCollectionPortal<T extends object, C extends BaseCollection<T>>(props: CollectionProps<T>, document?: Document<T, C>): ReactNode {
-  let ctx = useContext(CollectionDocumentContext);
-  let doc = document ?? ctx!;
-  let children = useCollectionChildren(props);
-  let wrappedChildren = useMemo(() => (
-    <ShallowRenderContext.Provider value>
-      {children}
-    </ShallowRenderContext.Provider>
-  ), [children]);
-  // During SSR, we render the content directly, and append nodes to the document during render.
-  // The collection children return null so that nothing is actually rendered into the HTML.
-  return useIsSSR()
-    ? <SSRContext.Provider value={doc}>{wrappedChildren}</SSRContext.Provider>
-    : createPortal(wrappedChildren, doc as unknown as Element);
-}
-
-export function CollectionPortal<T extends object>(props: CollectionProps<T>) {
-  return <>{useCollectionPortal(props)}</>;
-}
 
 export interface ItemRenderProps {
   /**
@@ -919,7 +923,30 @@ export function Collection<T extends object>(props: CollectionProps<T>): JSX.Ele
   let ctx = useContext(CollectionContext)!;
   props = mergeProps(ctx, props);
   props.dependencies = (ctx?.dependencies || []).concat(props.dependencies);
-  return <>{useCollectionChildren(props)}</>;
+  let children = useCollectionChildren(props);
+
+  let doc = useContext(CollectionDocumentContext);
+  if (doc) {
+    return <CollectionRoot>{children}</CollectionRoot>;
+  }
+
+  return <>{children}</>;
+}
+
+function CollectionRoot({children}) {
+  let doc = useContext(CollectionDocumentContext);
+  let wrappedChildren = useMemo(() => (
+    <CollectionDocumentContext.Provider value={null}>
+      <ShallowRenderContext.Provider value>
+        {children}
+      </ShallowRenderContext.Provider>
+    </CollectionDocumentContext.Provider>
+  ), [children]);
+  // During SSR, we render the content directly, and append nodes to the document during render.
+  // The collection children return null so that nothing is actually rendered into the HTML.
+  return useIsSSR()
+    ? <SSRContext.Provider value={doc}>{wrappedChildren}</SSRContext.Provider>
+    : createPortal(wrappedChildren, doc as unknown as Element);
 }
 
 export function createLeafComponent<T extends object, P extends object, E extends Element>(type: string, render: (props: P, ref: ForwardedRef<E>) => JSX.Element): (props: P & React.RefAttributes<T>) => React.ReactElement | null;
@@ -976,7 +1003,7 @@ export interface CollectionRenderer {
   CollectionBranch: React.ComponentType<CollectionBranchProps>
 }
 
-const DefaultCollectionRenderer: CollectionRenderer = {
+export const DefaultCollectionRenderer: CollectionRenderer = {
   CollectionRoot({collection}) {
     return useCachedChildren({
       items: collection,

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

@@ -12,8 +12,8 @@
 import {AriaComboBoxProps, useComboBox, useFilter} from 'react-aria';
 import {ButtonContext} from './Button';
 import {Collection, ComboBoxState, Node, useComboBoxState} from 'react-stately';
-import {CollectionDocumentContext, useCollectionDocument} from './Collection';
-import {ContextValue, forwardRefType, Hidden, Provider, RACValidation, removeDataAttributes, RenderProps, SlotProps, useContextProps, useRenderProps, useSlot, useSlottedContext} from './utils';
+import {CollectionBuilder} from './Collection';
+import {ContextValue, forwardRefType, Provider, RACValidation, removeDataAttributes, RenderProps, SlotProps, useContextProps, useRenderProps, useSlot, useSlottedContext} from './utils';
 import {FieldErrorContext} from './FieldError';
 import {filterDOMProps, useResizeObserver} from '@react-aria/utils';
 import {FormContext} from './Form';
@@ -67,35 +67,25 @@ export const ComboBoxStateContext = createContext<ComboBoxState<any> | null>(nul
 
 function ComboBox<T extends object>(props: ComboBoxProps<T>, ref: ForwardedRef<HTMLDivElement>) {
   [props, ref] = useContextProps(props, ref, ComboBoxContext);
-  let {collection, document} = useCollectionDocument();
   let {children, isDisabled = false, isInvalid = false, isRequired = false} = props;
-  children = useMemo(() => (
-    typeof children === 'function'
-      ? children({
-        isOpen: false,
-        isDisabled,
-        isInvalid,
-        isRequired,
-        defaultChildren: null
-      })
-      : children
-  ), [children, isDisabled, isInvalid, isRequired]);
+  let content = useMemo(() => (
+    <ListBoxContext.Provider value={{items: props.items ?? props.defaultItems}}>
+      {typeof children === 'function'
+        ? children({
+          isOpen: false,
+          isDisabled,
+          isInvalid,
+          isRequired,
+          defaultChildren: null
+        })
+        : children}
+    </ListBoxContext.Provider>
+  ), [children, isDisabled, isInvalid, isRequired, props.items, props.defaultItems]);
 
   return (
-    <>
-      {/* Render a hidden copy of the children so that we can build the collection even when the popover is not open.
-        * This should always come before the real DOM content so we have built the collection by the time it renders during SSR. */}
-      <Hidden>
-        <Provider
-          values={[
-            [CollectionDocumentContext, document],
-            [ListBoxContext, {items: props.items ?? props.defaultItems}]
-          ]}>
-          {children}
-        </Provider>
-      </Hidden>
-      <ComboBoxInner props={props} collection={collection} comboBoxRef={ref} />
-    </>
+    <CollectionBuilder content={content}>
+      {collection => <ComboBoxInner props={props} collection={collection} comboBoxRef={ref} />}
+    </CollectionBuilder>
   );
 }
 

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

@@ -12,10 +12,10 @@
 import {AriaGridListProps, DraggableItemResult, DragPreviewRenderer, DropIndicatorAria, DroppableCollectionResult, FocusScope, ListKeyboardDelegate, mergeProps, useCollator, useFocusRing, useGridList, useGridListItem, useGridListSelectionCheckbox, useHover, useLocale, useVisuallyHidden} from 'react-aria';
 import {ButtonContext} from './Button';
 import {CheckboxContext} from './RSPContexts';
-import {Collection, DraggableCollectionState, DroppableCollectionState, ListState, Node, SelectionBehavior, useListState} from 'react-stately';
-import {CollectionProps, CollectionRendererContext, createLeafComponent, ItemRenderProps, useCollection} from './Collection';
+import {Collection, CollectionBuilder, CollectionProps, CollectionRendererContext, createLeafComponent, ItemRenderProps} from './Collection';
 import {ContextValue, DEFAULT_SLOT, forwardRefType, Provider, RenderProps, ScrollableProps, SlotProps, StyleRenderProps, useContextProps, useRenderProps} from './utils';
 import {DragAndDropContext, DragAndDropHooks, DropIndicator, DropIndicatorContext, DropIndicatorProps} from './useDragAndDrop';
+import {DraggableCollectionState, DroppableCollectionState, Collection as ICollection, ListState, Node, SelectionBehavior, useListState} from 'react-stately';
 import {filterDOMProps, useObjectRef} from '@react-aria/utils';
 import {HoverEvents, Key, LinkDOMProps} from '@react-types/shared';
 import {ListStateContext} from './ListBox';
@@ -74,18 +74,17 @@ export const GridListContext = createContext<ContextValue<GridListProps<any>, HT
 function GridList<T extends object>(props: GridListProps<T>, ref: ForwardedRef<HTMLDivElement>) {
   // Render the portal first so that we have the collection by the time we render the DOM in SSR.
   [props, ref] = useContextProps(props, ref, GridListContext);
-  let {collection, portal} = useCollection(props);
+
   return (
-    <>
-      {portal}
-      <GridListInner props={props} collection={collection} gridListRef={ref} />
-    </>
+    <CollectionBuilder content={<Collection {...props} />}>
+      {collection => <GridListInner props={props} collection={collection} gridListRef={ref} />}
+    </CollectionBuilder>
   );
 }
 
 interface GridListInnerProps<T extends object> {
   props: GridListProps<T>,
-  collection: Collection<Node<T>>,
+  collection: ICollection<Node<object>>,
   gridListRef: RefObject<HTMLDivElement | null>
 }