fix: react refresh in async entrypoint

crates/rspack_plugin_javascript/src/ast/mod.rs

@@ -4,6 +4,5 @@ mod stringify;
 
 pub use minify::minify;
 pub use parse::parse;
-pub use parse::parse_js_code;
 pub use stringify::print;
 pub use stringify::stringify;

crates/rspack_plugin_javascript/src/ast/parse.rs

@@ -1,4 +1,3 @@
-use std::path::Path;
 use std::sync::Arc;
 
 use rspack_core::{ast::javascript::Ast, ModuleType};
@@ -11,7 +10,7 @@ use swc_core::ecma::parser::{
 };
 use swc_node_comments::SwcComments;
 
-use crate::utils::{ecma_parse_error_to_rspack_error, syntax_by_module_type};
+use crate::utils::ecma_parse_error_to_rspack_error;
 use crate::IsModule;
 
 fn module_type_to_is_module(value: &ModuleType) -> IsModule {
@@ -96,26 +95,3 @@ pub fn parse(
     )),
   }
 }
-
-pub fn parse_js_code(js_code: String, module_type: &ModuleType) -> Result<Program, Error> {
-  let filename = "".to_string();
-  let syntax = syntax_by_module_type(Path::new(&filename), module_type, false);
-  let cm: Arc<swc_core::common::SourceMap> = Default::default();
-  let fm = cm.new_source_file(FileName::Custom(filename), js_code);
-
-  match parse_js(
-    fm.clone(),
-    ast::EsVersion::Es2022,
-    syntax,
-    module_type_to_is_module(module_type),
-    None,
-  ) {
-    Ok(program) => Ok(program),
-    Err(errs) => Err(Error::BatchErrors(
-      errs
-        .into_iter()
-        .map(|err| ecma_parse_error_to_rspack_error(err, &fm, module_type))
-        .collect::<Vec<_>>(),
-    )),
-  }
-}

crates/rspack_plugin_javascript/src/visitors/mod.rs

@@ -93,7 +93,7 @@ pub fn run_before_pass(
         should_transform_by_react
       ),
       Optional::new(
-        swc_visitor::fold_react_refresh(),
+        swc_visitor::fold_react_refresh(unresolved_mark),
         should_transform_by_react && options.builtins.react.refresh.is_some()
       ),
       either!(

crates/rspack_plugin_javascript/src/visitors/swc_visitor/react.rs

@@ -1,15 +1,17 @@
 use std::sync::Arc;
 
-use once_cell::sync::Lazy;
-use rspack_core::{ModuleType, ReactOptions};
+use rspack_core::ReactOptions;
+use swc_core::common::DUMMY_SP;
 use swc_core::common::{comments::SingleThreadedComments, Mark, SourceMap};
-use swc_core::ecma::ast::{CallExpr, Callee, Expr, Ident, ModuleItem, Program, Script};
+use swc_core::ecma::ast::{
+  BinExpr, BinaryOp, BlockStmt, CallExpr, Callee, Expr, ExprOrSpread, ExprStmt, FnDecl, FnExpr,
+  Function, Ident, MemberExpr, ModuleItem, Program, Stmt,
+};
 use swc_core::ecma::transforms::react::RefreshOptions;
 use swc_core::ecma::transforms::react::{react as swc_react, Options};
+use swc_core::ecma::utils::{member_expr, quote_ident, quote_str};
 use swc_core::ecma::visit::{noop_visit_type, Fold, Visit, VisitWith};
 
-use crate::ast::parse_js_code;
-
 pub fn react<'a>(
   top_level_mark: Mark,
   comments: Option<&'a SingleThreadedComments>,
@@ -41,8 +43,8 @@ pub fn react<'a>(
   )
 }
 
-pub fn fold_react_refresh() -> impl Fold {
-  ReactHmrFolder {}
+pub fn fold_react_refresh(unresolved_mark: Mark) -> impl Fold {
+  ReactHmrFolder { unresolved_mark }
 }
 
 #[derive(Default)]
@@ -80,24 +82,146 @@ impl Visit for ReactRefreshUsageFinder {
   }
 }
 
