perf: use cleanup register with deps

ant-design · Aug 10, 2023 · 91fcc34 · 91fcc34
1 parent 4963179
commit 91fcc34
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 10 deletions.
diff --git a/src/hooks/useCacheToken.tsx b/src/hooks/useCacheToken.tsx
@@ -4,7 +4,6 @@ import { useContext } from 'react';
 import StyleContext, { ATTR_TOKEN, CSS_IN_JS_INSTANCE } from '../StyleContext';
 import type Theme from '../theme/Theme';
 import { flattenToken, token2key } from '../util';
-import useEffectCleanupRegister from './useEffectCleanupRegister';
 import useGlobalCache from './useGlobalCache';
 
 const EMPTY_OVERRIDE = {};
@@ -137,8 +136,6 @@ export default function useCacheToken<
     getComputedToken: compute,
   } = option;
 
-  const register = useEffectCleanupRegister();
-
   // Basic - We do basic cache here
   const mergedToken = React.useMemo(
     () => Object.assign({}, ...tokens),
@@ -175,10 +172,7 @@ export default function useCacheToken<
     },
     (cache) => {
       // Remove token will remove all related style
-      // Always remove styles in useEffect callback
-      register(() => {
-        cleanTokenStyle(cache[0]._tokenKey, instanceId);
-      });
+      cleanTokenStyle(cache[0]._tokenKey, instanceId);
     },
   );
 

diff --git a/src/hooks/useEffectCleanupRegister.ts b/src/hooks/useEffectCleanupRegister.ts
@@ -7,7 +7,7 @@ const fullClone = {
 const { useInsertionEffect } = fullClone;
 
 // DO NOT register functions in useEffect cleanup function, or functions that registered will never be called.
-const useCleanupRegister = () => {
+const useCleanupRegister = (deps?: React.DependencyList) => {
   const effectCleanups: (() => void)[] = [];
   let cleanupFlag = false;
   function register(fn: () => void) {
@@ -32,7 +32,7 @@ const useCleanupRegister = () => {
         effectCleanups.forEach((fn) => fn());
       }
     };
-  });
+  }, deps);
 
   return register;
 };

diff --git a/src/hooks/useGlobalCache.tsx b/src/hooks/useGlobalCache.tsx
@@ -2,6 +2,7 @@ import * as React from 'react';
 import type { KeyType } from '../Cache';
 import StyleContext from '../StyleContext';
 import useCompatibleInsertionEffect from './useCompatibleInsertionEffect';
+import useEffectCleanupRegister from './useEffectCleanupRegister';
 import useHMR from './useHMR';
 
 export default function useGlobalCache<CacheType>(
@@ -16,6 +17,8 @@ export default function useGlobalCache<CacheType>(
   const fullPath = [prefix, ...keyPath];
   const deps = fullPath.join('_');
 
+  const register = useEffectCleanupRegister([deps]);
+
   const HMRUpdate = useHMR();
 
   type UpdaterArgs = [times: number, cache: CacheType];
@@ -84,7 +87,8 @@ export default function useGlobalCache<CacheType>(
           const nextCount = times - 1;
 
           if (nextCount === 0) {
-            onCacheRemove?.(cache, false);
+            // Always remove styles in useEffect callback
+            register(() => onCacheRemove?.(cache, false));
             return null;
           }