refactor: useInsertionEffect if in React 18 (#121)

* fix: order of memo & effect

* chore: mv effect style in

* chore: add test case

* test: add 17 test

* chore: clean up

src/hooks/useGlobalCache.tsx

@@ -2,15 +2,19 @@ import * as React from 'react';
 import type { KeyType } from '../Cache';
 import StyleContext from '../StyleContext';
 import useHMR from './useHMR';
+import useInsertionEffect from './useInsertionEffect';
 
-export default function useClientCache<CacheType>(
+export default function useGlobalCache<CacheType>(
   prefix: string,
   keyPath: KeyType[],
   cacheFn: () => CacheType,
   onCacheRemove?: (cache: CacheType, fromHMR: boolean) => void,
+  // Add additional effect trigger by `useInsertionEffect`
+  onCacheEffect?: (cachedValue: CacheType) => void,
 ): CacheType {
   const { cache: globalCache } = React.useContext(StyleContext);
   const fullPath = [prefix, ...keyPath];
+  const deps = fullPath.join('_');
 
   const HMRUpdate = useHMR();
 
@@ -40,12 +44,16 @@ export default function useClientCache<CacheType>(
   React.useMemo(
     () => buildCache(),
     /* eslint-disable react-hooks/exhaustive-deps */
-    [fullPath.join('_')],
+    [deps],
     /* eslint-enable */
   );
 
+  const cacheContent = globalCache.get(fullPath)![1];
+
   // Remove if no need anymore
-  React.useEffect(() => {
+  useInsertionEffect(() => {
+    onCacheEffect?.(cacheContent);
+
     // It's bad to call build again in effect.
     // But we have to do this since StrictMode will call effect twice
     // which will clear cache on the first time.
@@ -64,7 +72,7 @@ export default function useClientCache<CacheType>(
         return [times - 1, cache];
       });
     };
-  }, fullPath);
+  }, [deps]);
 
-  return globalCache.get(fullPath)![1];
+  return cacheContent;
 }

src/hooks/useInsertionEffect.tsx

@@ -0,0 +1,13 @@
+// import canUseDom from 'rc-util/lib/Dom/canUseDom';
+import useLayoutEffect from 'rc-util/lib/hooks/useLayoutEffect';
+import * as React from 'react';
+
+// We need fully clone React function here
+// to avoid webpack warning React 17 do not export `useId`
+const fullClone = {
+  ...React,
+};
+const { useInsertionEffect } = fullClone;
+const useMergedInsertionEffect = useInsertionEffect || useLayoutEffect;
+
+export default useMergedInsertionEffect;

src/hooks/useStyleRegister.tsx

@@ -398,9 +398,22 @@ export default function useStyleRegister(
         transformers,
         linters,
       });
+
       const styleStr = normalizeStyle(parsedStyle);
       const styleId = uniqueHash(fullPath, styleStr);
 
+      return [styleStr, tokenKey, styleId, effectStyle];
+    },
+
+    // Remove cache if no need
+    ([, , styleId], fromHMR) => {
+      if ((fromHMR || autoClear) && isClientSide) {
+        removeCSS(styleId, { mark: ATTR_MARK });
+      }
+    },
+
+    // Inject style here
+    ([styleStr, _, styleId, effectStyle]) => {
       if (isMergedClientSide) {
         const mergedCSSConfig: Parameters<typeof updateCSS>[2] = {
           mark: ATTR_MARK,
@@ -435,14 +448,6 @@ export default function useStyleRegister(
           );
         });
       }
-
-      return [styleStr, tokenKey, styleId];
-    },
-    // Remove cache if no need
-    ([, , styleId], fromHMR) => {
-      if ((fromHMR || autoClear) && isClientSide) {
-        removeCSS(styleId, { mark: ATTR_MARK });
-      }
     },
   );
 

tests/legacy.spec.tsx

@@ -0,0 +1,109 @@
+import { render } from '@testing-library/react';
+import * as React from 'react';
+import type { CSSInterpolation } from '../src';
+import { Theme, useCacheToken, useStyleRegister } from '../src';
+
+interface DesignToken {
+  primaryColor: string;
+}
+
+interface DerivativeToken extends DesignToken {
+  primaryColorDisabled: string;
+}
+
+const derivative = (designToken: DesignToken): DerivativeToken => ({
+  ...designToken,
+  primaryColorDisabled: designToken.primaryColor,
+});
+
+const baseToken: DesignToken = {
+  primaryColor: '#1890ff',
+};
+
+const theme = new Theme(derivative);
+
+vi.mock('react', async () => {
+  const origin: any = await vi.importActual('react');
+
+  return {
+    ...origin,
+    useInsertionEffect: undefined,
+  };
+});
+
+// Same as `index.spec.tsx` but we hack to no to support `useInsertionEffect`
+describe('legacy React version', () => {
+  beforeEach(() => {
+    const styles = Array.from(document.head.querySelectorAll('style'));
+    styles.forEach((style) => {
+      style.parentNode?.removeChild(style);
+    });
+  });
+
+  const genStyle = (token: DerivativeToken): CSSInterpolation => ({
+    '.box': {
+      width: 93,
+      lineHeight: 1,
+      backgroundColor: token.primaryColor,
+    },
+  });
+
+  interface BoxProps {
+    propToken?: DesignToken;
+  }
+
+  const Box = ({ propToken = baseToken }: BoxProps) => {
+    const [token] = useCacheToken<DerivativeToken>(theme, [propToken]);
+
+    useStyleRegister({ theme, token, path: ['.box'] }, () => [genStyle(token)]);
+
+    return <div className="box" />;
+  };
+
+  // We will not remove style immediately,
+  // but remove when second style patched.
+  describe('remove old style to ensure style set only exist one', () => {
+    function test(
+      name: string,
+      wrapperFn?: (node: React.ReactElement) => React.ReactElement,
+      beforeFn?: () => void,
+    ) {
+      it(name, () => {
+        beforeFn?.();
+
+        const getBox = (props?: BoxProps) => {
+          const box: React.ReactElement = <Box {...props} />;
+
+          return wrapperFn?.(box) || box;
+        };
+
+        const { rerender } = render(getBox());
+        expect(document.head.querySelectorAll('style')).toHaveLength(1);
+
+        // First change
+        rerender(
+          getBox({
+            propToken: {
+              primaryColor: 'red',
+            },
+          }),
+        );
+        expect(document.head.querySelectorAll('style')).toHaveLength(1);
+
+        // Second change
+        rerender(
+          getBox({
+            propToken: {
+              primaryColor: 'green',
+            },
+          }),
+        );
+        expect(document.head.querySelectorAll('style')).toHaveLength(1);
+      });
+    }
+
+    test('normal');
+
+    test('StrictMode', (ele) => <React.StrictMode>{ele}</React.StrictMode>);
+  });
+});