fix: focus restoring in collections

packages/@react-spectrum/listbox/test/ListBox.test.js

@@ -1042,54 +1042,67 @@ describe('ListBox', function () {
       act(() => jest.runAllTimers());
       let listbox = tree.getByRole('listbox');
       let options = within(listbox).getAllByRole('option');
-      act(() => {options[2].focus();});
+      expect(options.length).toBe(6);
+      // Go to *Snake* and select it
+      await user.tab();
+      await user.keyboard('{ArrowDown}');
+      await user.keyboard('{ArrowDown}');
       expect(options[2]).toHaveAttribute('aria-posinset', '3');
       expect(options[2]).toHaveAttribute('aria-setsize', '6');
       expect(document.activeElement).toBe(options[2]);
-      fireEvent.keyDown(document.activeElement, {key: ' ', code: 32, charCode: 32});
+      await user.keyboard('{Enter}');
       expect(document.activeElement).toHaveAttribute('aria-selected', 'true');
+
+      // Remove *Snake*
       let removeButton = tree.getByRole('button');
       expect(removeButton).toBeInTheDocument();
-      act(() => {removeButton.focus();});
-      expect(document.activeElement).toBe(removeButton);
       await user.click(removeButton);
       act(() => jest.runAllTimers());
       let confirmationDialog = tree.getByRole('alertdialog');
       expect(document.activeElement).toBe(confirmationDialog);
       let confirmationDialogButton = within(confirmationDialog).getByRole('button');
-      expect(confirmationDialogButton).toBeInTheDocument();
       await user.click(confirmationDialogButton);
       act(() => jest.runAllTimers());
       options = within(listbox).getAllByRole('option');
       expect(options.length).toBe(5);
       act(() => jest.runAllTimers());
       expect(confirmationDialog).not.toBeInTheDocument();
+
+      // Dialog returns focus to the ListBox which forwards it to the options
       expect(document.activeElement).toBe(options[2]);
       expect(options[2]).toHaveAttribute('aria-posinset', '3');
       expect(options[2]).toHaveAttribute('aria-setsize', '5');
-      act(() => {options[1].focus();});
-      fireEvent.keyDown(document.activeElement, {key: ' ', code: 32, charCode: 32});
+
+      // Select option
+      await user.keyboard('{Enter}');
       expect(document.activeElement).toHaveAttribute('aria-selected', 'true');
-      act(() => {options[0].focus();});
-      fireEvent.keyDown(document.activeElement, {key: ' ', code: 32, charCode: 32});
+
+      // Go to option 0, *Aardvark* and select it too
+      await user.keyboard('{ArrowUp}');
+      await user.keyboard('{ArrowUp}');
+      await user.keyboard('{ArrowUp}');
+      await user.keyboard('{ArrowUp}');
+      expect(document.activeElement).toBe(options[0]);
+      await user.keyboard('{Enter}');
       expect(document.activeElement).toHaveAttribute('aria-selected', 'true');
-      act(() => {options[0].focus();});
+
+      // Remove the two selected items
       removeButton = tree.getByRole('button');
-      expect(removeButton).toBeInTheDocument();
-      act(() => {removeButton.focus();});
+      await user.tab({shift: true});
       expect(document.activeElement).toBe(removeButton);
       await user.click(removeButton);
       act(() => jest.runAllTimers());
       confirmationDialog = tree.getByRole('alertdialog');
       expect(document.activeElement).toBe(confirmationDialog);
       confirmationDialogButton = within(confirmationDialog).getByRole('button');
-      expect(confirmationDialogButton).toBeInTheDocument();
       await user.click(confirmationDialogButton);
       act(() => jest.runAllTimers());
       options = within(listbox).getAllByRole('option');
       expect(options.length).toBe(3);
       act(() => jest.runAllTimers());
       expect(confirmationDialog).not.toBeInTheDocument();
+
+      // Dialog returns focus to the ListBox which forwards it to the options
       expect(document.activeElement).toBe(options[0]);
       expect(options[0]).toHaveAttribute('aria-posinset', '1');
       expect(options[0]).toHaveAttribute('aria-setsize', '3');

packages/@react-stately/list/src/useListState.ts

@@ -89,47 +89,34 @@ function useFocusedKeyReset<T>(collection: Collection<Node<T>>, selectionManager
   const cachedCollection = useRef<Collection<Node<T>> | null>(null);
   useEffect(() => {
     if (selectionManager.focusedKey != null && !collection.getItem(selectionManager.focusedKey) && cachedCollection.current) {
-      const startItem = cachedCollection.current.getItem(selectionManager.focusedKey);
-      const cachedItemNodes = [...cachedCollection.current.getKeys()].map(
-        key => {
-          const itemNode = cachedCollection.current!.getItem(key);
-          return itemNode?.type === 'item' ? itemNode : null;
-        }
-      ).filter(node => node !== null);
-      const itemNodes = [...collection.getKeys()].map(
-        key => {
-          const itemNode = collection.getItem(key);
-          return itemNode?.type === 'item' ? itemNode : null;
-        }
-      ).filter(node => node !== null);
-      const diff: number = (cachedItemNodes?.length ?? 0) - (itemNodes?.length ?? 0);
-      let index = Math.min(
-        (
-          diff > 1 ?
-          Math.max((startItem?.index ?? 0) - diff + 1, 0) :
-          startItem?.index ?? 0
-        ),
-        (itemNodes?.length ?? 0) - 1);
-      let newNode: Node<T> | null = null;
-      let isReverseSearching = false;
-      while (index >= 0) {
-        if (!selectionManager.isDisabled(itemNodes[index].key)) {
-          newNode = itemNodes[index];
+      // Walk forward in the old collection to find the next key that still exists in the new collection.
+      let key = cachedCollection.current.getKeyAfter(selectionManager.focusedKey);
+      let nextFocusedKey: Key | null = null;
+      while (key != null) {
+        let node = collection.getItem(key);
+        if (node && node.type === 'item' && !selectionManager.isDisabled(key)) {
+          nextFocusedKey = key;
           break;
         }
-        // Find next, not disabled item.
-        if (index < itemNodes.length - 1 && !isReverseSearching) {
-          index++;
-        // Otherwise, find previous, not disabled item.
-        } else {
-          isReverseSearching = true;
-          if (index > (startItem?.index ?? 0)) {
-            index = (startItem?.index ?? 0);
+
+        key = cachedCollection.current.getKeyAfter(key);
+      }
+
+      // If no such key exists, walk backward.
+      if (nextFocusedKey == null) {
+        key = cachedCollection.current.getKeyBefore(selectionManager.focusedKey);
+        while (key != null) {
+          let node = collection.getItem(key);
+          if (node && node.type === 'item' && !selectionManager.isDisabled(key)) {
+            nextFocusedKey = key;
+            break;
           }
-          index--;
+
+          key = cachedCollection.current.getKeyBefore(key);
         }
       }
-      selectionManager.setFocusedKey(newNode ? newNode.key : null);
+
+      selectionManager.setFocusedKey(nextFocusedKey);
     }
     cachedCollection.current = collection;
   }, [collection, selectionManager]);

packages/react-aria-components/test/TagGroup.test.js

@@ -12,7 +12,7 @@
 
 import {act, fireEvent, mockClickDefault, pointerMap, render} from '@react-spectrum/test-utils-internal';
 import {Button, Label, RouterProvider, Tag, TagGroup, TagList, Text, Tooltip, TooltipTrigger} from '../';
-import React from 'react';
+import React, {useRef} from 'react';
 import {useListData} from '@react-stately/data';
 import {User} from '@react-aria/test-utils';
 import userEvent from '@testing-library/user-event';
@@ -572,6 +572,77 @@ describe('TagGroup', () => {
     expect(onRemove).toHaveBeenLastCalledWith(new Set(['dog']));
   });
 
+  it('should maintain item order when adding new items', async () => {
+    function MyTag(props) {
+      return (
+        <Tag
+          {...props}
+          style={({isSelected}) => ({border: '1px solid gray', borderRadius: 4, padding: '0 4px', background: isSelected ? 'black' : '', color: isSelected ? 'white' : '', cursor: props.href ? 'pointer' : 'default'})} />
+      );
+    }
+    function Example() {
+      const list = useListData({
+        initialItems: []
+      });
+
+      const nextIdRef = useRef(0);
+
+      const insertItem = () => {
+        const id = nextIdRef.current++;
+        list.insert(0, {
+          id,
+          label: `Item ${id + 1}`
+        });
+      };
+
+      return (
+        <div>
+          <Button onPress={insertItem}>Insert item</Button>
+          <TagGroup onRemove={keys => list.remove(...keys)}>
+            <Label>Categories</Label>
+            <TagList style={{display: 'flex', gap: 4}} items={list.items} renderEmptyState={() => 'No categories.'}>
+              {item => <MyTag textValue={item.label}>{item.label}<Button slot="remove">X</Button></MyTag>}
+            </TagList>
+          </TagGroup>
+        </div>
+      );
+    }
+    let {getAllByRole, queryAllByRole, getByRole} = render(<Example />);
+    let addButton = getAllByRole('button')[0];
+    let tagGroup = getByRole('group');
+
+    await user.click(addButton);
+    await user.click(addButton);
+    await user.click(addButton);
+    await user.click(addButton);
+    act(() => jest.runAllTimers());
+    let items = getAllByRole('row');
+    expect(items).toHaveLength(4);
+    expect(items[0]).toHaveTextContent('Item 4');
+
+    await user.tab();
+
+    await user.keyboard('{Delete}');
+    items = getAllByRole('row');
+    expect(items).toHaveLength(3);
+    expect(items[0]).toHaveTextContent('Item 3');
+
+    await user.keyboard('{Delete}');
+    items = getAllByRole('row');
+    expect(items).toHaveLength(2);
+    expect(items[0]).toHaveTextContent('Item 2');
+
+    await user.keyboard('{Delete}');
+    items = getAllByRole('row');
+    expect(items).toHaveLength(1);
+    expect(items[0]).toHaveTextContent('Item 1');
+
+    await user.keyboard('{Delete}');
+    let noItems = queryAllByRole('row');
+    expect(noItems).toHaveLength(0);
+    expect(document.activeElement).toBe(tagGroup);
+  });
+
   it('should support onAction', async () => {
     let onAction = jest.fn();
     let {getAllByRole} = renderTagGroup({onAction, selectionMode: 'none'});