(WIP) Add subdialog support to Menu and Autocomplete

packages/@react-aria/autocomplete/src/useAutocomplete.ts

@@ -66,6 +66,12 @@ export function UNSTABLE_useAutocomplete(props: AriaAutocompleteOptions, state:
   let delayNextActiveDescendant = useRef(false);
   let queuedActiveDescendant = useRef(null);
   let lastCollectionNode = useRef<HTMLElement>(null);
+  // Stores the previously focused item id if it was cleared via ArrowLeft/Right. Used to dispatch keyboard events to the proper item
+  // even though we've cleared state.focusedNodeId so that things like ArrowLeft/Right will still open the submenutrigger after it is closed
+  // TODO: ideally, we'd just preserve state.focusedNodeId if the user's ArrowLeft/Right was being used to trigger the focused submenutrigger but
+  // that would involve differentiating that event from a moving the text cursor in the input. Will be a moot point if/when NVDA announces
+  // moving the text cursor when an activedescendant is set properly.
+  let clearedFocusedId = useRef<string | null>(null);
 
   let updateActiveDescendant = useEffectEvent((e) => {
     let {target} = e;
@@ -82,12 +88,15 @@ export function UNSTABLE_useAutocomplete(props: AriaAutocompleteOptions, state:
         timeout.current = setTimeout(() => {
           state.setFocusedNodeId(target.id);
           queuedActiveDescendant.current = null;
+          clearedFocusedId.current = null;
         }, 500);
       } else {
         state.setFocusedNodeId(target.id);
+        clearedFocusedId.current = null;
       }
     } else {
       state.setFocusedNodeId(null);
+      clearedFocusedId.current = null;
     }
 
     delayNextActiveDescendant.current = false;
@@ -123,11 +132,20 @@ export function UNSTABLE_useAutocomplete(props: AriaAutocompleteOptions, state:
     );
   });
 
