Fix useListData remove method for all selection (#2206)

* Fix useListData remove method for all selection

* Handle empty items scenario for remove method

Author: jluyau

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

@@ -211,9 +211,15 @@ export function createListActions<T>(opts: ListOptions<T>, dispatch: (updater: (
         let keySet = new Set(keys);
         let items = state.items.filter(item => !keySet.has(getKey(item)));
 
-        let selection = new Set(state.selectedKeys);
-        for (let key of keys) {
-          selection.delete(key);
+        let selection: Selection = 'all';
+        if (state.selectedKeys !== 'all') {
+          selection = new Set(state.selectedKeys);
+          for (let key of keys) {
+            selection.delete(key);
+          }
+        }
+        if (selection === 'all' && items.length === 0) {
+          selection = new Set();
         }
 
         return {

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

@@ -225,6 +225,28 @@ describe('useListData', function () {
     expect(result.current.selectedKeys).toEqual(new Set(['Julia']));
   });
 
+  it('should preserve all selected value through a remove call', function () {
+    let {result} = renderHook(() => useListData({initialItems: initial, getKey, initialSelectedKeys: 'all'}));
+
+    act(() => {
+      result.current.remove('Sam');
+    });
+
+    expect(result.current.selectedKeys).toEqual('all');
+  });
+
+  it('should change all selection to empty set if last item is removed', function () {
+    let {result} = renderHook(() => useListData({initialItems: initial, getKey, initialSelectedKeys: 'all'}));
+
+    act(() => {
+      result.current.remove('Sam');
+      result.current.remove('David');
+      result.current.remove('Julia');
+    });
+
+    expect(result.current.selectedKeys).toEqual(new Set());
+  });
+
   it('should remove multiple items', function () {
     let {result} = renderHook(() => useListData({initialItems: initial, getKey, initialSelectedKeys: ['Sam', 'David', 'Julia']}));
     let initialResult = result.current;