fix: reset state when unmounted

src/hooks/useIsFirstRender.ts

@@ -0,0 +1,12 @@
+import * as React from 'react';
+
+export default function useIsFirstRender(): boolean {
+  const firstRenderRef = React.useRef(true);
+  React.useEffect(() => {
+    firstRenderRef.current = false;
+    return () => {
+      firstRenderRef.current = true;
+    };
+  }, []);
+  return firstRenderRef.current;
+}

src/hooks/useMergedState.ts

@@ -1,5 +1,6 @@
 import * as React from 'react';
 import useState from './useState';
+import useIsFirstRender from './useIsFirstRender';
 
 /**
  * Similar to `useState` but will use props value if provided.
@@ -49,10 +50,9 @@ export default function useControlledState<T, R = T>(
   );
 
   // Effect of reset value to `undefined`
-  const firstRenderRef = React.useRef(true);
+  const isFirstRender = useIsFirstRender();
   React.useEffect(() => {
-    if (firstRenderRef.current) {
-      firstRenderRef.current = false;
+    if (isFirstRender) {
       return;
     }
 

src/hooks/useState.ts

@@ -8,11 +8,11 @@ export type SetState<T> = (
    * Will not update state when destroyed.
    * Developer should make sure this is safe to ignore.
    */
-  ignoreDestroy?: boolean,
+  ignoreSetStateIfDestroyed?: boolean,
 ) => void;
 
 /**
- * Same as React.useState but `setState` accept `ignoreDestroy` param to not to setState after destroyed.
+ * Same as React.useState but `setState` accept `ignoreSetStateIfDestroyed` param to not to setState after destroyed.
  * We do not make this auto is to avoid real memory leak.
  * Developer should confirm it's safe to ignore themselves.
  */
@@ -22,15 +22,18 @@ export default function useState<T>(
   const destroyRef = React.useRef(false);
   const [value, setValue] = React.useState(defaultValue);
 
-  React.useEffect(
-    () => () => {
+  React.useEffect(() => {
+    destroyRef.current = false;
+    return () => {
       destroyRef.current = true;
-    },
-    [],
-  );
+    };
+  }, []);
 
-  function safeSetState(updater: Updater<T>, ignoreDestroy?: boolean) {
-    if (ignoreDestroy && destroyRef.current) {
+  function safeSetState(
+    updater: Updater<T>,
+    ignoreSetStateIfDestroyed?: boolean,
+  ) {
+    if (ignoreSetStateIfDestroyed && destroyRef.current) {
       return;
     }
 

tests/hooks.test.js

@@ -4,6 +4,7 @@ import useMemo from '../src/hooks/useMemo';
 import useMergedState from '../src/hooks/useMergedState';
 import useLayoutEffect from '../src/hooks/useLayoutEffect';
 import useState from '../src/hooks/useState';
+import useIsFirstRender from '../src/hooks/useIsFirstRender';
 
 describe('hooks', () => {
   it('useMemo', () => {
@@ -177,4 +178,47 @@ describe('hooks', () => {
       }, 50);
     });
   });
+
+  it('useState in StrictMode', () => {
+    const Demo = () => {
+      const [count, setCount] = useState(0);
+
+      React.useEffect(() => {
+        setCount(count => count + 1);
+      }, []);
+
+      return count;
+    };
+
+    const { container } = render(<Demo />);
+    const { container: strictContainer } = render(<Demo />, {
+      wrapper: React.StrictMode,
+    });
+
+    expect(container.textContent).toBe('1');
+    // useEffect will be triggered twice in StrictMode,
+    // useState should work in StrictMode
+    expect(strictContainer.textContent).toBe('2');
+  });
+
+  it('useIsFirstRender', async () => {
+    const Demo = () => {
+      const isFirstRender = useIsFirstRender();
+      const [, update] = React.useState();
+      return (
+        <>
+          <div id="state">{String(isFirstRender)}</div>
+          <button data-testid="update" onClick={update} />
+        </>
+      );
+    };
+
+    const { container, findByTestId } = render(<Demo />);
+    const button = await findByTestId('update');
+    expect(container.textContent).toEqual('true');
+    fireEvent.click(button);
+    expect(container.textContent).toEqual('false');
+    fireEvent.click(button);
+    expect(container.textContent).toEqual('false');
+  });
 });