fix(SelectionManager): support full collection for select all functionality

packages/@react-stately/selection/src/SelectionManager.ts

@@ -28,7 +28,8 @@ import {Selection} from './Selection';
 
 interface SelectionManagerOptions {
   allowsCellSelection?: boolean,
-  layoutDelegate?: LayoutDelegate
+  layoutDelegate?: LayoutDelegate,
+  fullCollection?: Collection<Node<unknown>>
 }
 
 /**
@@ -40,13 +41,15 @@ export class SelectionManager implements MultipleSelectionManager {
   private allowsCellSelection: boolean;
   private _isSelectAll: boolean | null;
   private layoutDelegate: LayoutDelegate | null;
+  private fullCollection: Collection<Node<unknown>> | null;
 
   constructor(collection: Collection<Node<unknown>>, state: MultipleSelectionState, options?: SelectionManagerOptions) {
     this.collection = collection;
     this.state = state;
     this.allowsCellSelection = options?.allowsCellSelection ?? false;
     this._isSelectAll = null;
     this.layoutDelegate = options?.layoutDelegate || null;
+    this.fullCollection = options?.fullCollection || null;
   }
 
   /**
@@ -388,26 +391,29 @@ export class SelectionManager implements MultipleSelectionManager {
   }
 
   private getSelectAllKeys() {
+    // Use the full (unfiltered) collection when available so that materializing
+    // the 'all' selection includes items that are currently filtered out (e.g. by Autocomplete).
+    let collection = this.fullCollection ?? this.collection;
     let keys: Key[] = [];
     let addKeys = (key: Key | null) => {
       while (key != null) {
-        if (this.canSelectItem(key)) {
-          let item = this.collection.getItem(key);
+        if (this.canSelectItemIn(key, collection)) {
+          let item = collection.getItem(key);
           if (item?.type === 'item') {
             keys.push(key);
           }
 
           // Add child keys. If cell selection is allowed, then include item children too.
           if (item?.hasChildNodes && (this.allowsCellSelection || item.type !== 'item')) {
-            addKeys(getFirstItem(getChildNodes(item, this.collection))?.key ?? null);
+            addKeys(getFirstItem(getChildNodes(item, collection))?.key ?? null);
           }
         }
 
-        key = this.collection.getKeyAfter(key);
+        key = collection.getKeyAfter(key);
       }
     };
 
-    addKeys(this.collection.getFirstKey());
+    addKeys(collection.getFirstKey());
     return keys;
   }
 
@@ -489,11 +495,15 @@ export class SelectionManager implements MultipleSelectionManager {
   }
 
   canSelectItem(key: Key): boolean {
+    return this.canSelectItemIn(key, this.collection);
+  }
+
+  private canSelectItemIn(key: Key, collection: Collection<Node<unknown>>): boolean {
     if (this.state.selectionMode === 'none' || this.state.disabledKeys.has(key)) {
       return false;
     }
 
-    let item = this.collection.getItem(key);
+    let item = collection.getItem(key);
     if (!item || item?.props?.isDisabled || (item.type === 'cell' && !this.allowsCellSelection)) {
       return false;
     }
@@ -516,7 +526,8 @@ export class SelectionManager implements MultipleSelectionManager {
   withCollection(collection: Collection<Node<unknown>>): SelectionManager {
     return new SelectionManager(collection, this.state, {
       allowsCellSelection: this.allowsCellSelection,
-      layoutDelegate: this.layoutDelegate || undefined
+      layoutDelegate: this.layoutDelegate || undefined,
+      fullCollection: this.fullCollection ?? this.collection
     });
   }
 }

packages/react-aria-components/stories/Autocomplete.stories.tsx

@@ -513,6 +513,30 @@ export const AutocompleteWithListbox: AutocompleteStory = {
   name: 'Autocomplete with ListBox + Popover'
 };
 
+export const AutocompleteSelectAllFiltering: AutocompleteStory = {
+  render: (args) => {
+    return (
+      <AutocompleteWrapper disableVirtualFocus={args.disableVirtualFocus}>
+        <div>
+          <SearchField autoFocus>
+            <Label style={{display: 'block'}}>Test</Label>
+            <Input />
+          </SearchField>
+          <ListBox<AutocompleteItem>
+            className={styles.menu}
+            items={items}
+            selectionMode="multiple"
+            defaultSelectedKeys="all"
+            onSelectionChange={action('onSelectionChange')}>
+            {(item: AutocompleteItem) => <MyListBoxItem id={item.id}>{item.name}</MyListBoxItem>}
+          </ListBox>
+        </div>
+      </AutocompleteWrapper>
+    );
+  },
+  name: 'Autocomplete, select all with filtering'
+};
+
 function VirtualizedListBox(props) {
   let items: {id: number, name: string}[] = [];
   for (let i = 0; i < 10000; i++) {

packages/react-aria-components/test/Autocomplete.test.tsx

@@ -1056,6 +1056,54 @@ describe('Autocomplete', () => {
     options = within(listbox).getAllByRole('option');
     expect(options).toHaveLength(3);
   });
+
+  it('should preserve select all selection when toggling an item in a filtered collection', async function () {
+    let onSelectionChange = jest.fn();
+    let {getByRole} = render(
+      <AutocompleteWrapper>
+        <StaticListbox selectionMode="multiple" defaultSelectedKeys="all" onSelectionChange={onSelectionChange} />
+      </AutocompleteWrapper>
+    );
+
+    let input = getByRole('searchbox');
+    let listbox = getByRole('listbox');
+
+    // All 3 items should be selected initially (Foo, Bar, Baz)
+    let options = within(listbox).getAllByRole('option');
+    expect(options).toHaveLength(3);
+    expect(options[0]).toHaveAttribute('aria-selected', 'true');
+    expect(options[1]).toHaveAttribute('aria-selected', 'true');
+    expect(options[2]).toHaveAttribute('aria-selected', 'true');
+
+    // Filter to show only "Ba" items (Bar, Baz)
+    await user.tab();
+    expect(document.activeElement).toBe(input);
+    await user.keyboard('Ba');
+    act(() => jest.runAllTimers());
+
+    options = within(listbox).getAllByRole('option');
+    expect(options).toHaveLength(2);
+    expect(options[0]).toHaveAttribute('aria-selected', 'true');
+    expect(options[1]).toHaveAttribute('aria-selected', 'true');
+
+    // Move down and deselect Baz
+    await user.keyboard('{ArrowDown}');
+    await user.keyboard('{Enter}');
+    expect(onSelectionChange).toHaveBeenCalledTimes(1);
+    // Should contain Foo and Bar
+    expect(new Set(onSelectionChange.mock.calls[0][0])).toEqual(new Set(['1', '2']));
+
+    // Clear the filter
+    await user.clear(input);
+    act(() => jest.runAllTimers());
+
+    // All items should be visible, with Foo and Bar still selected
+    options = within(listbox).getAllByRole('option');
+    expect(options).toHaveLength(3);
+    expect(options[0]).toHaveAttribute('aria-selected', 'true');
+    expect(options[1]).toHaveAttribute('aria-selected', 'true');
+    expect(options[2]).toHaveAttribute('aria-selected', 'false');
+  });
 });
 
 AriaAutocompleteTests({