From 8077f18a0069c98e04db6a8d87c51b8564728360 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=B0=BD=E9=BE=99?= <brickspert.fjl@antfin.com>
Date: Mon, 20 Dec 2021 15:39:30 +0800
Subject: [PATCH] fix: optimize useSetState ts

---
 packages/hooks/src/useSetState/index.ts | 17 ++++++++++++-----
 1 file changed, 12 insertions(+), 5 deletions(-)

diff --git a/packages/hooks/src/useSetState/index.ts b/packages/hooks/src/useSetState/index.ts
index dc031a6ea8..495bcfd243 100644
--- a/packages/hooks/src/useSetState/index.ts
+++ b/packages/hooks/src/useSetState/index.ts
@@ -1,13 +1,20 @@
 import { useCallback, useState } from 'react';
 import { isFunction } from '../utils';
 
-const useSetState = <T extends Record<string, any>>(
-  initialState: T = {} as T,
-): [T, (patch: Partial<T> | ((prevState: T) => Partial<T>)) => void] => {
-  const [state, setState] = useState<T>(initialState);
+export type SetState<S extends Record<string, any>> = <K extends keyof S>(
+  state: Pick<S, K> | null | ((prevState: Readonly<S>) => Pick<S, K> | S | null),
+) => void;
+
+const useSetState = <S extends Record<string, any>>(
+  initialState: S | (() => S),
+): [S, SetState<S>] => {
+  const [state, setState] = useState<S>(initialState);
 
   const setMergeState = useCallback((patch) => {
-    setState((prevState) => ({ ...prevState, ...(isFunction(patch) ? patch(prevState) : patch) }));
+    setState((prevState) => {
+      const newState = isFunction(patch) ? patch(prevState) : patch;
+      return newState ? { ...prevState, ...newState } : prevState;
+    });
   }, []);
 
   return [state, setMergeState];