refactor: flattenToken use WeakMap

src/util.ts

@@ -3,19 +3,28 @@ import canUseDom from 'rc-util/lib/Dom/canUseDom';
 import { removeCSS, updateCSS } from 'rc-util/lib/Dom/dynamicCSS';
 import { Theme } from './theme';
 
+// Create a cache here to avoid always loop generate
+const flattenTokenCache = new WeakMap<any, string>();
+
 export function flattenToken(token: any) {
-  let str = '';
-  Object.keys(token).forEach((key) => {
-    const value = token[key];
-    str += key;
-    if (value instanceof Theme) {
-      str += value.id;
-    } else if (value && typeof value === 'object') {
-      str += flattenToken(value);
-    } else {
-      str += value;
-    }
-  });
+  let str = flattenTokenCache.get(token) || '';
+
+  if (!str) {
+    Object.keys(token).forEach((key) => {
+      const value = token[key];
+      str += key;
+      if (value instanceof Theme) {
+        str += value.id;
+      } else if (value && typeof value === 'object') {
+        str += flattenToken(value);
+      } else {
+        str += value;
+      }
+    });
+
+    // Put in cache
+    flattenTokenCache.set(token, str);
+  }
   return str;
 }
 

tests/util.spec.tsx

@@ -1,7 +1,8 @@
 import { normalizeStyle, parseStyle } from '../src/hooks/useStyleRegister';
+import { flattenToken } from '../src/util';
 
-vi.mock('../src/util', () => {
-  const origin = vi.importActual('../src/util');
+vi.mock('../src/util', async () => {
+  const origin: any = await vi.importActual('../src/util');
   return {
     ...origin,
     supportLayer: () => true,
@@ -104,4 +105,25 @@ describe('util', () => {
       });
     });
   });
+
+  it('flattenToken should support cache', () => {
+    const token = {};
+
+    let checkTimes = 0;
+    Object.defineProperty(token, 'a', {
+      get() {
+        checkTimes += 1;
+        return 1;
+      },
+      enumerable: true,
+    });
+
+    // Repeat call flattenToken
+    for (let i = 0; i < 10000; i += 1) {
+      const tokenStr = flattenToken(token);
+      expect(tokenStr).toEqual('a1');
+    }
+
+    expect(checkTimes).toEqual(1);
+  });
 });