-  let clearVirtualFocus = useEffectEvent(() => {
+  let clearVirtualFocus = useEffectEvent((clearFocusKey?: boolean) => {
+    if (clearFocusKey === false && state.focusedNodeId) {
+      clearedFocusedId.current = state.focusedNodeId;
+    } else if (clearFocusKey) {
+      clearedFocusedId.current = null;
+    }
+
     state.setFocusedNodeId(null);
     let clearFocusEvent = new CustomEvent(CLEAR_FOCUS_EVENT, {
       cancelable: true,
-      bubbles: true
+      bubbles: true,
+      detail: {
+        clearFocusKey
+      }
     });
     clearTimeout(timeout.current);
     delayNextActiveDescendant.current = false;
@@ -141,7 +159,8 @@ export function UNSTABLE_useAutocomplete(props: AriaAutocompleteOptions, state:
     if (state.inputValue !== value && state.inputValue.length <= value.length) {
       focusFirstItem();
     } else {
-      clearVirtualFocus();
+      // Fully clear focused key when backspacing since the list may change and thus we'd want to start fresh again
+      clearVirtualFocus(true);
     }
 
     state.setInputValue(value);
@@ -196,10 +215,11 @@ export function UNSTABLE_useAutocomplete(props: AriaAutocompleteOptions, state:
       }
       case 'ArrowLeft':
       case 'ArrowRight':
-        // TODO: will need to special case this so it doesn't clear the focused key if we are currently
-        // focused on a submenutrigger? May not need to since focus would
-        // But what about wrapped grids where ArrowLeft and ArrowRight should navigate left/right
-        clearVirtualFocus();
+        // Clear the activedescendant so NVDA announcements aren't interrupted but retain the focused key in the collection so the
+        // user's keyboard navigation restarts from where they left off
+        // TODO: What about wrapped grids where ArrowLeft and ArrowRight should navigate left/right? Is it weird that the focused key will remain visible
+        // but activedescendant got cleared? If so, then we'll need to know if we are pressing Left/Right on a submenu/dialog trigger...
+        clearVirtualFocus(false);
         break;
     }
 
@@ -211,12 +231,13 @@ export function UNSTABLE_useAutocomplete(props: AriaAutocompleteOptions, state:
       e.stopPropagation();
     }
 
-    if (state.focusedNodeId == null) {
+    let focusedElementId = state.focusedNodeId ?? clearedFocusedId.current;
+    if (focusedElementId == null) {
       collectionRef.current?.dispatchEvent(
         new KeyboardEvent(e.nativeEvent.type, e.nativeEvent)
       );
     } else {
-      let item = document.getElementById(state.focusedNodeId);
+      let item = document.getElementById(focusedElementId);
       item?.dispatchEvent(
         new KeyboardEvent(e.nativeEvent.type, e.nativeEvent)
       );

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

@@ -224,7 +224,7 @@ export class BaseCollection<T> implements ICollection<Node<T>> {
         let clonedSection: Mutable<CollectionNode<T>> = (node as CollectionNode<T>).clone();
         let lastChildInSection: Mutable<CollectionNode<T>> | null = null;
         for (let child of this.getChildren(node.key)) {
-          if (filterFn(child.textValue) || child.type === 'header') {
+          if (shouldKeepNode(child, filterFn, this, newCollection)) {
             let clonedChild: Mutable<CollectionNode<T>> = (child as CollectionNode<T>).clone();
             // eslint-disable-next-line max-depth
             if (lastChildInSection == null) {
@@ -284,22 +284,25 @@ export class BaseCollection<T> implements ICollection<Node<T>> {
           lastNode = clonedSeparator;
           newCollection.addNode(clonedSeparator);
         }
-      } else if (filterFn(node.textValue)) {
+      } else {
+        // At this point, the node is either a subdialogtrigger node or a standard row/item
         let clonedNode: Mutable<CollectionNode<T>> = (node as CollectionNode<T>).clone();
-        if (newCollection.firstKey == null) {
-          newCollection.firstKey = clonedNode.key;
-        }
+        if (shouldKeepNode(clonedNode, filterFn, this, newCollection)) {
+          if (newCollection.firstKey == null) {
+            newCollection.firstKey = clonedNode.key;
+          }
 
-        if (lastNode != null && (lastNode.type !== 'section' && lastNode.type !== 'separator') && lastNode.parentKey === clonedNode.parentKey) {
-          lastNode.nextKey = clonedNode.key;
-          clonedNode.prevKey = lastNode.key;
-        } else {
-          clonedNode.prevKey = null;
-        }
+          if (lastNode != null && (lastNode.type !== 'section' && lastNode.type !== 'separator') && lastNode.parentKey === clonedNode.parentKey) {
+            lastNode.nextKey = clonedNode.key;
+            clonedNode.prevKey = lastNode.key;
+          } else {
+            clonedNode.prevKey = null;
+          }
 
-        clonedNode.nextKey = null;
-        newCollection.addNode(clonedNode);
-        lastNode = clonedNode;
+          clonedNode.nextKey = null;
+          newCollection.addNode(clonedNode);
+          lastNode = clonedNode;
+        }
       }
     }
 
@@ -318,3 +321,22 @@ export class BaseCollection<T> implements ICollection<Node<T>> {
     return newCollection;
   }
 }
+
+function shouldKeepNode<T>(node: Node<T>, filterFn: (nodeValue: string) => boolean, oldCollection: BaseCollection<T>, newCollection: BaseCollection<T>): boolean {
+  if (node.type === 'subdialogtrigger' || node.type === 'submenutrigger') {
+    // Subdialog wrapper should only have one child, if it passes the filter add it to the new collection since we don't need to
+    // do any extra handling for its first/next key
+    let triggerChild = [...oldCollection.getChildren(node.key)][0];
+    if (triggerChild && filterFn(triggerChild.textValue)) {
+      let clonedChild: Mutable<CollectionNode<T>> = (triggerChild as CollectionNode<T>).clone();
+      newCollection.addNode(clonedChild);
+      return true;
+    } else {
+      return false;
+    }
+  } else if (node.type === 'header') {
+    return true;
+  } else {
+    return filterFn(node.textValue);
+  }
+}

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

@@ -27,6 +27,7 @@
   },
   "dependencies": {
     "@react-aria/focus": "^3.19.1",
+    "@react-aria/interactions": "^3.22.5",
     "@react-aria/overlays": "^3.25.0",
     "@react-aria/utils": "^3.27.0",
     "@react-types/dialog": "^3.5.15",

packages/@react-aria/dialog/src/useDialog.ts

@@ -15,6 +15,7 @@ import {DOMAttributes, FocusableElement, RefObject} from '@react-types/shared';
 import {filterDOMProps, useSlotId} from '@react-aria/utils';
 import {focusSafely} from '@react-aria/focus';
 import {useEffect, useRef} from 'react';
+import {useKeyboard} from '@react-aria/interactions';
 import {useOverlayFocusContain} from '@react-aria/overlays';
 
 export interface DialogAria {
@@ -30,7 +31,11 @@ export interface DialogAria {
  * A dialog is an overlay shown above other content in an application.
  */
 export function useDialog(props: AriaDialogProps, ref: RefObject<FocusableElement | null>): DialogAria {
-  let {role = 'dialog'} = props;
+  let {
+    role = 'dialog',
+    onKeyUp,
+    onKeyDown
+  } = props;
   let titleId: string | undefined = useSlotId();
   titleId = props['aria-label'] ? undefined : titleId;
 
@@ -62,6 +67,8 @@ export function useDialog(props: AriaDialogProps, ref: RefObject<FocusableElemen
   }, [ref]);
 
   useOverlayFocusContain();
+  // TODO: keep in mind this will stop propagation and prevent the Esc handler in useOverlay from firing I believe
+  let {keyboardProps} = useKeyboard({onKeyDown, onKeyUp});
 
   // We do not use aria-modal due to a Safari bug which forces the first focusable element to be focused
   // on mount when inside an iframe, no matter which element we programmatically focus.
@@ -71,6 +78,7 @@ export function useDialog(props: AriaDialogProps, ref: RefObject<FocusableElemen
   return {
     dialogProps: {
       ...filterDOMProps(props, {labelable: true}),
+      ...keyboardProps,
       role,
       tabIndex: -1,
       'aria-labelledby': props['aria-labelledby'] || titleId,

packages/@react-aria/menu/src/useMenuItem.ts

@@ -17,6 +17,7 @@ import {isFocusVisible, useFocus, useHover, useKeyboard, usePress} from '@react-
 import {menuData} from './utils';
 import {SelectionManager} from '@react-stately/selection';
 import {TreeState} from '@react-stately/tree';
+import {useEffect} from 'react';
 import {useSelectableItem} from '@react-aria/selection';
 
 export interface MenuItemAria {
@@ -189,6 +190,12 @@ export function useMenuItem<T>(props: AriaMenuItemProps, state: TreeState<T>, re
   }
 
   let onPressStart = (e: PressEvent) => {
+    // TODO: ideally this would be done in useselectableItem but we don't apply that hooks onPress if the node is a menu trigger
+    if (data.shouldUseVirtualFocus && (e.pointerType === 'virtual' || e.pointerType === 'keyboard')) {
+      selectionManager.setFocused(true);
+      selectionManager.setFocusedKey(key);
+    }
+
     if (e.pointerType === 'keyboard') {
       performAction(e);
     }
@@ -205,6 +212,11 @@ export function useMenuItem<T>(props: AriaMenuItemProps, state: TreeState<T>, re
   };
 
   let onPressUp = (e: PressEvent) => {
+    if (data.shouldUseVirtualFocus && (e.pointerType === 'touch' || e.pointerType === 'mouse')) {
+      selectionManager.setFocused(true);
+      selectionManager.setFocusedKey(key);
+    }
+
     // If interacting with mouse, allow the user to mouse down on the trigger button,
     // drag, and release over an item (matching native behavior).
     if (e.pointerType === 'mouse') {
@@ -299,6 +311,14 @@ export function useMenuItem<T>(props: AriaMenuItemProps, state: TreeState<T>, re
   delete domProps.id;
   let linkProps = useLinkProps(item?.props);
 
+  useEffect(() => {
+    if (isTrigger && data.shouldUseVirtualFocus && isTriggerExpanded && key !== selectionManager.focusedKey) {
+      // If using virtual focus, we need to fake a blur event when virtual focus moves away from an open submenutrigger since we won't actual trigger a real
+      // blur event. This is so the submenu will close when the user hovers/keyboard navigates to another sibiling menu item
+      ref.current?.dispatchEvent(new FocusEvent('focusout', {bubbles: true}));
+    }
+  }, [data.shouldUseVirtualFocus, isTrigger, isTriggerExpanded, key, selectionManager, ref]);
+
   return {
     menuItemProps: {
       ...ariaProps,

packages/@react-aria/menu/src/useSubmenuTrigger.ts

@@ -38,7 +38,9 @@ export interface AriaSubmenuTriggerProps {
    * The delay time in milliseconds for the submenu to appear after hovering over the trigger.
    * @default 200
    */
-  delay?: number
+  delay?: number,
+  /** Whether the submenu trigger uses virtual focus. */
+  isVirtualFocus?: boolean
 }
 
 interface SubmenuTriggerProps extends Omit<AriaMenuItemProps, 'key'> {
@@ -67,7 +69,7 @@ export interface SubmenuTriggerAria<T> {
  * @param ref - Ref to the submenu trigger element.
  */
 export function useSubmenuTrigger<T>(props: AriaSubmenuTriggerProps, state: SubmenuTriggerState, ref: RefObject<FocusableElement | null>): SubmenuTriggerAria<T> {
-  let {parentMenuRef, submenuRef, type = 'menu', isDisabled, delay = 200} = props;
+  let {parentMenuRef, submenuRef, type = 'menu', isDisabled, delay = 200, isVirtualFocus} = props;
   let submenuTriggerId = useId();
   let overlayId = useId();
   let {direction} = useLocale();
@@ -101,14 +103,18 @@ export function useSubmenuTrigger<T>(props: AriaSubmenuTriggerProps, state: Subm
         if (direction === 'ltr' && e.currentTarget.contains(e.target as Element)) {
           e.stopPropagation();
           onSubmenuClose();
-          ref.current?.focus();
+          if (!isVirtualFocus) {
+            ref.current?.focus();
+          }
         }
         break;
       case 'ArrowRight':
         if (direction === 'rtl' && e.currentTarget.contains(e.target as Element)) {
           e.stopPropagation();
           onSubmenuClose();
-          ref.current?.focus();
+          if (!isVirtualFocus) {
+            ref.current?.focus();
+          }
         }
         break;
       case 'Escape':
@@ -118,10 +124,28 @@ export function useSubmenuTrigger<T>(props: AriaSubmenuTriggerProps, state: Subm
     }
   };
 
+  let subDialogKeyDown = (e: KeyboardEvent) => {
+    switch (e.key) {
+      case 'Escape':
+        onSubmenuClose();
+        if (!isVirtualFocus) {
+          ref.current?.focus();
+        }
+        return;
+      default:
+        // Ensure events like Tab are still handled by the FocusScope
+        if ('continuePropagation' in e) {
+          e.continuePropagation();
+        }
+
+    }
+  };
+
   let submenuProps = {
     id: overlayId,
     'aria-labelledby': submenuTriggerId,
     submenuLevel: state.submenuLevel,
+    onKeyDown: type === 'dialog' ? subDialogKeyDown : undefined,
     ...(type === 'menu' && {
       onClose: state.closeAll,
       autoFocus: state.focusStrategy ?? undefined,
@@ -186,6 +210,7 @@ export function useSubmenuTrigger<T>(props: AriaSubmenuTriggerProps, state: Subm
     if (!isDisabled && (e.pointerType === 'touch' || e.pointerType === 'mouse')) {
       // For touch or on a desktop device with a small screen open on press up to possible problems with
       // press up happening on the newly opened tray items
+      console.log('2')
       onSubmenuOpen();
     }
   };
@@ -205,7 +230,7 @@ export function useSubmenuTrigger<T>(props: AriaSubmenuTriggerProps, state: Subm
   };
 
   let onBlur = (e) => {
-    if (state.isOpen && parentMenuRef.current?.contains(e.relatedTarget)) {
+    if (state.isOpen && (parentMenuRef.current?.contains(e.relatedTarget) || isVirtualFocus)) {
       onSubmenuClose();
     }
   };
@@ -236,7 +261,8 @@ export function useSubmenuTrigger<T>(props: AriaSubmenuTriggerProps, state: Subm
     submenuProps,
     popoverProps: {
       isNonModal: true,
-      disableFocusManagement: true,
+      // TODO: maybe also include a check that we aren't in a screen reader experience? Kinda gross
+      disableFocusManagement: type === 'menu' && !isVirtualFocus,
       shouldCloseOnInteractOutside
     }
   };

packages/@react-aria/selection/src/useSelectableCollection.ts

@@ -454,10 +454,12 @@ export function useSelectableCollection(options: AriaSelectableCollectionOptions
     resetFocusFirstFlag();
   }, [manager.focusedKey, resetFocusFirstFlag]);
 
-  useEvent(ref, CLEAR_FOCUS_EVENT, !shouldUseVirtualFocus ? undefined : (e) => {
+  useEvent(ref, CLEAR_FOCUS_EVENT, !shouldUseVirtualFocus ? undefined : (e: any) => {
     e.stopPropagation();
     manager.setFocused(false);
-    manager.setFocusedKey(null);
+    if (e.detail?.clearFocusKey) {
+      manager.setFocusedKey(null);
+    }
   });
 
   const autoFocusRef = useRef(autoFocus);

packages/@react-aria/selection/src/useSelectableItem.ts

@@ -373,7 +373,10 @@ export function useSelectableItem(options: SelectableItemOptions): SelectableIte
     ),
     isPressed,
     isSelected: manager.isSelected(key),
-    isFocused: manager.isFocused && manager.focusedKey === key,
+    // TODO: an alternative to this would be to set manager.isFocused to true in useAutocomplete in the input's onFocus when returning back to a parent autocomplete menu
+    // after closing a submenu/subdialog, but that feels more iffy because we'd need to differentiate that from cases where the user
+    // is simply returning to the autocomplete menu from elsewhere
+    isFocused: (shouldUseVirtualFocus || manager.isFocused) && manager.focusedKey === key,
     isDisabled,
     allowsSelection,
     hasAction