fix: improve keyboard navigation when Table has expandable rows (#9842)

* fix: speed up keyboard navigation with Table expandable rows

* fix tests

* copy pasta

* cleanup

* fix lint?

Author: yihuiliao

packages/@adobe/react-spectrum/test/table/TreeGridTable.test.tsx

@@ -692,42 +692,78 @@ describe('TableView with expandable rows', function () {
         expect(document.activeElement).toBe(getCell(treegrid, 'Row 1, Lvl 2, Foo'));
       });
 
-      it('should properly wrap focus with ArrowRight (RTL)', function () {
+      it('should move focus to parent if already collapsed (RTL)', function () {
         let treegrid = render(<ManyRowsExpandableTable />, undefined, 'ar-AE');
-        let row = treegrid.getAllByRole('row')[2];
+        let row1 = treegrid.getAllByRole('row')[1];
+        let row2 = treegrid.getAllByRole('row')[2];
         focusCell(treegrid, 'Row 1, Lvl 2, Foo');
         moveFocus('ArrowRight');
+        expect(document.activeElement).toBe(row2);
+        expect(row2).toContainElement(getCell(treegrid, 'Row 1, Lvl 2, Foo'));
+        // Focus doesn't move because Arrow Right on the row will collapse it here
+        moveFocus('ArrowRight');
+        expect(document.activeElement).toBe(row2);
+        moveFocus('ArrowRight');
+        expect(document.activeElement).toBe(row1);
+        moveFocus('ArrowRight');
+        expect(document.activeElement).toBe(row1);
+        moveFocus('ArrowRight');
+      });
+
+      it('should properly wrap focus with ArrowRight when top level row is collapsed (RTL)', function () {
+        let treegrid = render(<ManyRowsExpandableTable />, undefined, 'ar-AE');
+        let row = treegrid.getAllByRole('row')[4];
+        focusCell(treegrid, 'Row 2, Lvl 1, Foo');
+        moveFocus('ArrowRight');
         expect(document.activeElement).toBe(row);
-        expect(row).toContainElement(getCell(treegrid, 'Row 1, Lvl 2, Foo'));
-        // Focus doesn't move because Arrow Right on the row will collapse it
+        expect(row).toContainElement(getCell(treegrid, 'Row 2, Lvl 1, Foo'));
+        // Focus doesn't move because Arrow Right on the row will collapse it here
         moveFocus('ArrowRight');
         expect(document.activeElement).toBe(row);
         moveFocus('ArrowRight');
-        expect(document.activeElement).toBe(getCell(treegrid, 'Row 1, Lvl 2, Baz'));
+        expect(document.activeElement).toBe(getCell(treegrid, 'Row 2, Lvl 1, Baz'));
         moveFocus('ArrowRight');
-        expect(document.activeElement).toBe(getCell(treegrid, 'Row 1, Lvl 2, Bar'));
+        expect(document.activeElement).toBe(getCell(treegrid, 'Row 2, Lvl 1, Bar'));
         moveFocus('ArrowRight');
-        expect(document.activeElement).toBe(getCell(treegrid, 'Row 1, Lvl 2, Foo'));
+        expect(document.activeElement).toBe(getCell(treegrid, 'Row 2, Lvl 1, Foo'));
       });
     });
 
     describe('ArrowLeft', function () {
-      it('should properly wrap focus with ArrowLeft', function () {
+      it('should move focus to parent if already collapsed', function () {
         let treegrid = render(<ManyRowsExpandableTable />);
-        let row = treegrid.getAllByRole('row')[2];
+        let row1 = treegrid.getAllByRole('row')[1];
+        let row2 = treegrid.getAllByRole('row')[2];
         focusCell(treegrid, 'Row 1, Lvl 2, Foo');
         moveFocus('ArrowLeft');
+        expect(document.activeElement).toBe(row2);
+        expect(row2).toContainElement(getCell(treegrid, 'Row 1, Lvl 2, Foo'));
+        // Focus doesn't move because Arrow Left on the row will collapse it here
+        moveFocus('ArrowLeft');
+        expect(document.activeElement).toBe(row2);
+        moveFocus('ArrowLeft');
+        expect(document.activeElement).toBe(row1);
+        moveFocus('ArrowLeft');
+        expect(document.activeElement).toBe(row1);
+        moveFocus('ArrowLeft');
+      });
+
+      it('should properly wrap focus with ArrowLeft when top level row is collapsed', function () {
+        let treegrid = render(<ManyRowsExpandableTable />);
+        let row = treegrid.getAllByRole('row')[4];
+        focusCell(treegrid, 'Row 2, Lvl 1, Foo');
+        moveFocus('ArrowLeft');
         expect(document.activeElement).toBe(row);
-        expect(row).toContainElement(getCell(treegrid, 'Row 1, Lvl 2, Foo'));
+        expect(row).toContainElement(getCell(treegrid, 'Row 2, Lvl 1, Foo'));
         // Focus doesn't move because Arrow Left on the row will collapse it here
         moveFocus('ArrowLeft');
         expect(document.activeElement).toBe(row);
         moveFocus('ArrowLeft');
-        expect(document.activeElement).toBe(getCell(treegrid, 'Row 1, Lvl 2, Baz'));
+        expect(document.activeElement).toBe(getCell(treegrid, 'Row 2, Lvl 1, Baz'));
         moveFocus('ArrowLeft');
-        expect(document.activeElement).toBe(getCell(treegrid, 'Row 1, Lvl 2, Bar'));
+        expect(document.activeElement).toBe(getCell(treegrid, 'Row 2, Lvl 1, Bar'));
         moveFocus('ArrowLeft');
-        expect(document.activeElement).toBe(getCell(treegrid, 'Row 1, Lvl 2, Foo'));
+        expect(document.activeElement).toBe(getCell(treegrid, 'Row 2, Lvl 1, Foo'));
       });
 
       it('should properly wrap focus with ArrowLeft (RTL)', function () {

packages/react-aria/src/table/useTableRow.ts

@@ -84,9 +84,20 @@ export function useTableRow<T>(props: GridRowProps<T>, state: TableState<T> | Tr
           if ((e.key === EXPANSION_KEYS['expand'][direction]) && state.selectionManager.focusedKey === treeNode.key && hasChildRows && state.expandedKeys !== 'all' && !state.expandedKeys.has(treeNode.key)) {
             state.toggleKey(treeNode.key);
             e.stopPropagation();
-          } else if ((e.key === EXPANSION_KEYS['collapse'][direction]) && state.selectionManager.focusedKey === treeNode.key && hasChildRows && (state.expandedKeys === 'all' || state.expandedKeys.has(treeNode.key))) {
-            state.toggleKey(treeNode.key);
-            e.stopPropagation();
+          } else if ((e.key === EXPANSION_KEYS['collapse'][direction]) && state.selectionManager.focusedKey === treeNode.key) {
+            if (state.expandedKeys !== 'all') {
+              if (hasChildRows && state.expandedKeys.has(treeNode.key)) {
+                state.toggleKey(treeNode.key);
+                e.stopPropagation();
+              } else if (!state.expandedKeys.has(treeNode.key) && treeNode.parentKey && treeNode.level > 0) {
+                // Item is a leaf or already collapsed, move focus to parent
+                state.selectionManager.setFocusedKey(treeNode.parentKey);
+                e.stopPropagation();
+              }
+            } else if (state.expandedKeys === 'all') {
+              state.toggleKey(treeNode.key);
+              e.stopPropagation();
+            }
           }
         },
         'aria-expanded': hasChildRows ? state.expandedKeys === 'all' || state.expandedKeys.has(node.key) : undefined,