progress with supporting virtual focus in table

still a bunch of bugs, see notes

Author: LFDanLu

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

@@ -231,12 +231,15 @@ export function useAutocomplete<T>(props: AriaAutocompleteOptions<T>, state: Aut
         // Backspace shouldn't trigger tag deletion either
         return;
       case 'Tab':
-        // Don't propogate Tab down to the collection, otherwise we will try to focus the collection via useSelectableCollection's Tab handler (aka shift tab logic)
-        // We want FocusScope to handle Tab if one exists (aka sub dialog), so special casepropogate
-        if ('continuePropagation' in e) {
+        // Propagate Tab down to the collection so that tabbing foward will hit our special logic to treat the collection
+        // as a single tab stop. We want FocusScope to handle Shift Tab if one exists (aka sub dialog), so special case propogate
+        // Otherwise, we don't want useSeletableCollection to handle that anyways since focus is actually on an input outside the
+        // wrapped collection and thus the browser can handle that for us
+        if ('continuePropagation' in e && e.shiftKey) {
           e.continuePropagation();
+          return;
         }
-        return;
+        break;
       case 'Home':
       case 'End':
       case 'PageDown':

packages/@react-aria/button/src/useButton.ts

@@ -105,8 +105,15 @@ export function useButton(props: AriaButtonOptions<ElementType>, ref: RefObject<
   if (allowFocusWhenDisabled) {
     focusableProps.tabIndex = isDisabled ? -1 : focusableProps.tabIndex;
   }
+
+  // TODO: need to make the button (and really any elements within a virtual focus collection) non tabbable
+  // if we don't do this then real focus can move inside the collection and breaks the virtual focus logic
+  // (can get cases where multiple things get focus visible styles)
+  // Ideally useSelectableCollection can do the same thing where it coerces focus past the entire collection
+  // focusableProps.tabIndex = -1;
+
   let buttonProps = mergeProps(focusableProps, pressProps, filterDOMProps(props, {labelable: true}));
-  // TODO: will need to support virtual focus on the button
+
   return {
     isPressed, // Used to indicate press state for visual
     buttonProps: mergeProps(additionalProps, buttonProps, {

packages/@react-aria/focus/src/useFocusRing.ts

@@ -50,12 +50,14 @@ export function useFocusRing(props: AriaFocusRingProps = {}): FocusRingAria {
   let updateState = useCallback(() => setFocusVisible(state.current.isFocused && state.current.isFocusVisible), []);
 
   let onFocusChange = useCallback(isFocused => {
+    // console.log('calling focus change', isFocused, state.current.isFocusVisible)
     state.current.isFocused = isFocused;
     setFocused(isFocused);
     updateState();
   }, [updateState]);
 
   useFocusVisibleListener((isFocusVisible) => {
+    // console.log('calling focus visible listern', isFocusVisible)
     state.current.isFocusVisible = isFocusVisible;
     updateState();
   }, [], {isTextInput});

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

@@ -64,7 +64,11 @@ export interface GridProps extends DOMProps, AriaLabelingProps {
    */
   escapeKeyBehavior?: 'clearSelection' | 'none',
   /** Whether selection should occur on press up instead of press down. */
-  shouldSelectOnPressUp?: boolean
+  shouldSelectOnPressUp?: boolean,
+  /**
+   * Whether the table cells should use virtual focus instead of being focused directly.
+   */
+  shouldUseVirtualFocus?: boolean
 }
 
 export interface GridAria {
@@ -90,7 +94,8 @@ export function useGrid<T>(props: GridProps, state: GridState<T, GridCollection<
     onRowAction,
     onCellAction,
     escapeKeyBehavior = 'clearSelection',
-    shouldSelectOnPressUp
+    shouldSelectOnPressUp,
+    shouldUseVirtualFocus
   } = props;
   let {selectionManager: manager} = state;
 
@@ -120,11 +125,12 @@ export function useGrid<T>(props: GridProps, state: GridState<T, GridCollection<
     isVirtualized,
     scrollRef,
     disallowTypeAhead,
-    escapeKeyBehavior
+    escapeKeyBehavior,
+    shouldUseVirtualFocus
   });
 
   let id = useId(props.id);
-  gridMap.set(state, {keyboardDelegate: delegate, actions: {onRowAction, onCellAction}, shouldSelectOnPressUp});
+  gridMap.set(state, {keyboardDelegate: delegate, actions: {onRowAction, onCellAction}, shouldSelectOnPressUp, shouldUseVirtualFocus});
 
   let descriptionProps = useHighlightSelectionDescription({
     selectionManager: manager,
@@ -169,6 +175,7 @@ export function useGrid<T>(props: GridProps, state: GridState<T, GridCollection<
       'aria-multiselectable': manager.selectionMode === 'multiple' ? 'true' : undefined
     },
     state.isKeyboardNavigationDisabled ? navDisabledHandlers : collectionProps,
+    // TODO: may need to change this tabIndex here for virtual focus
     // If collection is empty, make sure the grid is tabbable unless there is a child tabbable element.
     (state.collection.size === 0 && {tabIndex: hasTabbableChild ? -1 : 0}) || undefined,
     descriptionProps

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

@@ -12,7 +12,7 @@
 
 import {DOMAttributes, FocusableElement, Key, RefObject} from '@react-types/shared';
 import {focusSafely, isFocusVisible} from '@react-aria/interactions';
-import {getFocusableTreeWalker} from '@react-aria/focus';
+import {getFocusableTreeWalker, getVirtuallyFocusedElement, moveVirtualFocus} from '@react-aria/focus';
 import {getScrollParent, mergeProps, scrollIntoViewport} from '@react-aria/utils';
 import {GridCollection, GridNode} from '@react-types/grid';
 import {gridMap} from './utils';
@@ -62,37 +62,40 @@ export function useGridCell<T, C extends GridCollection<T>>(props: GridCellProps
   } = props;
 
   let {direction} = useLocale();
-  let {keyboardDelegate, actions: {onCellAction}} = gridMap.get(state)!;
+  let {keyboardDelegate, actions: {onCellAction}, shouldUseVirtualFocus} = gridMap.get(state)!;
 
   // We need to track the key of the item at the time it was last focused so that we force
   // focus to go to the item when the DOM node is reused for a different item in a virtualizer.
   let keyWhenFocused = useRef<Key | null>(null);
 
   // Handles focusing the cell. If there is a focusable child,
   // it is focused, otherwise the cell itself is focused.
+  // TODO: throughly test these changes
   let focus = () => {
     if (ref.current) {
       let treeWalker = getFocusableTreeWalker(ref.current);
+      let activeElement = shouldUseVirtualFocus ? getVirtuallyFocusedElement(document) : document.activeElement;
       if (focusMode === 'child') {
         // If focus is already on a focusable child within the cell, early return so we don't shift focus
-        if (ref.current.contains(document.activeElement) && ref.current !== document.activeElement) {
+        if (ref.current.contains(activeElement) && ref.current !== activeElement) {
           return;
         }
 
         let focusable = state.selectionManager.childFocusStrategy === 'last'
           ? last(treeWalker)
           : treeWalker.firstChild() as FocusableElement;
+
         if (focusable) {
-          focusSafely(focusable);
+          focusElement(focusable);
           return;
         }
       }
 
       if (
         (keyWhenFocused.current != null && node.key !== keyWhenFocused.current) ||
-        !ref.current.contains(document.activeElement)
+        !ref.current.contains(activeElement)
       ) {
-        focusSafely(ref.current);
+        focusElement(ref.current);
       }
     }
   };
@@ -105,16 +108,25 @@ export function useGridCell<T, C extends GridCollection<T>>(props: GridCellProps
     focus,
     shouldSelectOnPressUp,
     onAction: onCellAction ? () => onCellAction(node.key) : onAction,
-    isDisabled: state.collection.size === 0
+    isDisabled: state.collection.size === 0,
+    shouldUseVirtualFocus
   });
 
+  let focusElement = (element: FocusableElement) => {
+    if (!shouldUseVirtualFocus) {
+      focusSafely(element);
+    } else {
+      moveVirtualFocus(element);
+    }
+  };
+
   let onKeyDownCapture = (e: ReactKeyboardEvent) => {
     if (!e.currentTarget.contains(e.target as Element) || state.isKeyboardNavigationDisabled || !ref.current || !document.activeElement) {
       return;
     }
 
     let walker = getFocusableTreeWalker(ref.current);
-    walker.currentNode = document.activeElement;
+    walker.currentNode = shouldUseVirtualFocus ? getVirtuallyFocusedElement(document) as FocusableElement : document.activeElement;
 
     switch (e.key) {
       case 'ArrowLeft': {
@@ -131,7 +143,7 @@ export function useGridCell<T, C extends GridCollection<T>>(props: GridCellProps
         e.preventDefault();
         e.stopPropagation();
         if (focusable) {
-          focusSafely(focusable);
+          focusElement(focusable);
           scrollIntoViewport(focusable, {containingElement: getScrollParent(ref.current)});
         } else {
           // If there is no next focusable child, then move to the next cell to the left of this one.
@@ -150,16 +162,17 @@ export function useGridCell<T, C extends GridCollection<T>>(props: GridCellProps
             break;
           }
 
+          // TODO: may need the same handling as in GridList with the currentNode check
           if (focusMode === 'cell' && direction === 'rtl') {
-            focusSafely(ref.current);
+            focusElement(ref.current);
             scrollIntoViewport(ref.current, {containingElement: getScrollParent(ref.current)});
           } else {
             walker.currentNode = ref.current;
             focusable = direction === 'rtl'
               ? walker.firstChild() as FocusableElement
               : last(walker);
             if (focusable) {
-              focusSafely(focusable);
+              focusElement(focusable);
               scrollIntoViewport(focusable, {containingElement: getScrollParent(ref.current)});
             }
           }
@@ -178,7 +191,7 @@ export function useGridCell<T, C extends GridCollection<T>>(props: GridCellProps
         e.preventDefault();
         e.stopPropagation();
         if (focusable) {
-          focusSafely(focusable);
+          focusElement(focusable);
           scrollIntoViewport(focusable, {containingElement: getScrollParent(ref.current)});
         } else {
           let next = keyboardDelegate.getKeyRightOf?.(node.key);
@@ -192,16 +205,17 @@ export function useGridCell<T, C extends GridCollection<T>>(props: GridCellProps
             break;
           }
 
+           // TODO: may need the same handling as in GridList with the currentNode check
           if (focusMode === 'cell' && direction === 'ltr') {
-            focusSafely(ref.current);
+            focusElement(ref.current);
             scrollIntoViewport(ref.current, {containingElement: getScrollParent(ref.current)});
           } else {
             walker.currentNode = ref.current;
             focusable = direction === 'rtl'
               ? last(walker)
               : walker.firstChild() as FocusableElement;
             if (focusable) {
-              focusSafely(focusable);
+              focusElement(focusable);
               scrollIntoViewport(focusable, {containingElement: getScrollParent(ref.current)});
             }
           }

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

@@ -52,7 +52,7 @@ export function useGridRow<T, C extends GridCollection<T>, S extends GridState<T
     onAction
   } = props;
 
-  let {actions, shouldSelectOnPressUp: gridShouldSelectOnPressUp} = gridMap.get(state)!;
+  let {actions, shouldSelectOnPressUp: gridShouldSelectOnPressUp, shouldUseVirtualFocus} = gridMap.get(state)!;
   let onRowAction = actions.onRowAction ? () => actions.onRowAction?.(node.key) : onAction;
   let {itemProps, ...states} = useSelectableItem({
     selectionManager: state.selectionManager,
@@ -61,7 +61,8 @@ export function useGridRow<T, C extends GridCollection<T>, S extends GridState<T
     isVirtualized,
     shouldSelectOnPressUp: gridShouldSelectOnPressUp || shouldSelectOnPressUp,
     onAction: onRowAction || node?.props?.onAction ? chain(node?.props?.onAction, onRowAction) : undefined,
-    isDisabled: state.collection.size === 0
+    isDisabled: state.collection.size === 0,
+    shouldUseVirtualFocus
   });
 
   let isSelected = state.selectionManager.isSelected(node.key);

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

@@ -20,7 +20,8 @@ interface GridMapShared {
     onRowAction?: (key: Key) => void,
     onCellAction?: (key: Key) => void
   },
-  shouldSelectOnPressUp?: boolean
+  shouldSelectOnPressUp?: boolean,
+  shouldUseVirtualFocus?: boolean
 }
 
 // Used to share:

packages/@react-aria/gridlist/src/useGridListItem.ts

@@ -73,15 +73,17 @@ export function useGridListItem<T>(props: AriaGridListItemOptions, state: ListSt
   // We need to track the key of the item at the time it was last focused so that we force
   // focus to go to the item when the DOM node is reused for a different item in a virtualizer.
   let keyWhenFocused = useRef<Key | null>(null);
+  // TODO: need to update this to handle virtual focus
   let focus = () => {
     // Don't shift focus to the row if the active element is a element within the row already
     // (e.g. clicking on a row button)
+    let activeElement = shouldUseVirtualFocus ? getVirtuallyFocusedElement(document) : document.activeElement;
     if (
       ref.current !== null &&
       ((keyWhenFocused.current != null && node.key !== keyWhenFocused.current) ||
-      !ref.current?.contains(document.activeElement))
+      !ref.current?.contains(activeElement))
     ) {
-      focusSafely(ref.current);
+      focusElement(ref.current);
     }
   };
 

packages/@react-aria/interactions/src/focusSafely.ts

@@ -31,6 +31,7 @@ export function focusSafely(element: FocusableElement): void {
   // from off the screen.
   const ownerDocument = getOwnerDocument(element);
   const activeElement = getActiveElement(ownerDocument);
+
   if (getInteractionModality() === 'virtual') {
     let lastFocusedElement = activeElement;
     runAfterTransition(() => {

packages/@react-aria/interactions/src/useFocus.ts

@@ -18,6 +18,8 @@
 import {DOMAttributes, FocusableElement, FocusEvents} from '@react-types/shared';
 import {FocusEvent, useCallback} from 'react';
 import {getActiveElement, getEventTarget, getOwnerDocument} from '@react-aria/utils';
+// TODO: circular dependency
+import {getVirtuallyFocusedElement} from '../../focus';
 import {useSyntheticBlurEvent} from './utils';
 
 export interface FocusProps<Target = FocusableElement> extends FocusEvents<Target> {
@@ -63,9 +65,18 @@ export function useFocus<Target extends FocusableElement = FocusableElement>(pro
     // Double check that document.activeElement actually matches e.target in case a previously chained
     // focus handler already moved focus somewhere else.
 
+    // if (e.target.nodeName === 'BUTTON') {
+    //   console.log('focusing specific node', e.target)
+    //   console.trace()
+    // }
     const ownerDocument = getOwnerDocument(e.target);
     const activeElement = ownerDocument ? getActiveElement(ownerDocument) : getActiveElement();
-    if (e.target === e.currentTarget && activeElement === getEventTarget(e.nativeEvent)) {
+
+    // TODO the below check on getVirtuallyFocusedElement isn't accurate because at this point when focus is called, the virtually focused element
+    // that we detect is still the old one and thus we cant really distinguish if the element recieveing focus here will be the newly virtually focused element
+    // or if we are outside the component using virtual focus and focus is landing on something else? Maybe that doesn't matter since if we are landing on
+    // an element with real focus outside the component using virtual focus, then activeElement === getEventTarget(e.nativeEvent) resolves as true
+    if (e.target === e.currentTarget && (activeElement === getEventTarget(e.nativeEvent) || getVirtuallyFocusedElement(ownerDocument))) {
       if (onFocusProp) {
         onFocusProp(e);
       }