From cc60ca31dacd5f36a7482fcf8c2ec8a43a2066db Mon Sep 17 00:00:00 2001
From: unix <unix.bio@gmail.com>
Date: Sun, 24 May 2020 01:51:27 +0800
Subject: [PATCH] feat: support for functional initializer mode

---
 .../use-current-state/__tests__/current-state.test.tsx    | 8 ++++++++
 components/use-current-state/use-current-state.ts         | 6 ++++--
 2 files changed, 12 insertions(+), 2 deletions(-)

diff --git a/components/use-current-state/__tests__/current-state.test.tsx b/components/use-current-state/__tests__/current-state.test.tsx
index 8f90c12..f65ca1f 100644
--- a/components/use-current-state/__tests__/current-state.test.tsx
+++ b/components/use-current-state/__tests__/current-state.test.tsx
@@ -13,6 +13,14 @@ describe('UseCurrentState', () => {
     expect(result.current[2].current).toEqual('test')
   })
 
+  it('functional initial mode should be supported', () => {
+    const { result } = renderHook(() => useCurrentState(() => 1 + 1))
+    expect(result.current[0]).toEqual(2)
+    act(() => result.current[1](0))
+    expect(result.current[0]).toEqual(0)
+    expect(result.current[2].current).toEqual(0)
+  })
+
   it('functional update mode should be supported', () => {
     const { result } = renderHook(() => useCurrentState(0))
     expect(result.current[0]).toEqual(0)
diff --git a/components/use-current-state/use-current-state.ts b/components/use-current-state/use-current-state.ts
index d0a60af..0d8dce6 100644
--- a/components/use-current-state/use-current-state.ts
+++ b/components/use-current-state/use-current-state.ts
@@ -2,8 +2,10 @@ import { Dispatch, MutableRefObject, SetStateAction, useEffect, useRef, useState
 
 export type CurrentStateType<S> = [S, Dispatch<SetStateAction<S>>, MutableRefObject<S>]
 
-const useCurrentState = <S>(initialState: S): CurrentStateType<S> => {
-  const [state, setState] = useState<S>(initialState as S)
+const useCurrentState = <S>(initialState: S | (() => S)): CurrentStateType<S> => {
+  const [state, setState] = useState<S>(() => {
+    return typeof initialState === 'function' ? (initialState as () => S)() : initialState
+  })
   const ref = useRef<S>(initialState as S)
 
   useEffect(() => {