From 7266ab0383097a9ec4ca238a838063685aeb1540 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E9=A9=AC=E6=99=A8=E6=9D=B0?= <309272985@qq.com>
Date: Tue, 10 Jan 2023 14:46:58 +0800
Subject: [PATCH] feat: extract style node to support Next.js _document.tsx
 render style tags

---
 src/hooks/useStyleRegister.tsx | 22 ++++++++++++++++++
 src/index.ts                   |  6 ++++-
 tests/server.spec.tsx          | 42 +++++++++++++++++++++++++++-------
 3 files changed, 61 insertions(+), 9 deletions(-)

diff --git a/src/hooks/useStyleRegister.tsx b/src/hooks/useStyleRegister.tsx
index 0dc2d1f9..ce8de15d 100644
--- a/src/hooks/useStyleRegister.tsx
+++ b/src/hooks/useStyleRegister.tsx
@@ -483,3 +483,25 @@ export function extractStyle(cache: Cache) {
 
   return styleText;
 }
+
+export function extractStyleNode(cache: Cache) {
+  const styleKeys = Array.from(cache.cache.keys()).filter((key) =>
+    key.startsWith('style%'),
+  );
+
+  return styleKeys.map((key) => {
+    const [styleStr, tokenKey, styleId]: [string, string, string] =
+      cache.cache.get(key)![1];
+
+    return (
+      <style
+        key={key}
+        {...{
+          [ATTR_TOKEN]: tokenKey,
+          [ATTR_MARK]: styleId,
+        }}
+        dangerouslySetInnerHTML={{ __html: styleStr }}
+      />
+    );
+  });
+}
diff --git a/src/index.ts b/src/index.ts
index 6a605bcf..a98f859e 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -1,6 +1,9 @@
 import useCacheToken from './hooks/useCacheToken';
 import type { CSSInterpolation, CSSObject } from './hooks/useStyleRegister';
-import useStyleRegister, { extractStyle } from './hooks/useStyleRegister';
+import useStyleRegister, {
+  extractStyle,
+  extractStyleNode,
+} from './hooks/useStyleRegister';
 import Keyframes from './Keyframes';
 import type { Linter } from './linters';
 import { logicalPropertiesLinter } from './linters';
@@ -19,6 +22,7 @@ export {
   StyleProvider,
   Keyframes,
   extractStyle,
+  extractStyleNode,
 
   // Transformer
   legacyLogicalPropertiesTransformer,
diff --git a/tests/server.spec.tsx b/tests/server.spec.tsx
index 56016917..1f27e0b5 100644
--- a/tests/server.spec.tsx
+++ b/tests/server.spec.tsx
@@ -1,23 +1,23 @@
+import { render } from '@testing-library/react';
 import * as React from 'react';
 import { renderToString } from 'react-dom/server';
-import { render } from '@testing-library/react';
+import type { CSSInterpolation } from '../src';
 import {
+  createCache,
+  extractStyle,
+  extractStyleNode,
+  StyleProvider,
   Theme,
   useCacheToken,
   useStyleRegister,
-  StyleProvider,
-  extractStyle,
-  createCache,
 } from '../src';
-import type { CSSInterpolation } from '../src';
 // eslint-disable-next-line @typescript-eslint/no-unused-vars
-import canUseDom from 'rc-util/lib/Dom/canUseDom';
+import classNames from 'classnames';
 import {
+  ATTR_MARK,
   CSS_IN_JS_INSTANCE,
   CSS_IN_JS_INSTANCE_ID,
-  ATTR_MARK,
 } from '../src/StyleContext';
-import classNames from 'classnames';
 
 interface DesignToken {
   primaryColor: string;
@@ -156,6 +156,32 @@ describe('SSR', () => {
     expect(errorSpy).not.toHaveBeenCalled();
   });
 
+  it('ssr extract style node', () => {
+    // >>> SSR
+    const cache = createCache();
+
+    const html = renderToString(
+      <StyleProvider cache={cache}>
+        <IdHolder />
+        <Box>
+          <IdHolder />
+        </Box>
+        <IdHolder />
+      </StyleProvider>,
+    );
+
+    const styleNode = extractStyleNode(cache);
+    const styleText = extractStyle(cache);
+
+    const style = renderToString(<>{styleNode}</>);
+
+    expect(html).toEqual(
+      '<div id=":R1:" class="id">:R1:</div><div class="box"><div id=":Ra:" class="id">:Ra:</div></div><div id=":R3:" class="id">:R3:</div>',
+    );
+
+    expect(style).toEqual(styleText);
+  });
+
   it('default hashPriority', () => {
     // >>> SSR
     const cache = createCache();