import React, { cloneElement, createContext, ForwardedRef, ReactElement, ReactNode, useCallback, useContext, useMemo, } from "react"; import { createPortal } from "react-dom"; // This Collection implementation is perhaps a little unusual. It works by rendering the React tree into a // Portal to a fake DOM implementation. This gives us efficient access to the tree of rendered objects, and // supports React features like composition and context. We use this fake DOM to access the full set of elements // before we render into the real DOM, which allows us to render a subset of the elements (e.g. virtualized scrolling), // and compute properties like the total number of items. It also enables keyboard navigation, selection, and other features. // React takes care of efficiently rendering components and updating the collection for us via this fake DOM. // // The DOM is a mutable API, and React expects the node instances to remain stable over time. So the implementation is split // into two parts. Each mutable fake DOM node owns an instance of an immutable collection node. When a fake DOM node is updated, // it queues a second render for the collection. Multiple updates to a collection can be queued at once. Collection nodes are // lazily copied on write, so only the changed nodes need to be cloned. During the second render, the new immutable collection // is finalized by updating the map of Key -> Node with the new cloned nodes. Then the new collection is frozen so it can no // longer be mutated, and returned to the calling component to render. type Mutable = { -readonly [P in keyof T]: T[P]; }; /** An immutable object representing a Node in a Collection. */ class NodeValue implements Node { readonly type: string; readonly key: Key; readonly value: T | null = null; readonly level: number = 0; readonly hasChildNodes: boolean = false; readonly rendered: ReactNode = null; readonly textValue: string = ""; readonly "aria-label"?: string = undefined; readonly index: number = 0; readonly parentKey: Key | null = null; readonly prevKey: Key | null = null; readonly nextKey: Key | null = null; readonly firstChildKey: Key | null = null; readonly lastChildKey: Key | null = null; readonly props: any = {}; constructor(type: string, key: Key) { this.type = type; this.key = key; } get childNodes(): Iterable> { throw new Error("childNodes is not supported"); } clone(): NodeValue { const node: Mutable> = new NodeValue(this.type, this.key); node.value = this.value; node.level = this.level; node.hasChildNodes = this.hasChildNodes; node.rendered = this.rendered; node.textValue = this.textValue; node["aria-label"] = this["aria-label"]; node.index = this.index; node.parentKey = this.parentKey; node.prevKey = this.prevKey; node.nextKey = this.nextKey; node.firstChildKey = this.firstChildKey; node.lastChildKey = this.lastChildKey; node.props = this.props; return node; } } /** * A mutable node in the fake DOM tree. When mutated, it marks itself as dirty * and queues an update with the owner document. */ class BaseNode { private _firstChild: ElementNode | null = null; private _lastChild: ElementNode | null = null; private _previousSibling: ElementNode | null = null; private _nextSibling: ElementNode | null = null; private _parentNode: BaseNode | null = null; ownerDocument: Document; constructor(ownerDocument: Document) { this.ownerDocument = ownerDocument; } *[Symbol.iterator]() { let node = this.firstChild; while (node) { yield node; node = node.nextSibling; } } get firstChild() { return this._firstChild; } set firstChild(firstChild) { this._firstChild = firstChild; this.ownerDocument.markDirty(this); } get lastChild() { return this._lastChild; } set lastChild(lastChild) { this._lastChild = lastChild; this.ownerDocument.markDirty(this); } get previousSibling() { return this._previousSibling; } set previousSibling(previousSibling) { this._previousSibling = previousSibling; this.ownerDocument.markDirty(this); } get nextSibling() { return this._nextSibling; } set nextSibling(nextSibling) { this._nextSibling = nextSibling; this.ownerDocument.markDirty(this); } get parentNode() { return this._parentNode; } set parentNode(parentNode) { this._parentNode = parentNode; this.ownerDocument.markDirty(this); } get isConnected() { return this.parentNode?.isConnected || false; } appendChild(child: ElementNode) { this.ownerDocument.startTransaction(); if (child.parentNode) { child.parentNode.removeChild(child); } if (this.firstChild == null) { this.firstChild = child; } if (this.lastChild) { this.lastChild.nextSibling = child; child.index = this.lastChild.index + 1; child.previousSibling = this.lastChild; } else { child.previousSibling = null; child.index = 0; } child.parentNode = this; child.nextSibling = null; this.lastChild = child; this.ownerDocument.markDirty(this); if (child.hasSetProps) { // Only add the node to the collection if we already received props for it. // Otherwise wait until then so we have the correct id for the node. this.ownerDocument.addNode(child); } this.ownerDocument.endTransaction(); this.ownerDocument.queueUpdate(); } insertBefore(newNode: ElementNode, referenceNode: ElementNode) { if (referenceNode == null) { return this.appendChild(newNode); } this.ownerDocument.startTransaction(); if (newNode.parentNode) { newNode.parentNode.removeChild(newNode); } newNode.nextSibling = referenceNode; newNode.previousSibling = referenceNode.previousSibling; newNode.index = referenceNode.index; if (this.firstChild === referenceNode) { this.firstChild = newNode; } else if (referenceNode.previousSibling) { referenceNode.previousSibling.nextSibling = newNode; } referenceNode.previousSibling = newNode; newNode.parentNode = referenceNode.parentNode; let node: ElementNode | null = referenceNode; while (node) { node.index++; node = node.nextSibling; } if (newNode.hasSetProps) { this.ownerDocument.addNode(newNode); } this.ownerDocument.endTransaction(); this.ownerDocument.queueUpdate(); } removeChild(child: ElementNode) { if (child.parentNode !== this) { return; } this.ownerDocument.startTransaction(); let node = child.nextSibling; while (node) { node.index--; node = node.nextSibling; } if (child.nextSibling) { child.nextSibling.previousSibling = child.previousSibling; } if (child.previousSibling) { child.previousSibling.nextSibling = child.nextSibling; } if (this.firstChild === child) { this.firstChild = child.nextSibling; } if (this.lastChild === child) { this.lastChild = child.previousSibling; } child.parentNode = null; child.nextSibling = null; child.previousSibling = null; child.index = 0; this.ownerDocument.removeNode(child); this.ownerDocument.endTransaction(); this.ownerDocument.queueUpdate(); } addEventListener() {} removeEventListener() {} } /** * A mutable element node in the fake DOM tree. It owns an immutable * Collection Node which is copied on write. */ class ElementNode extends BaseNode { nodeType = 8; // COMMENT_NODE (we'd use ELEMENT_NODE but React DevTools will fail to get its dimensions) node: NodeValue; private _index: number = 0; hasSetProps = false; constructor(type: string, ownerDocument: Document) { super(ownerDocument); this.node = new NodeValue(type, `react-aria-${++ownerDocument.nodeId}`); // Start a transaction so that no updates are emitted from the collection // until the props for this node are set. We don't know the real id for the // node until then, so we need to avoid emitting collections in an inconsistent state. this.ownerDocument.startTransaction(); } get index() { return this._index; } set index(index) { this._index = index; this.ownerDocument.markDirty(this); } get level(): number { if (this.parentNode instanceof ElementNode) { return this.parentNode.level + (this.node.type === "item" ? 1 : 0); } return 0; } updateNode() { const node = this.ownerDocument.getMutableNode(this); node.index = this.index; node.level = this.level; node.parentKey = this.parentNode instanceof ElementNode ? this.parentNode.node.key : null; node.prevKey = this.previousSibling?.node.key ?? null; node.nextKey = this.nextSibling?.node.key ?? null; node.hasChildNodes = !!this.firstChild; node.firstChildKey = this.firstChild?.node.key ?? null; node.lastChildKey = this.lastChild?.node.key ?? null; } setProps(obj: any, ref: ForwardedRef, rendered?: any) { const node = this.ownerDocument.getMutableNode(this); const { value, textValue, id, ...props } = obj; props.ref = ref; node.props = obj; node.rendered = rendered; node.value = value; node.textValue = textValue || (typeof rendered === "string" ? rendered : "") || obj["aria-label"] || ""; if (id != null && id !== node.key) { if (this.hasSetProps) { throw new Error("Cannot change the id of an item"); } node.key = id; } // If this is the first time props have been set, end the transaction started in the constructor // so this node can be emitted. if (!this.hasSetProps) { this.ownerDocument.addNode(this); this.ownerDocument.endTransaction(); this.hasSetProps = true; } this.ownerDocument.queueUpdate(); } get style() { return {}; } hasAttribute() {} setAttribute() {} setAttributeNS() {} removeAttribute() {} } /** * An immutable Collection implementation. Updates are only allowed * when it is not marked as frozen. */ export class BaseCollection { private keyMap: Map> = new Map(); private firstKey: Key | null = null; private lastKey: Key | null = null; private frozen = false; get size() { return this.keyMap.size; } getKeys() { return this.keyMap.keys(); } *[Symbol.iterator]() { let node: Node | undefined = this.firstKey != null ? this.keyMap.get(this.firstKey) : undefined; while (node) { yield node; node = node.nextKey != null ? this.keyMap.get(node.nextKey) : undefined; } } getChildren(key: Key): Iterable> { const keyMap = this.keyMap; return { *[Symbol.iterator]() { const parent = keyMap.get(key); let node = parent?.firstChildKey != null ? keyMap.get(parent.firstChildKey) : null; while (node) { yield node as Node; node = node.nextKey != null ? keyMap.get(node.nextKey) : undefined; } }, }; } getKeyBefore(key: Key) { let node = this.keyMap.get(key); if (!node) { return null; } if (node.prevKey != null) { node = this.keyMap.get(node.prevKey); while (node && node.type !== "item" && node.lastChildKey != null) { node = this.keyMap.get(node.lastChildKey); } return node?.key ?? null; } return node.parentKey; } getKeyAfter(key: Key) { let node = this.keyMap.get(key); if (!node) { return null; } if (node.type !== "item" && node.firstChildKey != null) { return node.firstChildKey; } while (node) { if (node.nextKey != null) { return node.nextKey; } if (node.parentKey != null) { node = this.keyMap.get(node.parentKey); } else { return null; } } return null; } getFirstKey() { return this.firstKey; } getLastKey() { let node = this.lastKey != null ? this.keyMap.get(this.lastKey) : null; while (node?.lastChildKey != null) { node = this.keyMap.get(node.lastChildKey); } return node?.key ?? null; } getItem(key: Key): Node | null { return this.keyMap.get(key) ?? null; } at(): Node { throw new Error("Not implemented"); } clone(): this { // We need to clone using this.constructor so that subclasses have the right prototype. // TypeScript isn't happy about this yet. // https://github.com/microsoft/TypeScript/issues/3841 const Constructor: any = this.constructor; const collection: this = new Constructor(); collection.keyMap = new Map(this.keyMap); collection.firstKey = this.firstKey; collection.lastKey = this.lastKey; return collection; } addNode(node: NodeValue) { if (this.frozen) { throw new Error("Cannot add a node to a frozen collection"); } this.keyMap.set(node.key, node); } removeNode(key: Key) { if (this.frozen) { throw new Error("Cannot remove a node to a frozen collection"); } this.keyMap.delete(key); } commit(firstKey: Key | null, lastKey: Key | null, isSSR = false) { if (this.frozen) { throw new Error("Cannot commit a frozen collection"); } this.firstKey = firstKey; this.lastKey = lastKey; this.frozen = !isSSR; } } /** * A mutable Document in the fake DOM. It owns an immutable Collection instance, * which is lazily copied on write during updates. */ class Document< T, C extends BaseCollection = BaseCollection, > extends BaseNode { nodeType = 11; // DOCUMENT_FRAGMENT_NODE ownerDocument = this; dirtyNodes: Set> = new Set(); isSSR = false; nodeId = 0; nodesByProps = new WeakMap>(); private collection: C; private collectionMutated: boolean; private mutatedNodes: Set> = new Set(); private subscriptions: Set<() => void> = new Set(); private transactionCount = 0; constructor(collection: C) { // @ts-expect-error - we don't want to call the base constructor super(null); this.collection = collection; this.collectionMutated = true; } get isConnected() { return true; } createElement(type: string) { return new ElementNode(type, this); } /** * Lazily gets a mutable instance of a Node. If the node has already * been cloned during this update cycle, it just returns the existing one. */ getMutableNode(element: ElementNode): Mutable> { let node = element.node; if (!this.mutatedNodes.has(element)) { node = element.node.clone(); this.mutatedNodes.add(element); element.node = node; } this.markDirty(element); return node; } private getMutableCollection() { if (!this.isSSR && !this.collectionMutated) { this.collection = this.collection.clone(); this.collectionMutated = true; } return this.collection; } markDirty(node: BaseNode) { this.dirtyNodes.add(node); } startTransaction() { this.transactionCount++; } endTransaction() { this.transactionCount--; } addNode(element: ElementNode) { const collection = this.getMutableCollection(); if (!collection.getItem(element.node.key)) { collection.addNode(element.node); for (const child of element) { this.addNode(child); } } this.markDirty(element); } removeNode(node: ElementNode) { for (const child of node) { this.removeNode(child); } const collection = this.getMutableCollection(); collection.removeNode(node.node.key); this.markDirty(node); } /** Finalizes the collection update, updating all nodes and freezing the collection. */ getCollection(): C { if (this.transactionCount > 0) { return this.collection; } this.updateCollection(); return this.collection; } updateCollection() { for (const element of this.dirtyNodes) { if (element instanceof ElementNode && element.isConnected) { element.updateNode(); } } this.dirtyNodes.clear(); if (this.mutatedNodes.size) { const collection = this.getMutableCollection(); for (const element of this.mutatedNodes) { if (element.isConnected) { collection.addNode(element.node); } } collection.commit( this.firstChild?.node.key ?? null, this.lastChild?.node.key ?? null, this.isSSR, ); this.mutatedNodes.clear(); } this.collectionMutated = false; } queueUpdate() { // Don't emit any updates if there is a transaction in progress. // queueUpdate should be called again after the transaction. if (this.dirtyNodes.size === 0 || this.transactionCount > 0) { return; } for (const fn of this.subscriptions) { fn(); } } subscribe(fn: () => void) { this.subscriptions.add(fn); return () => this.subscriptions.delete(fn); } resetAfterSSR() { if (this.isSSR) { this.isSSR = false; this.firstChild = null; this.lastChild = null; this.nodeId = 0; } } } interface CollectionProps { items?: Iterable; /** The contents of the collection. */ children?: ReactNode | ((item: T) => ReactNode); /** Values that should invalidate the item cache when using dynamic collections. */ dependencies?: any[]; } interface CachedChildrenOptions extends CollectionProps { idScope?: Key; addIdAndValue?: boolean; dependencies?: any[]; } export function useCachedChildren( props: CachedChildrenOptions, ): any { const { children, items, idScope, addIdAndValue, dependencies = [] } = props; // Invalidate the cache whenever the parent value changes. // eslint-disable-next-line react-hooks/exhaustive-deps const cache = useMemo(() => new WeakMap(), dependencies); return useMemo(() => { if (items && typeof children === "function") { const res: ReactElement[] = []; for (const item of items) { let rendered = cache.get(item); if (!rendered) { rendered = children(item); // @ts-expect-error - unknown let key = rendered.props.id ?? item.key ?? item.id; // eslint-disable-next-line max-depth if (key == null) { throw new Error("Could not determine key for item"); } // eslint-disable-next-line max-depth if (idScope) { key = idScope + ":" + key; } // Note: only works if wrapped Item passes through id... rendered = cloneElement( rendered, addIdAndValue ? { key, id: key, value: item } : { key }, ); cache.set(item, rendered); } res.push(rendered); } return res; } else if (typeof children !== "function") { return children; } }, [children, items, cache, idScope, addIdAndValue]); } function useCollectionChildren( props: CachedChildrenOptions, ) { return useCachedChildren({ ...props, addIdAndValue: true }); } const ShallowRenderContext = createContext(false); interface CollectionResult { portal: ReactNode; collection: C; } export function useCollection>( props: CollectionProps, initialCollection?: C, ): CollectionResult { const { collection, document } = useCollectionDocument( initialCollection, ); const portal = useCollectionPortal(props, document); return { portal, collection }; } interface CollectionDocumentResult> { collection: C; document: Document; } function useCollectionDocument>( initialCollection?: C, ): CollectionDocumentResult { // The document instance is mutable, and should never change between renders. // useSyncExternalStore is used to subscribe to updates, which vends immutable Collection objects. const document = useMemo( () => new Document(initialCollection || (new BaseCollection() as C)), [initialCollection], ); const subscribe = useCallback( (fn: () => void) => document.subscribe(fn), [document], ); const getSnapshot = useCallback(() => { const collection = document.getCollection(); if (document.isSSR) { // After SSR is complete, reset the document to empty so it is ready for React to render the portal into. // We do this _after_ getting the collection above so that the collection still has content in it from SSR // during the current render, before React has finished the client render. document.resetAfterSSR(); } return collection; }, [document]); const getServerSnapshot = useCallback(() => { document.isSSR = true; return document.getCollection(); }, [document]); const collection = React.useSyncExternalStore( subscribe, getSnapshot, getServerSnapshot, ); return { collection, document }; } const SSRContext = createContext | null>(null); const CollectionDocumentContext = createContext > | null>(null); function useCollectionPortal>( props: CollectionProps, document?: Document, ): ReactNode { const ctx = useContext(CollectionDocumentContext); const doc = document ?? ctx!; const children = useCollectionChildren(props); const wrappedChildren = useMemo( () => ( {children} ), [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() ? ( {wrappedChildren} ) : ( createPortal(wrappedChildren, doc as unknown as Element) ); } export function CollectionPortal(props: CollectionProps) { return <>{useCollectionPortal(props)}; } /** Renders a DOM element (e.g. separator or header) shallowly when inside a collection. */ export function useShallowRender

