feat: extractStyle support once (#228)

Author: MadCcc

src/Cache.ts

@@ -17,6 +17,8 @@ class Entity {
   /** @private Internal cache map. Do not access this directly */
   cache = new Map<string, ValueType>();
 
+  extracted: Set<string> = new Set();
+
   get(keys: KeyType[]): ValueType | null {
     return this.opGet(pathKey(keys));
   }

src/extractStyle.ts

@@ -36,9 +36,10 @@ export default function extractStyle(
     | {
         plain?: boolean;
         types?: ExtractStyleType | ExtractStyleType[];
+        once?: boolean;
       },
 ) {
-  const { plain = false, types = ['style', 'token', 'cssVar'] } =
+  const { plain = false, types = ['style', 'token', 'cssVar'], once = false } =
     typeof options === 'boolean' ? { plain: options } : options || {};
 
   const matchPrefixRegexp = new RegExp(
@@ -60,6 +61,9 @@ export default function extractStyle(
 
   styleKeys
     .map<[number, string] | null>((key) => {
+      if (once && cache.extracted.has(key)) {
+        return null; // Skip if already extracted
+      }
       const cachePath = key.replace(matchPrefixRegexp, '').replace(/%/g, '|');
       const [prefix] = key.split('%');
       const extractFn = ExtractStyleFns[prefix as keyof typeof ExtractStyleFns];
@@ -73,6 +77,10 @@ export default function extractStyle(
       if (key.startsWith('style')) {
         cachePathMap[cachePath] = styleId;
       }
+
+      // record that this style has been extracted
+      cache.extracted.add(key);
+
       return [order, styleStr];
     })
     .filter(isNotNull)

tests/server.spec.tsx

@@ -69,6 +69,12 @@ describe('SSR', () => {
     },
   });
 
+  const genCardStyle = (token: DerivativeToken): CSSInterpolation => ({
+    '.card': {
+      backgroundColor: token.primaryColor,
+    },
+  });
+
   const Box = ({ children }: { children?: React.ReactNode }) => {
     const [token] = useCacheToken<DerivativeToken>(theme, [baseToken]);
 
@@ -79,6 +85,16 @@ describe('SSR', () => {
     return wrapSSR(<div className="box">{children}</div>);
   };
 
+  const Card = ({ children }: { children?: React.ReactNode }) => {
+    const [token] = useCacheToken<DerivativeToken>(theme, [baseToken], {
+      cssVar: { key: 'css-var-test' },
+    });
+
+    useStyleRegister({ theme, token, path: ['.card'] }, () => [genCardStyle(token)]);
+
+    return <div className="card">{children}</div>;
+  };
+
   const IdHolder = () => {
     const id = React.useId();
     return (
@@ -441,4 +457,33 @@ describe('SSR', () => {
       ].join(''),
     );
   });
+
+  it('should extract once when once option is true', () => {
+    const cache = createCache();
+
+    renderToString(
+      <StyleProvider cache={cache}>
+        <IdHolder />
+        <Box>
+          <IdHolder />
+        </Box>
+        <IdHolder />
+      </StyleProvider>,
+    );
+
+    const style = extractStyle(cache, {plain: true, once: true});
+
+    renderToString(
+      <StyleProvider cache={cache}>
+        <Card />
+      </StyleProvider>,
+    );
+    const style2 = extractStyle(cache, {plain: true, once: true});
+
+    expect(style).toContain('.box');
+    expect(style).not.toContain('.card');
+
+    expect(style2).toContain('.card');
+    expect(style2).not.toContain('.box');
+  })
 });