From fc38fd3cd475153d79963649e064147108747de7 Mon Sep 17 00:00:00 2001
From: Hiroshi Ogawa <hi.ogawa.zz@gmail.com>
Date: Tue, 12 Nov 2024 17:29:24 +0900
Subject: [PATCH] refactor(viteroll): inline `react-refresh/runtime` (#72)

---
 .github/workflows/ci.yml                   |   1 +
 viteroll/examples/ssr/src/entry-client.tsx |   2 -
 viteroll/viteroll.ts                       | 116 ++++++++++-----------
 3 files changed, 54 insertions(+), 65 deletions(-)

diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index 71b2646..f2f451e 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -117,3 +117,4 @@ jobs:
       - run: pnpm test-e2e
       - run: pnpm -C examples/react test-e2e
       - run: pnpm -C examples/mpa test-e2e
+      - run: pnpm -C examples/ssr test-e2e
diff --git a/viteroll/examples/ssr/src/entry-client.tsx b/viteroll/examples/ssr/src/entry-client.tsx
index f50a730..3e5c48f 100644
--- a/viteroll/examples/ssr/src/entry-client.tsx
+++ b/viteroll/examples/ssr/src/entry-client.tsx
@@ -1,5 +1,3 @@
-// TODO: how to inject automatically?
-import "virtual:react-refresh/entry";
 import React from "react";
 import ReactDOMClient from "react-dom/client";
 import { App } from "./app";
diff --git a/viteroll/viteroll.ts b/viteroll/viteroll.ts
index e785301..dc142b5 100644
--- a/viteroll/viteroll.ts
+++ b/viteroll/viteroll.ts
@@ -63,8 +63,10 @@ export function viteroll(viterollOptions: ViterollOptions = {}): Plugin {
 					},
 					ssr: {
 						dev: {
-							createEnvironment:
-								RolldownEnvironment.createFactory(viterollOptions),
+							createEnvironment: RolldownEnvironment.createFactory({
+								...viterollOptions,
+								reactRefresh: false,
+							}),
 						},
 					},
 				},
@@ -315,9 +317,6 @@ function viterollEntryPlugin(
 				htmlEntryMap.set(id, htmlOutput);
 
 				let jsOutput = ``;
-				if (viterollOptions?.reactRefresh) {
-					jsOutput += `import "virtual:react-refresh/entry";\n`;
-				}
 
 				// extract <script src="...">
 				const matches = code.matchAll(
@@ -360,6 +359,9 @@ function viterollEntryPlugin(
 						}
 						for (var i = 0; i < module.parents.length; i++) {`,
 					);
+				if (viterollOptions.reactRefresh) {
+					output.prepend(getReactRefreshRuntimeCode());
+				}
 				return { code: output.toString(), map: output.generateMap() };
 			}
 		},
@@ -396,7 +398,6 @@ function viterollEntryPlugin(
 	};
 }
 
-// TODO: workaround rolldownExperimental.reactPlugin which injects js to html via `load` hook
 function reactRefreshPlugin(): rolldown.Plugin {
 	return {
 		name: "react-hmr",
@@ -407,65 +408,54 @@ function reactRefreshPlugin(): rolldown.Plugin {
 				},
 			},
 			handler(code, id) {
-				const output = new MagicString(code);
-				output.prepend(`
-					import * as __$refresh from 'virtual:react-refresh';
-					const [$RefreshSig$, $RefreshReg$] = __$refresh.create(${JSON.stringify(id)});
-				`);
-				output.append(`
-					__$refresh.setupHot(module.hot);
-				`);
-				return { code: output.toString(), map: output.generateMap() };
+				return [
+					`const [$RefreshSig$, $RefreshReg$] = __react_refresh_transform_define(${JSON.stringify(id)})`,
+					code,
+					`__react_refresh_transform_setupHot(module.hot)`,
+				].join(";");
 			},
 		},
-		resolveId: {
-			filter: {
-				id: {
-					include: [/^virtual:react-refresh/],
-				},
-			},
-			handler: (source) => "\0" + source,
-		},
-		load: {
-			filter: {
-				id: {
-					include: [/^\0virtual:react-refresh/],
-				},
-			},
-			async handler(id) {
-				const resolved = require.resolve("react-refresh/runtime");
-				if (id === "\0virtual:react-refresh/entry") {
-					return `
-						import runtime from ${JSON.stringify(resolved)};
-						runtime.injectIntoGlobalHook(window);
-					`;
-				}
-				if (id === "\0virtual:react-refresh") {
-					return `
-						import runtime from ${JSON.stringify(resolved)};
-
-						export const create = (file) => [
-							runtime.createSignatureFunctionForTransform,
-							(type, id) => runtime.register(type, file + '_' + id),
-						];
-
-						function debounce(fn, delay) {
-							let handle
-							return () => {
-								clearTimeout(handle)
-								handle = setTimeout(fn, delay)
-							}
-						}
-						const debouncedRefresh = debounce(runtime.performReactRefresh, 16);
+	};
+}
 
-						export function setupHot(hot) {
-							hot.accept((prev) => {
-								debouncedRefresh();
-							});
-						}
-					`;
+// inject react refresh runtime in client runtime to ensure initialized early
+function getReactRefreshRuntimeCode() {
+	let code = fs.readFileSync(
+		path.resolve(
+			require.resolve("react-refresh/runtime"),
+			"..",
+			"cjs/react-refresh-runtime.development.js",
+		),
+		"utf-8",
+	);
+	const output = new MagicString(code);
+	output.prepend("self.__react_refresh_runtime = {};\n");
+	output.replaceAll('process.env.NODE_ENV !== "production"', "true");
+	output.replaceAll(/\bexports\./g, "__react_refresh_runtime.");
+	output.append(`
+		(() => {
+			__react_refresh_runtime.injectIntoGlobalHook(self);
+
+			__react_refresh_transform_define = (file) => [
+				__react_refresh_runtime.createSignatureFunctionForTransform,
+				(type, id) => __react_refresh_runtime.register(type, file + '_' + id)
+			];
+
+			__react_refresh_transform_setupHot = (hot) => {
+				hot.accept((prev) => {
+					debouncedRefresh();
+				});
+			};
+
+			function debounce(fn, delay) {
+				let handle
+				return () => {
+					clearTimeout(handle)
+					handle = setTimeout(fn, delay)
 				}
-			},
-		},
-	};
+			}
+			const debouncedRefresh = debounce(__react_refresh_runtime.performReactRefresh, 16);
+		})()
+	`);
+	return output.toString();
 }