( type: string, props: P, ref: ForwardedRef, ): ReactElement | null { const isShallow = useContext(ShallowRenderContext); if (isShallow) { // Elements cannot be re-parented, so the context will always be there. return ( // eslint-disable-next-line react-hooks/rules-of-hooks useSSRCollectionNode( type, props, ref, "children" in props ? props.children : null, ) ?? <> ); } return null; } function useCollectionItemRef( props: any, ref: ForwardedRef, rendered?: any, ) { // Return a callback ref that sets the props object on the fake DOM node. return useCallback( (element) => { element?.setProps(props, ref, rendered); }, [props, ref, rendered], ); } export function useSSRCollectionNode( Type: string, props: object, ref: ForwardedRef, rendered?: any, children?: ReactNode, ) { // During SSR, portals are not supported, so the collection children will be wrapped in an SSRContext. // Since SSR occurs only once, we assume that the elements are rendered in order and never re-render. // Therefore we can create elements in our collection document during render so that they are in the // collection by the time we need to use the collection to render to the real DOM. // After hydration, we switch to client rendering using the portal. const itemRef = useCollectionItemRef(props, ref, rendered); const parentNode = useContext(SSRContext); if (parentNode) { // Guard against double rendering in strict mode. let element = parentNode.ownerDocument.nodesByProps.get(props); if (!element) { element = parentNode.ownerDocument.createElement(Type); element.setProps(props, ref, rendered); parentNode.appendChild(element); parentNode.ownerDocument.updateCollection(); parentNode.ownerDocument.nodesByProps.set(props, element); } return children ? ( {children} ) : null; } // @ts-expect-error - unknown return {children}; } const IsSSRContext = React.createContext(false); function getSnapshot() { return false; } function getServerSnapshot() { return true; } // eslint-disable-next-line @typescript-eslint/no-unused-vars function subscribe(onStoreChange: () => void): () => void { // noop return () => {}; } /** * Returns whether the component is currently being server side rendered or * hydrated on the client. Can be used to delay browser-specific rendering * until after hydration. */ function useIsSSR(): boolean { // In React 18, we can use useSyncExternalStore to detect if we're server rendering or hydrating. if (typeof React["useSyncExternalStore"] === "function") { return React["useSyncExternalStore"]( subscribe, getSnapshot, getServerSnapshot, ); } // eslint-disable-next-line react-hooks/rules-of-hooks return React.useContext(IsSSRContext); } type Key = string | number; export interface Node { /** The type of item this node represents. */ type: string; /** A unique key for the node. */ key: Key; /** The object value the node was created from. */ value: T | null; /** The level of depth this node is at in the heirarchy. */ level: number; /** Whether this item has children, even if not loaded yet. */ hasChildNodes: boolean; /** * The loaded children of this node. * @deprecated Use `collection.getChildren(node.key)` instead. */ childNodes: Iterable>; /** The rendered contents of this node (e.g. JSX). */ rendered: ReactNode; /** A string value for this node, used for features like typeahead. */ textValue: string; /** An accessibility label for this node. */ "aria-label"?: string; /** The index of this node within its parent. */ index?: number; /** A function that should be called to wrap the rendered node. */ wrapper?: (element: ReactElement) => ReactElement; /** The key of the parent node. */ parentKey?: Key | null; /** The key of the node before this node. */ prevKey?: Key | null; /** The key of the node after this node. */ nextKey?: Key | null; /** Additional properties specific to a particular node type. */ props?: any; /** @private */ shouldInvalidate?: (context: unknown) => boolean; }