-// __webpack_require__.$ReactRefreshRuntime$ is injected by the react-refresh additional entry
-// See https://github.com/web-infra-dev/rspack/pull/2714 why we have a promise here
-static RUNTIME_CODE: &str = r#"
-function $RefreshReg$(type, id) {
-  __webpack_modules__.$ReactRefreshRuntime$.register(type, __webpack_module__.id+ "_" + id);
+// $ReactRefreshRuntime$ is injected by provide
+fn create_react_refresh_runtime_stmts(unresolved_mark: Mark) -> Vec<Stmt> {
+  fn create_react_refresh_runtime_ident(unresolved_mark: Mark) -> Box<Expr> {
+    Box::new(Expr::Ident(Ident {
+      span: DUMMY_SP.apply_mark(unresolved_mark),
+      sym: "$ReactRefreshRuntime$".into(),
+      optional: false,
+    }))
+  }
+  vec![
+    FnDecl {
+      ident: quote_ident!("$RefreshReg$"),
+      declare: false,
+      function: Box::new(Function {
+        params: vec![quote_ident!("type").into(), quote_ident!("id").into()],
+        decorators: Vec::new(),
+        span: DUMMY_SP,
+        body: Some(BlockStmt {
+          span: DUMMY_SP,
+          stmts: vec![Stmt::Expr(ExprStmt {
+            span: DUMMY_SP,
+            expr: CallExpr {
+              span: DUMMY_SP,
+              callee: Callee::Expr(
+                MemberExpr {
+                  span: DUMMY_SP,
+                  obj: create_react_refresh_runtime_ident(unresolved_mark),
+                  prop: quote_ident!("register").into(),
+                }
+                .into(),
+              ),
+              args: vec![
+                ExprOrSpread {
+                  spread: None,
+                  expr: quote_ident!("type").into(),
+                },
+                ExprOrSpread {
+                  spread: None,
+                  expr: BinExpr {
+                    span: DUMMY_SP,
+                    op: BinaryOp::Add,
+                    left: BinExpr {
+                      span: DUMMY_SP,
+                      op: BinaryOp::Add,
+                      left: member_expr!(DUMMY_SP, __webpack_module__),
+                      right: quote_str!("_").into(),
+                    }
+                    .into(),
+                    right: quote_ident!("id").into(),
+                  }
+                  .into(),
+                },
+              ],
+              type_args: None,
+            }
+            .into(),
+          })],
+        }),
+        is_generator: false,
+        is_async: false,
+        type_params: None,
+        return_type: None,
+      }),
+    }
+    .into(),
+    ExprStmt {
+      span: DUMMY_SP,
+      expr: CallExpr {
+        span: DUMMY_SP,
+        callee: Callee::Expr(
+          MemberExpr {
+            span: DUMMY_SP,
+            obj: CallExpr {
+              span: DUMMY_SP,
+              // See https://github.com/web-infra-dev/rspack/pull/2714 why we have a promise here
+              callee: member_expr!(DUMMY_SP, Promise.resolve).into(),
+              args: vec![],
+              type_args: None,
+            }
+            .into(),
+            prop: quote_ident!("then").into(),
+          }
+          .into(),
+        ),
+        args: vec![ExprOrSpread {
+          spread: None,
+          expr: FnExpr {
+            ident: None,
+            function: Box::new(Function {
+              params: Vec::new(),
+              decorators: Vec::new(),
+              span: DUMMY_SP,
+              body: Some(BlockStmt {
+                span: DUMMY_SP,
+                stmts: vec![Stmt::Expr(ExprStmt {
+                  span: DUMMY_SP,
+                  expr: CallExpr {
+                    span: DUMMY_SP,
+                    callee: Callee::Expr(
+                      MemberExpr {
+                        span: DUMMY_SP,
+                        obj: create_react_refresh_runtime_ident(unresolved_mark),
+                        prop: quote_ident!("refresh").into(),
+                      }
+                      .into(),
+                    ),
+                    args: vec![
+                      ExprOrSpread {
+                        spread: None,
+                        expr: member_expr!(DUMMY_SP, __webpack_module__.id),
+                      },
+                      ExprOrSpread {
+                        spread: None,
+                        expr: member_expr!(DUMMY_SP, __webpack_module__.hot),
+                      },
+                    ],
+                    type_args: None,
+                  }
+                  .into(),
+                })],
+              }),
+              is_generator: false,
+              is_async: false,
+              type_params: None,
+              return_type: None,
+            }),
+          }
+          .into(),
+        }],
+        type_args: None,
+      }
+      .into(),
+    }
+    .into(),
+  ]
 }
-Promise.resolve().then(function(){
-  __webpack_modules__.$ReactRefreshRuntime$.refresh(__webpack_module__.id, __webpack_module__.hot);
-})
-"#;
-
-static RUNTIME_CODE_AST: Lazy<Script> = Lazy::new(|| {
-  parse_js_code(RUNTIME_CODE.to_string(), &ModuleType::Js)
-    .expect("TODO:")
-    .expect_script()
-});
 
-pub struct ReactHmrFolder;
+pub struct ReactHmrFolder {
+  unresolved_mark: Mark,
+}
 
 impl Fold for ReactHmrFolder {
   fn fold_program(&mut self, mut program: Program) -> Program {
@@ -108,7 +232,7 @@ impl Fold for ReactHmrFolder {
       return program;
     }
 
-    let runtime_stmts = RUNTIME_CODE_AST.body.clone();
+    let runtime_stmts = create_react_refresh_runtime_stmts(self.unresolved_mark);
 
     match program {
       Program::Module(ref mut m) => m

packages/playground/cases/react/worker/index.test.ts

@@ -0,0 +1,13 @@
+import { test, expect } from "@/fixtures";
+
+test("render should work", async ({ page }) => {
+	expect(await page.textContent("button")).toBe("+");
+});
+
+test("worker should work", async ({ page, fileAction, rspack }) => {
+	expect(await page.textContent("h1")).toBe("0");
+	await page.click("button");
+	expect(await page.textContent("h1")).toBe("1");
+	await page.click("button");
+	expect(await page.textContent("h1")).toBe("2");
+});

packages/playground/cases/react/worker/rspack.config.js

@@ -0,0 +1,24 @@
+/** @type { import('@rspack/core').RspackOptions } */
+module.exports = {
+	context: __dirname,
+	mode: "development",
+	entry: "./src/index.jsx",
+	devServer: {
+		hot: true
+	},
+	cache: false,
+	stats: "none",
+	infrastructureLogging: {
+		debug: false
+	},
+	builtins: {
+		html: [
+			{
+				template: "./src/index.html"
+			}
+		]
+	},
+	watchOptions: {
+		poll: 1000
+	}
+};

packages/playground/cases/react/worker/src/App.jsx

@@ -0,0 +1,21 @@
+import React, { useEffect, useReducer, useState } from "react";
+import Button from './Button'
+
+let updateRenderTimes;
+
+const worker = new Worker(new URL("./worker", import.meta.url));
+
+worker.onmessage = (e) => {
+	updateRenderTimes(e.data)
+}
+
+export const App = () => {
+	const [renderTimes, setRenderTimes] = useState(0)
+	updateRenderTimes = setRenderTimes;
+	return (
+		<div className="App">
+			<h1>{renderTimes}</h1>
+			<Button onClick={() => worker.postMessage('add')} />
+		</div>
+	);
+};

packages/playground/cases/react/worker/src/Button.jsx

@@ -0,0 +1,15 @@
+import React from "react";
+
+export default function Button({ onClick }) {
+	return <button onClick={onClick}>+</button>;
+};
+
+Button.count = 0;
+
+Button.get = () => {
+  return Button.count;
+}
+
+Button.add = () => {
+  Button.count += 1;
+}

packages/playground/cases/react/worker/src/index.html

@@ -0,0 +1,12 @@
+<!DOCTYPE html>
+<html lang="en">
+	<head>
+		<meta charset="UTF-8" />
+		<meta http-equiv="X-UA-Compatible" content="IE=edge" />
+		<meta name="viewport" content="width=device-width, initial-scale=1.0" />
+		<title>Document</title>
+	</head>
+	<body>
+		<div id="root"></div>
+	</body>
+</html>

packages/playground/cases/react/worker/src/index.jsx

@@ -0,0 +1,6 @@
+import React from "react";
+import { createRoot } from "react-dom/client";
+import { App } from './App';
+
+const container = createRoot(document.getElementById("root"));
+container.render(<App />);

packages/playground/cases/react/worker/src/worker.js

@@ -0,0 +1,6 @@
+import Button from "./Button";
+
+onmessage = e => {
+	Button.add();
+	postMessage(Button.get());
+};

packages/rspack-dev-client/package.json

@@ -4,12 +4,12 @@
   "license": "MIT",
   "description": "Development client for rspack",
   "scripts": {
-    "build": "tsc",
-    "dev": "tsc -w",
+    "build": "echo success",
+    "dev": "echo success",
     "test": "echo success"
   },
   "files": [
-    "dist"
+    "src"
   ],
   "homepage": "https://rspack.dev",
   "bugs": "https://github.com/web-infra-dev/rspack/issues",
@@ -34,6 +34,7 @@
     }
   },
   "exports": {
-    "./react-refresh": "./dist/reactRefresh.js"
+    "./react-refresh": "./src/reactRefresh.js",
+    "./react-refresh-entry": "./src/reactRefreshEntry.js"
   }
 }