feat: css var support ssr (#150)

* feat: css var support ssr

* chore: code clean

* chore: code clean

* chore: fix test

src/extractStyle.ts

@@ -0,0 +1,81 @@
+import type Cache from './Cache';
+import {
+  extract as tokenExtractStyle,
+  TOKEN_PREFIX,
+} from './hooks/useCacheToken';
+import {
+  CSS_VAR_PREFIX,
+  extract as cssVarExtractStyle,
+} from './hooks/useCSSVarRegister';
+import {
+  extract as styleExtractStyle,
+  STYLE_PREFIX,
+} from './hooks/useStyleRegister';
+import { toStyleStr } from './util';
+import {
+  ATTR_CACHE_MAP,
+  serialize as serializeCacheMap,
+} from './util/cacheMapUtil';
+
+const ExtractStyleFns = {
+  [STYLE_PREFIX]: styleExtractStyle,
+  [TOKEN_PREFIX]: tokenExtractStyle,
+  [CSS_VAR_PREFIX]: cssVarExtractStyle,
+};
+
+function isNotNull<T>(value: T | null): value is T {
+  return value !== null;
+}
+
+export default function extractStyle(cache: Cache, plain = false) {
+  const matchPrefixRegexp = new RegExp(
+    `^(${Object.keys(ExtractStyleFns).join('|')})%`,
+  );
+
+  // prefix with `style` is used for `useStyleRegister` to cache style context
+  const styleKeys = Array.from(cache.cache.keys()).filter((key) =>
+    matchPrefixRegexp.test(key),
+  );
+
+  // Common effect styles like animation
+  const effectStyles: Record<string, boolean> = {};
+
+  // Mapping of cachePath to style hash
+  const cachePathMap: Record<string, string> = {};
+
+  let styleText = '';
+
+  styleKeys
+    .map<[number, string] | null>((key) => {
+      const cachePath = key.replace(matchPrefixRegexp, '').replace(/%/g, '|');
+      const [prefix] = key.split('%');
+      const extractFn = ExtractStyleFns[prefix as keyof typeof ExtractStyleFns];
+      const extractedStyle = extractFn(cache.cache.get(key)![1], effectStyles, {
+        plain,
+      });
+      if (!extractedStyle) {
+        return null;
+      }
+      const [order, styleId, styleStr] = extractedStyle;
+      cachePathMap[cachePath] = styleId;
+      return [order, styleStr];
+    })
+    .filter(isNotNull)
+    .sort(([o1], [o2]) => o1 - o2)
+    .forEach(([, style]) => {
+      styleText += style;
+    });
+
+  // ==================== Fill Cache Path ====================
+  styleText += toStyleStr(
+    `.${ATTR_CACHE_MAP}{content:"${serializeCacheMap(cachePathMap)}";}`,
+    undefined,
+    undefined,
+    {
+      [ATTR_CACHE_MAP]: ATTR_CACHE_MAP,
+    },
+    plain,
+  );
+
+  return styleText;
+}

src/hooks/useCSSVarRegister.ts

@@ -1,15 +1,25 @@
-import hash from '@emotion/hash';
 import { removeCSS, updateCSS } from 'rc-util/lib/Dom/dynamicCSS';
 import { useContext } from 'react';
 import StyleContext, {
   ATTR_MARK,
   ATTR_TOKEN,
   CSS_IN_JS_INSTANCE,
 } from '../StyleContext';
-import { isClientSide } from '../util';
+import { isClientSide, toStyleStr } from '../util';
 import type { TokenWithCSSVar } from '../util/css-variables';
 import { transformToken } from '../util/css-variables';
+import type { ExtractStyle } from './useGlobalCache';
 import useGlobalCache from './useGlobalCache';
+import { uniqueHash } from './useStyleRegister';
+
+export const CSS_VAR_PREFIX = 'cssVar';
+
+type CSSVarCacheValue<T> = [
+  cssVarToken: TokenWithCSSVar<T>,
+  cssVarStr: string,
+  styleId: string,
+  cssVarKey: string,
+];
 
 const useCSSVarRegister = <V, T extends Record<string, V>>(
   config: {
@@ -29,24 +39,24 @@ const useCSSVarRegister = <V, T extends Record<string, V>>(
   } = useContext(StyleContext);
   const { _tokenKey: tokenKey } = token;
 
-  const cache = useGlobalCache<[TokenWithCSSVar<T>, string, T, string]>(
-    'variables',
+  const cache = useGlobalCache<CSSVarCacheValue<T>>(
+    CSS_VAR_PREFIX,
     [...config.path, key, tokenKey],
     () => {
-      const styleId = hash([...config.path, key].join('%'));
       const originToken = fn();
       const [mergedToken, cssVarsStr] = transformToken(originToken, key, {
         prefix,
         unitless,
       });
-      return [mergedToken, cssVarsStr, originToken, styleId];
+      const styleId = uniqueHash([...config.path, key], cssVarsStr);
+      return [mergedToken, cssVarsStr, styleId, key];
     },
-    ([, , , styleId], fromHMR) => {
+    ([, , styleId], fromHMR) => {
       if ((fromHMR || autoClear) && isClientSide) {
         removeCSS(styleId, { mark: ATTR_MARK });
       }
     },
-    ([, cssVarsStr, , styleId]) => {
+    ([, cssVarsStr, styleId]) => {
       if (!cssVarsStr) {
         return;
       }
@@ -67,4 +77,36 @@ const useCSSVarRegister = <V, T extends Record<string, V>>(
   return cache;
 };
 
+export const extract: ExtractStyle<CSSVarCacheValue<any>> = (
+  cache,
+  effectStyles,
+  options,
+) => {
+  const [, styleStr, styleId, cssVarKey] = cache;
+  const { plain } = options || {};
+
+  if (!styleStr) {
+    return null;
+  }
+
+  const order = -999;
+
+  // ====================== Style ======================
+  // Used for rc-util
+  const sharedAttrs = {
+    'data-rc-order': 'prependQueue',
+    'data-rc-priority': `${order}`,
+  };
+
+  const styleText = toStyleStr(
+    styleStr,
+    cssVarKey,
+    styleId,
+    sharedAttrs,
+    plain,
+  );
+
+  return [order, styleId, styleText];
+};
+
 export default useCSSVarRegister;

src/hooks/useCacheToken.tsx

@@ -7,8 +7,9 @@ import StyleContext, {
   CSS_IN_JS_INSTANCE,
 } from '../StyleContext';
 import type Theme from '../theme/Theme';
-import { flattenToken, memoResult, token2key } from '../util';
+import { flattenToken, memoResult, token2key, toStyleStr } from '../util';
 import { transformToken } from '../util/css-variables';
+import type { ExtractStyle } from './useGlobalCache';
 import useGlobalCache from './useGlobalCache';
 
 const EMPTY_OVERRIDE = {};
@@ -130,6 +131,16 @@ export const getComputedToken = <
   return mergedDerivativeToken;
 };
 
+export const TOKEN_PREFIX = 'token';
+
+type TokenCacheValue<DerivativeToken> = [
+  token: DerivativeToken & { _tokenKey: string; _themeKey: string },
+  hashId: string,
+  realToken: DerivativeToken & { _tokenKey: string },
+  cssVarStr: string,
+  cssVarKey: string,
+];
+
 /**
  * Cache theme derivative token as global shared one
  * @param theme Theme entity
@@ -144,12 +155,7 @@ export default function useCacheToken<
   theme: Theme<any, any>,
   tokens: Partial<DesignToken>[],
   option: Option<DerivativeToken, DesignToken> = {},
-): [
-  DerivativeToken & { _tokenKey: string; _themeKey: string },
-  string,
-  DerivativeToken,
-  string,
-] {
+): TokenCacheValue<DerivativeToken> {
   const {
     cache: { instanceId },
     container,
@@ -170,15 +176,8 @@ export default function useCacheToken<
 
   const cssVarStr = cssVar ? flattenToken(cssVar) : '';
 
-  const cachedToken = useGlobalCache<
-    [
-      DerivativeToken & { _tokenKey: string; _themeKey: string },
-      string,
-      DerivativeToken,
-      string,
-    ]
-  >(
-    'token',
+  const cachedToken = useGlobalCache<TokenCacheValue<DerivativeToken>>(
+    TOKEN_PREFIX,
     [salt, theme.id, tokenStr, overrideTokenStr, cssVarStr],
     () => {
       let mergedDerivativeToken = compute
@@ -214,7 +213,13 @@ export default function useCacheToken<
         : `${hashPrefix}-${hash(tokenKey)}`;
       mergedDerivativeToken._hashId = hashId; // Not used
 
-      return [mergedDerivativeToken, hashId, actualToken, cssVarsStr];
+      return [
+        mergedDerivativeToken,
+        hashId,
+        actualToken,
+        cssVarsStr,
+        cssVar?.key || '',
+      ];
     },
     (cache) => {
       // Remove token will remove all related style
@@ -243,3 +248,36 @@ export default function useCacheToken<
 
   return cachedToken;
 }
+
+export const extract: ExtractStyle<TokenCacheValue<any>> = (
+  cache,
+  effectStyles,
+  options,
+) => {
+  const [, , realToken, styleStr, cssVarKey] = cache;
+  const { plain } = options || {};
+
+  if (!styleStr) {
+    return null;
+  }
+
+  const styleId = realToken._tokenKey;
+  const order = -999;
+
+  // ====================== Style ======================
+  // Used for rc-util
+  const sharedAttrs = {
+    'data-rc-order': 'prependQueue',
+    'data-rc-priority': `${order}`,
+  };
+
+  const styleText = toStyleStr(
+    styleStr,
+    cssVarKey,
+    styleId,
+    sharedAttrs,
+    plain,
+  );
+
+  return [order, styleId, styleText];
+};

src/hooks/useGlobalCache.tsx

@@ -5,6 +5,14 @@ import useCompatibleInsertionEffect from './useCompatibleInsertionEffect';
 import useEffectCleanupRegister from './useEffectCleanupRegister';
 import useHMR from './useHMR';
 
+export type ExtractStyle<CacheValue> = (
+  cache: CacheValue,
+  effectStyles: Record<string, boolean>,
+  options?: {
+    plain?: boolean;
+  },
+) => [order: number, styleId: string, style: string] | null;
+
 export default function useGlobalCache<CacheType>(
   prefix: string,
   keyPath: KeyType[],
@@ -25,7 +33,7 @@ export default function useGlobalCache<CacheType>(
 
   const buildCache = (updater?: (data: UpdaterArgs) => UpdaterArgs) => {
     globalCache.update(fullPath, (prevCache) => {
-      const [times = 0, cache] = prevCache || [];
+      const [times = 0, cache] = prevCache || [undefined, undefined];
 
       // HMR should always ignore cache since developer may change it
       let tmpCache = cache;