feat: allow useSetState setter to accept function

* make setState of useSetState accept function argument

* update storybook

Author: clumsyme

docs/useSetState.md

@@ -18,6 +18,15 @@ const Demo = () => {
       <pre>{JSON.stringify(state, null, 2)}</pre>
       <button onClick={() => setState({hello: 'world'})}>hello</button>
       <button onClick={() => setState({foo: 'bar'})}>foo</button>
+      <button 
+        onClick={() => {
+          setState((prevState) => ({
+            count: prevState.count === undefined ? 0 : prevState.count + 1,
+          }))
+        }}
+      >
+        count
+      </button>
     </div>
   );
 };

src/__stories__/useSetState.story.tsx

@@ -11,6 +11,15 @@ const Demo = () => {
       <pre>{JSON.stringify(state, null, 2)}</pre>
       <button onClick={() => setState({hello: 'world'})}>hello</button>
       <button onClick={() => setState({foo: 'bar'})}>foo</button>
+      <button 
+        onClick={() => {
+          setState((prevState) => ({
+            count: prevState.count === undefined ? 0 : prevState.count + 1,
+          }))
+        }}
+      >
+        count
+      </button>
     </div>
   );
 };

src/react.ts

@@ -1,6 +1,6 @@
 import * as React from 'react';
 
-export type UseState = <T>(initialState: T | (() => T)) => [T, (newState: T) => void];
+export type UseState = <T>(initialState: T | (() => T)) => [T, (newState: T | ((newState) => T)) => void];
 export const useState: UseState = (React as any).useState;
 
 export type UseEffect = (didUpdate: () => ((() => void) | void), params?: any[]) => void;

src/useSetState.ts

@@ -1,10 +1,16 @@
 import {useState} from './react';
 
-const useSetState = <T extends object>(initialState: T = {} as T): [T, (patch: Partial<T>) => void]=> {
+const useSetState = <T extends object>(initialState: T = {} as T): [T, (patch: Partial<T> | Function) => void]=> {
   const [state, set] = useState<T>(initialState);
   const setState = (patch) => {
-    Object.assign(state, patch);
-    set(state);
+    if (patch instanceof Function) {
+      set((prevState) => {
+          return Object.assign(state, patch(prevState))
+      })
+    } else {
+      Object.assign(state, patch);
+      set(state);
+    }
   };
 
   return [state, setState];