Update useListData to handle select all w/ unloaded lists (#2225)

* handle removal and selection for scenarios with unloaded items

* better solution + tests

* Revert removeSelectedItems functionality and add extra test

* remove outdated comment

Author: jluyau

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

@@ -342,7 +342,7 @@ export function useAsyncList<T, C = string>(options: AsyncListOptions<T, C>): As
     sort(sortDescriptor: SortDescriptor) {
       dispatchFetch({type: 'sorting', sortDescriptor}, sort || load);
     },
-    ...createListActions({...options, getKey}, fn => {
+    ...createListActions({...options, getKey, cursor: data.cursor}, fn => {
       dispatch({type: 'update', updater: fn});
     }),
     setFilterText(filterText: string) {

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

@@ -128,6 +128,10 @@ export interface ListState<T> {
   filterText: string
 }
 
+interface CreateListOptions<T, C> extends ListOptions<T> {
+  cursor?: C
+}
+
 /**
  * Manages state for an immutable list data structure, and provides convenience methods to
  * update the data over time.
@@ -162,8 +166,8 @@ export function useListData<T>(options: ListOptions<T>): ListData<T> {
   };
 }
 
-export function createListActions<T>(opts: ListOptions<T>, dispatch: (updater: (state: ListState<T>) => ListState<T>) => void): Omit<ListData<T>, 'items' | 'selectedKeys' | 'getItem' | 'filterText'> {
-  let {getKey} = opts;
+export function createListActions<T, C>(opts: CreateListOptions<T, C>, dispatch: (updater: (state: ListState<T>) => ListState<T>) => void): Omit<ListData<T>, 'items' | 'selectedKeys' | 'getItem' | 'filterText'> {
+  let {cursor, getKey} = opts;
   return {
     setSelectedKeys(selectedKeys: Selection) {
       dispatch(state => ({
@@ -218,7 +222,7 @@ export function createListActions<T>(opts: ListOptions<T>, dispatch: (updater: (
             selection.delete(key);
           }
         }
-        if (selection === 'all' && items.length === 0) {
+        if (cursor == null && items.length === 0) {
           selection = new Set();
         }
 

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

@@ -29,6 +29,12 @@ function getItems2() {
   });
 }
 
+function getItemsEnd() {
+  return new Promise(resolve => {
+    setTimeout(() => resolve({items: ITEMS, cursor: null}), 100);
+  });
+}
+
 describe('useAsyncList', () => {
   beforeAll(() => {
     jest.useFakeTimers();
@@ -870,6 +876,59 @@ describe('useAsyncList', () => {
     expect(result.current.selectedKeys).toEqual('all');
   });
 
+  it('should maintain all selection if last visible item removed and unloaded items still exist', async () => {
+    let load = jest.fn()
+      .mockImplementationOnce(getItems);
+    let {result, waitForNextUpdate} = renderHook(
+      () => useAsyncList({load})
+    );
+    await act(async () => {
+      result.current.loadMore();
+      jest.runAllTimers();
+      await waitForNextUpdate();
+      result.current.setSelectedKeys('all');
+      result.current.remove(1);
+      result.current.remove(2);
+      jest.runAllTimers();
+    });
+    expect(result.current.selectedKeys).toEqual('all');
+  });
+
+  it('should change selection to empty set if last item removed with no unloaded items left', async () => {
+    let load = jest.fn()
+      .mockImplementationOnce(getItemsEnd);
+    let {result, waitForNextUpdate} = renderHook(
+      () => useAsyncList({load})
+    );
+    await act(async () => {
+      result.current.loadMore();
+      jest.runAllTimers();
+      await waitForNextUpdate();
+      result.current.setSelectedKeys('all');
+      result.current.remove(1);
+      result.current.remove(2);
+      jest.runAllTimers();
+    });
+    expect(result.current.selectedKeys).toEqual(new Set());
+  });
+
+  it('should change selection to empty set if all items removed', async () => {
+    let load = jest.fn()
+      .mockImplementationOnce(getItemsEnd);
+    let {result, waitForNextUpdate} = renderHook(
+      () => useAsyncList({load})
+    );
+    await act(async () => {
+      result.current.loadMore();
+      jest.runAllTimers();
+      await waitForNextUpdate();
+      result.current.setSelectedKeys('all');
+      result.current.removeSelectedItems();
+      jest.runAllTimers();
+    });
+    expect(result.current.selectedKeys).toEqual(new Set());
+  });
+
   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'}];