feat: Add support for grid edit mode w/ nested interactive widgets

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

@@ -14,7 +14,7 @@ import {BaseCollection} from './BaseCollection';
 import {BaseNode, Document, ElementNode} from './Document';
 import {CachedChildrenOptions, useCachedChildren} from './useCachedChildren';
 import {createPortal} from 'react-dom';
-import {forwardRefType, Node} from '@react-types/shared';
+import {forwardRefType, Key, Node} from '@react-types/shared';
 import {Hidden} from './Hidden';
 import React, {createContext, ForwardedRef, forwardRef, JSX, ReactElement, ReactNode, useCallback, useContext, useMemo, useRef, useState} from 'react';
 import {useIsSSR} from '@react-aria/ssr';
@@ -25,6 +25,7 @@ const ShallowRenderContext = createContext(false);
 const CollectionDocumentContext = createContext<Document<any, BaseCollection<any>> | null>(null);
 
 export interface CollectionBuilderProps<C extends BaseCollection<object>> {
+  id?: Key,
   content: ReactNode,
   children: (collection: C) => ReactNode,
   createCollection?: () => C
@@ -51,6 +52,11 @@ 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);
+
+  if (props.id) {
+    document.key = props.id;
+  }
+
   return (
     <>
       <Hidden>

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

@@ -12,7 +12,7 @@
 
 import {BaseCollection, CollectionNode, Mutable} from './BaseCollection';
 import {ForwardedRef, ReactElement} from 'react';
-import {Node} from '@react-types/shared';
+import {Key, Node} from '@react-types/shared';
 
 // 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
@@ -226,7 +226,7 @@ export class ElementNode<T> extends BaseNode<T> {
 
   constructor(type: string, ownerDocument: Document<T, any>) {
     super(ownerDocument);
-    this.node = new CollectionNode(type, `react-aria-${++ownerDocument.nodeId}`);
+    this.node = new CollectionNode(type, `${ownerDocument.key}-${++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.
@@ -304,6 +304,7 @@ export class ElementNode<T> extends BaseNode<T> {
  * which is lazily copied on write during updates.
  */
 export class Document<T, C extends BaseCollection<T> = BaseCollection<T>> extends BaseNode<T> {
+  key: Key = 'react-aria';
   nodeType = 11; // DOCUMENT_FRAGMENT_NODE
   ownerDocument = this;
   dirtyNodes: Set<BaseNode<T>> = new Set();

packages/@react-aria/combobox/src/useComboBox.ts

@@ -14,7 +14,7 @@ import {announce} from '@react-aria/live-announcer';
 import {AriaButtonProps} from '@react-types/button';
 import {AriaComboBoxProps} from '@react-types/combobox';
 import {ariaHideOutside} from '@react-aria/overlays';
-import {AriaListBoxOptions, getItemId, listData} from '@react-aria/listbox';
+import {AriaListBoxOptions, getItemId, listMap} from '@react-aria/listbox';
 import {BaseEvent, DOMAttributes, KeyboardDelegate, LayoutDelegate, PressEvent, RefObject, RouterOptions, ValidationResult} from '@react-types/shared';
 import {chain, isAppleDevice, mergeProps, useLabels, useRouter} from '@react-aria/utils';
 import {ComboBoxState} from '@react-stately/combobox';
@@ -95,7 +95,7 @@ export function useComboBox<T>(props: AriaComboBoxOptions<T>, state: ComboBoxSta
   );
 
   // Set listbox id so it can be used when calling getItemId later
-  listData.set(state, {id: menuProps.id});
+  listMap.set(state, {id: menuProps.id});
 
   // By default, a KeyboardDelegate is provided which uses the DOM to query layout information (e.g. for page up/page down).
   // When virtualized, the layout object will be passed in as a prop and override this.

packages/@react-aria/grid/package.json

@@ -30,6 +30,7 @@
     "@react-aria/utils": "^3.25.3",
     "@react-stately/collections": "^3.11.0",
     "@react-stately/grid": "^3.9.3",
+    "@react-stately/group": "3.0.0-alpha.1",
     "@react-stately/selection": "^3.17.0",
     "@react-types/checkbox": "^3.8.4",
     "@react-types/grid": "^3.2.9",

packages/@react-aria/grid/src/GridKeyboardDelegate.ts

@@ -13,7 +13,7 @@
 import {Direction, DisabledBehavior, Key, KeyboardDelegate, LayoutDelegate, Node, Rect, RefObject, Size} from '@react-types/shared';
 import {DOMLayoutDelegate} from '@react-aria/selection';
 import {getChildNodes, getFirstItem, getLastItem, getNthItem} from '@react-stately/collections';
-import {GridCollection} from '@react-types/grid';
+import {GridCollection, GridNode} from '@react-types/grid';
 
 export interface GridKeyboardDelegateOptions<C> {
   collection: C,
@@ -48,15 +48,11 @@ export class GridKeyboardDelegate<T, C extends GridCollection<T>> implements Key
   }
 
   protected isCell(node: Node<T>) {
-    return node.type === 'cell';
+    return node?.type === 'cell' || node?.type === 'rowheader' || node?.type === 'column';
   }
 
   protected isRow(node: Node<T>) {
-    return node.type === 'row' || node.type === 'item';
-  }
-
-  private isDisabled(item: Node<unknown>) {
-    return this.disabledBehavior === 'all' && (item.props?.isDisabled || this.disabledKeys.has(item.key));
+    return node?.type === 'row' || node?.type === 'item';
   }
 
   protected findPreviousKey(fromKey?: Key, pred?: (item: Node<T>) => boolean) {
@@ -66,7 +62,7 @@ export class GridKeyboardDelegate<T, C extends GridCollection<T>> implements Key
 
     while (key != null) {
       let item = this.collection.getItem(key);
-      if (!this.isDisabled(item) && (!pred || pred(item))) {
+      if (item && !this.isDisabled(item.key) && (!pred || pred(item))) {
         return key;
       }
 
@@ -81,7 +77,7 @@ export class GridKeyboardDelegate<T, C extends GridCollection<T>> implements Key
 
     while (key != null) {
       let item = this.collection.getItem(key);
-      if (!this.isDisabled(item) && (!pred || pred(item))) {
+      if (item && !this.isDisabled(item.key) && (!pred || pred(item))) {
         return key;
       }
 
@@ -101,19 +97,30 @@ export class GridKeyboardDelegate<T, C extends GridCollection<T>> implements Key
     }
 
     // Find the next item
-    key = this.findNextKey(key, (item => item.type === 'item'));
-    if (key != null) {
-      // If focus was on a cell, focus the cell with the same index in the next row.
-      if (this.isCell(startItem)) {
-        let item = this.collection.getItem(key);
-        return getNthItem(getChildNodes(item, this.collection), startItem.index).key;
+    let row = key;
+    let item;
+    let next;
+
+    do {
+      row = this.findNextKey(row, item => item.type === 'item');
+      if (row == null) {
+        break;
       }
 
-      // Otherwise, focus the next row
-      if (this.focusMode === 'row') {
-        return key;
+      // If focus was on a cell, try to focus the cell with the same index in the next row
+      if (this.isCell(startItem)) {
+        item = this.collection.getItem(row);
+        next = getNthItem(getChildNodes(item, this.collection), startItem.index);
+
+        // If we found a non-disabled cell at the same index, use it
+        if (next && !this.isDisabled(next.key)) {
+          return next.key;
+        }
+      } else if (this.focusMode === 'row') {
+        // If in row mode and not starting from a cell, just return the next row
+        return row;
       }
-    }
+    } while (row != null);
   }
 
   getKeyAbove(key: Key) {
@@ -128,19 +135,30 @@ export class GridKeyboardDelegate<T, C extends GridCollection<T>> implements Key
     }
 
     // Find the previous item
-    key = this.findPreviousKey(key, item => item.type === 'item');
-    if (key != null) {
-      // If focus was on a cell, focus the cell with the same index in the previous row.
-      if (this.isCell(startItem)) {
-        let item = this.collection.getItem(key);
-        return getNthItem(getChildNodes(item, this.collection), startItem.index).key;
+    let row = key;
+    let item;
+    let prev;
+
+    do {
+      row = this.findPreviousKey(row, item => item.type === 'item');
+      if (row == null) {
+        break;
       }
 
-      // Otherwise, focus the previous row
-      if (this.focusMode === 'row') {
-        return key;
+      // If focus was on a cell, try to focus the cell with the same index in the previous row
+      if (this.isCell(startItem)) {
+        item = this.collection.getItem(row);
+        prev = getNthItem(getChildNodes(item, this.collection), startItem.index);
+
+        // If we found a non-disabled cell at the same index, use it
+        if (prev && !this.isDisabled(prev.key)) {
+          return prev.key;
+        }
+      } else if (this.focusMode === 'row') {
+        // If in row mode and not starting from a cell, just return the previous row
+        return row;
       }
-    }
+    } while (row != null);
   }
 
   getKeyRightOf(key: Key) {
@@ -153,20 +171,31 @@ export class GridKeyboardDelegate<T, C extends GridCollection<T>> implements Key
     if (this.isRow(item)) {
       let children = getChildNodes(item, this.collection);
       return this.direction === 'rtl'
-        ? getLastItem(children).key
-        : getFirstItem(children).key;
+        ? [...children].reverse().find(child => !this.isDisabled(child.key))?.key
+        : [...children].find(child => !this.isDisabled(child.key))?.key;
     }
 
     // If focus is on a cell, focus the next cell if any,
     // otherwise focus the parent row.
     if (this.isCell(item)) {
       let parent = this.collection.getItem(item.parentKey);
       let children = getChildNodes(parent, this.collection);
-      let next = this.direction === 'rtl'
-        ? getNthItem(children, item.index - 1)
-        : getNthItem(children, item.index + 1);
 
-      if (next) {
+      let next;
+      let offset = 0;
+
+      do {
+        offset++;
+        next = this.direction === 'rtl'
+          ? getNthItem(children, item.index - offset)
+          : getNthItem(children, item.index + offset);
+      } while (next && 
+        this.isDisabled(next.key) && 
+        (this.direction === 'rtl' 
+          ? item.index - offset >= 0
+          : item.index + offset < this.collection.columnCount));
+
+      if (next && !this.isDisabled(next.key)) {
         return next.key;
       }
 
@@ -189,18 +218,28 @@ export class GridKeyboardDelegate<T, C extends GridCollection<T>> implements Key
     if (this.isRow(item)) {
       let children = getChildNodes(item, this.collection);
       return this.direction === 'rtl'
-        ? getFirstItem(children).key
-        : getLastItem(children).key;
+        ? [...children].find(child => !this.isDisabled(child.key))?.key 
+        : [...children].reverse().find(child => !this.isDisabled(child.key))?.key;
     }
 
     // If focus is on a cell, focus the previous cell if any,
     // otherwise focus the parent row.
     if (this.isCell(item)) {
       let parent = this.collection.getItem(item.parentKey);
       let children = getChildNodes(parent, this.collection);
-      let prev = this.direction === 'rtl'
-        ? getNthItem(children, item.index + 1)
-        : getNthItem(children, item.index - 1);
+
+      let prev;
+      let offset = 0;
+
+      do {
+        offset++;
+        prev = this.direction === 'rtl'
+          ? getNthItem(children, item.index + offset)
+          : getNthItem(children, item.index - offset);
+      } while (prev && this.isDisabled(prev.key) && 
+        (this.direction === 'rtl' 
+          ? item.index + offset < this.collection.columnCount
+          : item.index - offset >= 0));
 
       if (prev) {
         return prev.key;
@@ -232,7 +271,7 @@ export class GridKeyboardDelegate<T, C extends GridCollection<T>> implements Key
     }
 
     // Find the first row
-    key = this.findNextKey(null, item => item.type === 'item');
+    key = this.findNextKey(null, item => this.isRow(item));
 
     // If global flag is set (or if focus mode is cell), focus the first cell in the first row.
     if ((key != null && item && this.isCell(item) && global) || this.focusMode === 'cell') {
@@ -262,7 +301,7 @@ export class GridKeyboardDelegate<T, C extends GridCollection<T>> implements Key
     }
 
     // Find the last row
-    key = this.findPreviousKey(null, item => item.type === 'item');
+    key = this.findPreviousKey(null, item => this.isRow(item));
 
     // If global flag is set (or if focus mode is cell), focus the last cell in the last row.
     if ((key != null && item && this.isCell(item) && global) || this.focusMode === 'cell') {
@@ -345,7 +384,7 @@ export class GridKeyboardDelegate<T, C extends GridCollection<T>> implements Key
         }
       }
 
-      key = this.findNextKey(key, item => item.type === 'item');
+      key = this.findNextKey(key, item => this.isRow(item));
 
       // Wrap around when reaching the end of the collection
       if (key == null && !hasWrapped) {
@@ -356,6 +395,20 @@ export class GridKeyboardDelegate<T, C extends GridCollection<T>> implements Key
 
     return null;
   }
+
+  isDisabled(key: Key) {
+    let item = this.collection.getItem(key) as GridNode<T>;
+
+    if (this.disabledBehavior === 'all' && (this.disabledKeys.has(key) || !!item?.props?.isDisabled)) {
+      return true;
+    }
+
+    if (key === item?.column?.key) {
+      return false;
+    }
+
+    return this.isCell(item) && (this.isDisabled(item.parentKey) || this.isDisabled(item.column?.key));
+  }
 }
 
 /* Backward compatibility for old Virtualizer Layout interface. */