Enable propagation of all value through useAsyncList methods (#2200)

* Enable propagation of all value through useAsyncList methods

* review fix

* making table story use list.selectedKeys

Co-authored-by: Daniel Lu <[email protected]>

Author: jluyau

packages/@react-spectrum/table/stories/Table.stories.tsx

@@ -1008,7 +1008,7 @@ function AsyncLoadingExample() {
   return (
     <div>
       <ActionButton marginBottom={10} onPress={() => list.remove(list.items[0].data.id)}>Remove first item</ActionButton>
-      <TableView aria-label="Top news from Reddit" selectionMode="multiple" width={1000} height={400} isQuiet sortDescriptor={list.sortDescriptor} onSortChange={list.sort}>
+      <TableView aria-label="Top news from Reddit" selectionMode="multiple" width={1000} height={400} isQuiet sortDescriptor={list.sortDescriptor} onSortChange={list.sort} selectedKeys={list.selectedKeys} onSelectionChange={list.setSelectedKeys}>
         <TableHeader>
           <Column key="score" width={100} allowsSorting>Score</Column>
           <Column key="title" isRowHeader allowsSorting>Title</Column>

packages/@react-stately/data/src/useAsyncList.ts

@@ -111,6 +111,7 @@ interface AsyncListData<T> extends ListData<T> {
 }
 
 function reducer<T, C>(data: AsyncListState<T, C>, action: Action<T, C>): AsyncListState<T, C> {
+  let selectedKeys;
   switch (data.state) {
     case 'idle':
     case 'error':
@@ -151,12 +152,13 @@ function reducer<T, C>(data: AsyncListState<T, C>, action: Action<T, C>): AsyncL
             return data;
           }
 
+          selectedKeys = action.selectedKeys ?? data.selectedKeys;
           return {
             ...data,
             filterText: action.filterText ?? data.filterText,
             state: 'idle',
             items: [...action.items],
-            selectedKeys: new Set(action.selectedKeys ?? data.selectedKeys),
+            selectedKeys: selectedKeys === 'all' ? 'all' : new Set(selectedKeys),
             sortDescriptor: action.sortDescriptor ?? data.sortDescriptor,
             abortController: null,
             cursor: action.cursor
@@ -200,12 +202,15 @@ function reducer<T, C>(data: AsyncListState<T, C>, action: Action<T, C>): AsyncL
     case 'loadingMore':
       switch (action.type) {
         case 'success':
+          selectedKeys = (data.selectedKeys === 'all' || action.selectedKeys === 'all')
+            ? 'all'
+            : new Set([...data.selectedKeys, ...(action.selectedKeys ?? [])]);
           // Append the new items
           return {
             ...data,
             state: 'idle',
             items: [...data.items, ...action.items],
-            selectedKeys: new Set([...data.selectedKeys, ...(action.selectedKeys ?? [])]),
+            selectedKeys,
             sortDescriptor: action.sortDescriptor ?? data.sortDescriptor,
             abortController: null,
             cursor: action.cursor
@@ -274,7 +279,7 @@ export function useAsyncList<T, C = string>(options: AsyncListOptions<T, C>): As
     state: 'idle',
     error: null,
     items: [],
-    selectedKeys: new Set(initialSelectedKeys),
+    selectedKeys: initialSelectedKeys === 'all' ? 'all' : new Set(initialSelectedKeys),
     sortDescriptor: initialSortDescriptor,
     filterText: initialFilterText
   });

packages/@react-stately/data/test/useAsyncList.test.js

@@ -831,6 +831,45 @@ describe('useAsyncList', () => {
     expect(result.current.items).toEqual(ITEMS2);
   });
 
+  it('should maintain all selection through a loadMore call', async () => {
+    let load = jest.fn()
+      .mockImplementationOnce(getItems)
+      .mockImplementationOnce(getItems2);
+    let {result} = renderHook(
+      () => useAsyncList({load})
+    );
+
+    expect(load).toHaveBeenCalledTimes(1);
+    expect(result.current.isLoading).toBe(true);
+    expect(result.current.items).toEqual([]);
+
+    await act(async () => {
+      jest.runAllTimers();
+      result.current.setSelectedKeys('all');
+    });
+
+    expect(result.current.selectedKeys).toEqual('all');
+
+    await act(async () => {
+      result.current.loadMore();
+      jest.runAllTimers();
+    });
+
+    expect(load).toHaveBeenCalledTimes(2);
+    expect(result.current.isLoading).toBe(false);
+    expect(result.current.items).toEqual([...ITEMS, ...ITEMS2]);
+    expect(result.current.selectedKeys).toEqual('all');
+  });
+
+  it('should accept all for initialSelectedKeys', () => {
+    let load = jest.fn()
+      .mockImplementationOnce(getItems);
+    let {result} = renderHook(
+      () => useAsyncList({load, initialSelectedKeys: 'all'})
+    );
+    expect(result.current.selectedKeys).toEqual('all');
+  });
+
   describe('filtering', function () {
     const filterItems = [{id: 1, name: 'Bob'}, {id: 2, name: 'Joe'}, {id: 3, name: 'Bob Joe'}];
     const itemsFirstCall = [{id: 1, name: 'Bob'}, {id: 3, name: 'Bob Joe'}];
@@ -889,6 +928,29 @@ describe('useAsyncList', () => {
       expect(result.current.filterText).toEqual('Blah');
     });
 
+    it('should preserve all selectedKeys through filtering', async () => {
+      let load = jest.fn().mockImplementation(getFilterItems);
+      let initialFilterText = 'Blah';
+      let {result, waitForNextUpdate} = renderHook(() => useAsyncList({load, initialFilterText}));
+
+      await act(async () => {
+        result.current.setSelectedKeys('all');
+      });
+
+      expect(result.current.selectedKeys).toEqual('all');
+
+      await act(async () => {
+        jest.runAllTimers();
+        await waitForNextUpdate();
+      });
+
+      expect(result.current.loadingState).toBe('idle');
+      expect(result.current.isLoading).toBe(false);
+      expect(result.current.items).toEqual(filterItems);
+      expect(result.current.filterText).toEqual('Blah');
+      expect(result.current.selectedKeys).toEqual('all');
+    });
+
     it('should update the list of items when the filter text changes (server side filtering)', async () => {
       let load = jest
         .fn()