fix: focus the closest visible row when the focused Tree/Table item is collapsed

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

@@ -13,6 +13,7 @@
 import {
   ChildrenOrFunction,
   ClassNameOrFunction,
+  closestVisibleKey,
   ContextValue,
   dom,
   Provider,
@@ -320,26 +321,6 @@ function findKeyForRoute(collection: Collection<Node<unknown>>, route: string):
   return null;
 }
 
-// Walks up from `key` to the closest ancestor that is actually rendered (all ancestors expanded).
-// Returns `key` unchanged when already visible. A collapsed ancestor hides everything beneath it,
-// so the highest collapsed ancestor is the closest visible row.
-function closestVisibleKey(
-  collection: Collection<Node<unknown>>,
-  expandedKeys: Set<Key>,
-  key: Key
-): Key {
-  let target = key;
-  let node = collection.getItem(key);
-  while (node?.parentKey != null) {
-    let parent = collection.getItem(node.parentKey);
-    if (parent?.type === 'item' && !expandedKeys.has(node.parentKey)) {
-      target = node.parentKey;
-    }
-    node = parent;
-  }
-  return target;
-}
-
 // Moves the tree's focused key to the item matching selectedRoute. Runs when the route or the
 // collection changes; the shared syncedRouteRef dedupes across items so it fires once per change.
 function useRouteFocusSync({state}: {state: TreeState<unknown>}): void {

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

@@ -18,6 +18,7 @@ import {ButtonContext} from './Button';
 import {CheckboxContext, CheckboxFieldContext} from './Checkbox';
 import {
   ClassNameOrFunction,
+  closestVisibleKey,
   ContextValue,
   DEFAULT_SLOT,
   dom,
@@ -850,6 +851,22 @@ function TableInner({props, forwardedRef: ref, selectionState, collection}: Tabl
     isRootDropTarget = dropState.isDropTarget({type: 'root'});
   }
 
+  // Dragged rows keep focus so it can follow them when they are dropped.
+  useEffect(() => {
+    let node =
+      selectionState.focusedKey != null ? collection.getItem(selectionState.focusedKey) : null;
+    // A cell is visible whenever its row is.
+    let rowKey = node?.type === 'cell' ? node.parentKey : node?.key;
+    if (rowKey == null || dragState?.isDragging(rowKey)) {
+      return;
+    }
+    let visibleKey = closestVisibleKey(collection, expandedKeys, rowKey);
+    if (visibleKey !== rowKey) {
+      selectionState.setFocusedKey(visibleKey);
+    }
+    // eslint-disable-next-line react-hooks/exhaustive-deps
+  }, [collection, expandedKeys, selectionState.focusedKey]);
+
   let {focusProps, isFocused, isFocusVisible} = useFocusRing();
   let renderProps = useRenderProps({
     ...props,

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

@@ -24,6 +24,7 @@ import {CheckboxContext, CheckboxFieldContext} from './Checkbox';
 import {
   ChildrenOrFunction,
   ClassNameOrFunction,
+  closestVisibleKey,
   ContextValue,
   DEFAULT_SLOT,
   dom,
@@ -569,6 +570,19 @@ function TreeInner<T>({props, collection, treeRef: ref}: TreeInnerProps<T>) {
     isRootDropTarget = dropState.isDropTarget({type: 'root'});
   }
 
+  // Dragged items keep focus so it can follow them when they are dropped.
+  useEffect(() => {
+    let focusedKey = state.selectionManager.focusedKey;
+    if (focusedKey == null || dragState?.isDragging(focusedKey)) {
+      return;
+    }
+    let visibleKey = closestVisibleKey(state.collection, expandedKeys, focusedKey);
+    if (visibleKey !== focusedKey) {
+      state.selectionManager.setFocusedKey(visibleKey);
+    }
+    // eslint-disable-next-line react-hooks/exhaustive-deps
+  }, [state.collection, expandedKeys, state.selectionManager.focusedKey]);
+
   let isTreeDraggable = !!(hasDragHooks && !dragState?.isDisabled);
 
   let {focusProps, isFocused, isFocusVisible} = useFocusRing();

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

@@ -10,7 +10,14 @@
  * governing permissions and limitations under the License.
  */
 
-import {AriaLabelingProps, RefObject, DOMProps as SharedDOMProps} from '@react-types/shared';
+import {
+  AriaLabelingProps,
+  Collection,
+  Key,
+  Node,
+  RefObject,
+  DOMProps as SharedDOMProps
+} from '@react-types/shared';
 import {mergeProps} from 'react-aria/mergeProps';
 import {mergeRefs} from 'react-aria/mergeRefs';
 import React, {
@@ -404,6 +411,24 @@ export function removeDataAttributes<T>(props: T): T {
   return filteredProps;
 }
 
+// A collapsed ancestor hides everything below it, so the highest one is the closest visible row.
+export function closestVisibleKey<T>(
+  collection: Collection<Node<T>>,
+  expandedKeys: Set<Key>,
+  key: Key
+): Key {
+  let target = key;
+  let node = collection.getItem(key);
+  while (node?.parentKey != null) {
+    let parent = collection.getItem(node.parentKey);
+    if (parent?.type === 'item' && !expandedKeys.has(node.parentKey)) {
+      target = node.parentKey;
+    }
+    node = parent;
+  }
+  return target;
+}
+
 // Override base type to change the default.
 export interface RACValidation {
   /**