Fix transfer selected keys (#51523)

* fix: transfer selected keys

* test: add transfer test case

* test: remove unused code

* refactor: use cal data

---------

Co-authored-by: karos <chenwenfan@harmonycloud.cn>
Co-authored-by: 二货机器人 <smith3816@gmail.com>

components/transfer/__tests__/index.test.tsx

@@ -7,6 +7,7 @@ import Transfer from '..';
 import mountTest from '../../../tests/shared/mountTest';
 import rtlTest from '../../../tests/shared/rtlTest';
 import Button from '../../button';
+import { waitFakeTimer } from '../../../tests/utils';
 
 const listCommonProps: {
   dataSource: { key: string; title: string; disabled?: boolean }[];
@@ -767,6 +768,70 @@ describe('Transfer', () => {
 
     errSpy.mockRestore();
   });
+  it('it checks correctly after changing the dataSource', async () => {
+    const mockData = Array.from({ length: 10 }).map((_, i) => ({
+      key: i.toString(),
+      title: `content${i + 1}`,
+      description: `description of content${i + 1}`,
+    }));
+
+    const initialTargetKeys = mockData
+      .filter((item) => Number(item.key) > 4)
+      .map((item) => item.key);
+
+    const defaultCheckedKeys = ['1', '2'];
+    const handleSelectChange = jest.fn();
+    const App: React.FC = () => {
+      const [targetKeys, setTargetKeys] = useState<TransferProps['targetKeys']>(initialTargetKeys);
+      const [selectedKeys, setSelectedKeys] = useState<TransferProps['targetKeys']>([]);
+
+      const [dataSource, setDataSource] = useState(mockData);
+
+      const onChange: TransferProps['onChange'] = (nextTargetKeys) => {
+        setTargetKeys(nextTargetKeys);
+      };
+
+      return (
+        <>
+          <Button
+            className="update-btn"
+            onClick={() => {
+              setSelectedKeys(defaultCheckedKeys);
+              setDataSource([]);
+              setDataSource([...mockData]);
+              // setTimeout(() => {
+              //   setDataSource([...mockData]);
+              // });
+            }}
+          >
+            update
+          </Button>
+          <Transfer
+            dataSource={dataSource}
+            titles={['Source', 'Target']}
+            targetKeys={targetKeys}
+            selectedKeys={selectedKeys}
+            onChange={onChange}
+            onSelectChange={handleSelectChange}
+            render={(item) => item.title}
+          />
+        </>
+      );
+    };
+    const { container } = render(<App />);
+
+    fireEvent.click(container.querySelector('.update-btn')!);
+    await waitFakeTimer();
+
+    defaultCheckedKeys.forEach((item) => {
+      expect(
+        container
+          ?.querySelectorAll('.ant-transfer-list-content-item')
+          ?.item(Number(item))
+          ?.querySelector('input[type="checkbox"]')!,
+      ).toBeChecked();
+    });
+  });
 });
 
 describe('immutable data', () => {

components/transfer/hooks/useSelection.ts

@@ -1,4 +1,5 @@
 import * as React from 'react';
+import { useEvent, useMergedState } from 'rc-util';
 
 import type { TransferKey } from '../interface';
 
@@ -16,12 +17,12 @@ function flattenKeys(keys: Set<TransferKey>) {
 export default function useSelection<T extends { key: TransferKey }>(
   leftDataSource: T[],
   rightDataSource: T[],
-  selectedKeys: TransferKey[] = EMPTY_KEYS,
+  selectedKeys?: TransferKey[],
 ): [
   sourceSelectedKeys: TransferKey[],
   targetSelectedKeys: TransferKey[],
-  setSourceSelectedKeys: React.Dispatch<React.SetStateAction<TransferKey[]>>,
-  setTargetSelectedKeys: React.Dispatch<React.SetStateAction<TransferKey[]>>,
+  setSourceSelectedKeys: (srcKeys: TransferKey[]) => void,
+  setTargetSelectedKeys: (srcKeys: TransferKey[]) => void,
 ] {
   // Prepare `dataSource` keys
   const [leftKeys, rightKeys] = React.useMemo(
@@ -33,25 +34,35 @@ export default function useSelection<T extends { key: TransferKey }>(
   );
 
   // Selected Keys
-  const [sourceSelectedKeys, setSourceSelectedKeys] = React.useState(() =>
-    filterKeys(selectedKeys, leftKeys),
+  const [mergedSelectedKeys, setMergedSelectedKeys] = useMergedState(EMPTY_KEYS, {
+    value: selectedKeys,
+  });
+
+  const sourceSelectedKeys = React.useMemo(
+    () => filterKeys(mergedSelectedKeys, leftKeys),
+    [mergedSelectedKeys, leftKeys],
   );
-  const [targetSelectedKeys, setTargetSelectedKeys] = React.useState(() =>
-    filterKeys(selectedKeys, rightKeys),
+  const targetSelectedKeys = React.useMemo(
+    () => filterKeys(mergedSelectedKeys, rightKeys),
+    [mergedSelectedKeys, rightKeys],
   );
 
-  // Fill selected keys
+  // // Reset when data changed
   React.useEffect(() => {
-    setSourceSelectedKeys(filterKeys(selectedKeys, leftKeys));
-    setTargetSelectedKeys(filterKeys(selectedKeys, rightKeys));
-  }, [selectedKeys]);
-
-  // Reset when data changed
-  React.useEffect(() => {
-    setSourceSelectedKeys(filterKeys(sourceSelectedKeys, leftKeys));
-    setTargetSelectedKeys(filterKeys(targetSelectedKeys, rightKeys));
+    setMergedSelectedKeys([
+      ...filterKeys(mergedSelectedKeys, leftKeys),
+      ...filterKeys(mergedSelectedKeys, rightKeys),
+    ]);
   }, [flattenKeys(leftKeys), flattenKeys(rightKeys)]);
 
+  // Update keys
+  const setSourceSelectedKeys = useEvent((nextSrcKeys: TransferKey[]) => {
+    setMergedSelectedKeys([...nextSrcKeys, ...targetSelectedKeys]);
+  });
+  const setTargetSelectedKeys = useEvent((nextTargetKeys: TransferKey[]) => {
+    setMergedSelectedKeys([...sourceSelectedKeys, ...nextTargetKeys]);
+  });
+
   return [
     // Keys
     sourceSelectedKeys,