From b72366eb41218816535aff3c15447c46b7259202 Mon Sep 17 00:00:00 2001
From: KangXinzhi <50164727+KangXinzhi@users.noreply.github.com>
Date: Thu, 23 Mar 2023 20:45:29 +0800
Subject: [PATCH] refactor(useMap/useSet): refactoring useMap and useSet for
 improved initialValue checking. (#2116)

* fix(useMap): map initialValue supports undefined

* fix(useSet): set initialValue supports undefined

* fix: use memoized state in useMap and useSet

* fix: getInitValue is used as a function
---
 .../hooks/src/useMap/__tests__/index.test.ts  | 28 ++++++++++++++++++-
 packages/hooks/src/useMap/index.ts            |  7 ++---
 .../hooks/src/useSet/__tests__/index.test.ts  |  4 ++-
 packages/hooks/src/useSet/index.ts            |  7 ++---
 4 files changed, 34 insertions(+), 12 deletions(-)

diff --git a/packages/hooks/src/useMap/__tests__/index.test.ts b/packages/hooks/src/useMap/__tests__/index.test.ts
index 4c79ef6b5a..019c8b0893 100644
--- a/packages/hooks/src/useMap/__tests__/index.test.ts
+++ b/packages/hooks/src/useMap/__tests__/index.test.ts
@@ -26,8 +26,10 @@ describe('useMap', () => {
 
   it('should init empty map if not initial object provided', () => {
     const { result } = setup();
-
     expect([...result.current[0]]).toEqual([]);
+
+    const { result: result2 } = setup(undefined);
+    expect([...result2.current[0]]).toEqual([]);
   });
 
   it('should get corresponding value for initial provided key', () => {
@@ -131,6 +133,12 @@ describe('useMap', () => {
       ['foo', 'foo'],
       ['a', 2],
     ]);
+
+    act(() => {
+      // @ts-ignore
+      utils.setAll();
+    });
+    expect([...result.current[0]]).toEqual([]);
   });
 
   it('remove should be work', () => {
@@ -141,6 +149,24 @@ describe('useMap', () => {
       remove('msg');
     });
     expect(result.current[0].size).toBe(0);
+
+    const { result: result2 } = setup([
+      ['foo', 'bar'],
+      ['a', 1],
+      ['b', 2],
+      ['c', 3],
+    ]);
+    const [, utils] = result2.current;
+
+    act(() => {
+      utils.remove('a');
+    });
+
+    expect([...result2.current[0]]).toEqual([
+      ['foo', 'bar'],
+      ['b', 2],
+      ['c', 3],
+    ]);
   });
 
   it('reset should be work', () => {
diff --git a/packages/hooks/src/useMap/index.ts b/packages/hooks/src/useMap/index.ts
index 464ff3e4e9..4f0407fb57 100644
--- a/packages/hooks/src/useMap/index.ts
+++ b/packages/hooks/src/useMap/index.ts
@@ -2,11 +2,8 @@ import { useState } from 'react';
 import useMemoizedFn from '../useMemoizedFn';
 
 function useMap<K, T>(initialValue?: Iterable<readonly [K, T]>) {
-  const getInitValue = () => {
-    return initialValue === undefined ? new Map() : new Map(initialValue);
-  };
-
-  const [map, setMap] = useState<Map<K, T>>(() => getInitValue());
+  const getInitValue = () => new Map(initialValue);
+  const [map, setMap] = useState<Map<K, T>>(getInitValue);
 
   const set = (key: K, entry: T) => {
     setMap((prev) => {
diff --git a/packages/hooks/src/useSet/__tests__/index.test.ts b/packages/hooks/src/useSet/__tests__/index.test.ts
index 8039b53fa4..50ed21ff53 100644
--- a/packages/hooks/src/useSet/__tests__/index.test.ts
+++ b/packages/hooks/src/useSet/__tests__/index.test.ts
@@ -18,8 +18,10 @@ describe('useSet', () => {
 
   it('should init empty set if no initial set provided', () => {
     const { result } = setUp();
-
     expect(result.current[0]).toEqual(new Set());
+
+    const { result: result1 } = setUp(undefined);
+    expect(result1.current[0]).toEqual(new Set());
   });
 
   it('should have an initially provided key', () => {
diff --git a/packages/hooks/src/useSet/index.ts b/packages/hooks/src/useSet/index.ts
index e7bd7cadb6..59f18debcb 100644
--- a/packages/hooks/src/useSet/index.ts
+++ b/packages/hooks/src/useSet/index.ts
@@ -2,11 +2,8 @@ import { useState } from 'react';
 import useMemoizedFn from '../useMemoizedFn';
 
 function useSet<K>(initialValue?: Iterable<K>) {
-  const getInitValue = () => {
-    return initialValue === undefined ? new Set<K>() : new Set(initialValue);
-  };
-
-  const [set, setSet] = useState<Set<K>>(() => getInitValue());
+  const getInitValue = () => new Set(initialValue);
+  const [set, setSet] = useState<Set<K>>(getInitValue);
 
   const add = (key: K) => {
     if (set.has(key)) {