feat: useLayoutEffect support firstMount param (#447)

* refactor: useLayoutEffectSupport

* test: add test case

Author: zombieJ

src/hooks/useLayoutEffect.ts

@@ -4,30 +4,39 @@ import canUseDom from '../Dom/canUseDom';
 /**
  * Wrap `React.useLayoutEffect` which will not throw warning message in test env
  */
-const useLayoutEffect =
+const useInternalLayoutEffect =
   process.env.NODE_ENV !== 'test' && canUseDom()
     ? React.useLayoutEffect
     : React.useEffect;
 
-export default useLayoutEffect;
-
-export const useLayoutUpdateEffect: typeof React.useEffect = (
-  callback,
-  deps,
+const useLayoutEffect = (
+  callback: (mount: boolean) => void | VoidFunction,
+  deps?: React.DependencyList,
 ) => {
   const firstMountRef = React.useRef(true);
 
-  useLayoutEffect(() => {
-    if (!firstMountRef.current) {
-      return callback();
-    }
+  useInternalLayoutEffect(() => {
+    return callback(firstMountRef.current);
   }, deps);
 
   // We tell react that first mount has passed
-  useLayoutEffect(() => {
+  useInternalLayoutEffect(() => {
     firstMountRef.current = false;
     return () => {
       firstMountRef.current = true;
     };
   }, []);
 };
+
+export const useLayoutUpdateEffect: typeof React.useEffect = (
+  callback,
+  deps,
+) => {
+  useLayoutEffect(firstMount => {
+    if (!firstMount) {
+      return callback();
+    }
+  }, deps);
+};
+
+export default useLayoutEffect;

tests/hooks.test.js

@@ -343,6 +343,30 @@ describe('hooks', () => {
       expect(errorSpy).not.toHaveBeenCalled();
       errorSpy.mockRestore();
     });
+
+    it('can get mount state', () => {
+      const Demo = () => {
+        const timesRef = React.useRef(0);
+        const [, forceUpdate] = React.useState(0);
+
+        useLayoutEffect(firstMount => {
+          if (timesRef.current === 0) {
+            expect(firstMount).toBeTruthy();
+            forceUpdate(1);
+          } else {
+            expect(firstMount).toBeFalsy();
+            forceUpdate(2);
+          }
+
+          timesRef.current += 1;
+        });
+
+        return <p>{timesRef.current}</p>;
+      };
+
+      const { container } = render(<Demo />);
+      expect(container.querySelector('p').textContent).toEqual('2');
+    });
   });
 
   describe('useState', () => {