refactor: improve runtime module macro (#12941)

* refactor: attach chunk ukey to runtime module in macro

* refactor: generate runtime module identifier in macro

* refactor: generate runtime module identifier in macro

* refactor: generate runtime module identifier in macro

* refactor: generate runtime module identifier in macro

* refactor: generate runtime module identifier in macro

Author: LingyuCoder

Cargo.lock

@@ -772,6 +772,7 @@ pub struct JsAddingRuntimeModule {
 impl From<JsAddingRuntimeModule> for RuntimeModuleFromJs {
   fn from(value: JsAddingRuntimeModule) -> Self {
     Self {
+      chunk: None,
       id: Identifier::from(value.name),
       full_hash: value.full_hash,
       dependent_hash: value.dependent_hash,

crates/rspack_binding_api/src/module.rs

@@ -19,6 +19,7 @@ dyn-clone = { workspace = true }
 either = { workspace = true }
 futures = { workspace = true }
 hashlink = { workspace = true }
+heck = { workspace = true }
 hex = { workspace = true }
 indexmap = { workspace = true, features = ["rayon"] }
 indoc = { workspace = true }

crates/rspack_core/Cargo.toml

@@ -7,9 +7,9 @@ use rspack_collections::Identifier;
 use crate::{ChunkUkey, Compilation, Module, RuntimeGlobals};
 
 #[async_trait]
-pub trait RuntimeModule: Module + CustomSourceRuntimeModule {
-  fn name(&self) -> Identifier;
-  fn attach(&mut self, _chunk: ChunkUkey) {}
+pub trait RuntimeModule:
+  Module + CustomSourceRuntimeModule + AttachableRuntimeModule + NamedRuntimeModule
+{
   fn stage(&self) -> RuntimeModuleStage {
     RuntimeModuleStage::Normal
   }
@@ -39,6 +39,16 @@ pub trait RuntimeModule: Module + CustomSourceRuntimeModule {
   }
 }
 
+#[async_trait]
+pub trait AttachableRuntimeModule {
+  fn attach(&mut self, chunk: ChunkUkey);
+}
+
+#[async_trait]
+pub trait NamedRuntimeModule {
+  fn name(&self) -> Identifier;
+}
+
 #[async_trait]
 pub trait CustomSourceRuntimeModule {
   fn set_custom_source(&mut self, source: String);

crates/rspack_core/src/runtime_module.rs

@@ -4,6 +4,7 @@ use std::{
 };
 
 use cow_utils::CowUtils;
+use heck::ToSnakeCase;
 use itertools::Itertools;
 use regex::{Captures, Regex};
 use rspack_collections::Identifier;
@@ -229,6 +230,23 @@ impl RuntimeTemplate {
     "webpack/runtime/"
   }
 
+  pub fn create_runtime_module_identifier(&self, name: &str) -> Identifier {
+    let module_name = if let Some(name) = name.strip_suffix("RuntimeModule") {
+      name
+    } else {
+      name
+    };
+    Identifier::from(format!(
+      "{}{}",
+      self.runtime_module_prefix(),
+      module_name.to_snake_case()
+    ))
+  }
+
+  pub fn create_custom_runtime_module_identifier(&self, custom: &str) -> Identifier {
+    Identifier::from(format!("{}{custom}", self.runtime_module_prefix()))
+  }
+
   pub fn create_module_codegen_runtime_template(&self) -> ModuleCodegenRuntimeTemplate {
     ModuleCodegenRuntimeTemplate::new(self.compiler_options.clone())
   }

crates/rspack_core/src/runtime_template.rs

@@ -12,6 +12,16 @@ pub fn impl_runtime_module(
   let origin_fields = input.fields.clone();
 
   if let syn::Fields::Named(ref mut fields) = input.fields {
+    fields.named.push(
+      syn::Field::parse_named
+        .parse2(quote! { pub id: ::rspack_collections::Identifier })
+        .expect("Failed to parse new field for id"),
+    );
+    fields.named.push(
+      syn::Field::parse_named
+        .parse2(quote! { pub chunk: Option<::rspack_core::ChunkUkey> })
+        .expect("Failed to parse new field for chunk"),
+    );
     fields.named.push(
       syn::Field::parse_named
         .parse2(quote! { pub source_map_kind: ::rspack_util::source_map::SourceMapKind })
@@ -41,23 +51,40 @@ pub fn impl_runtime_module(
   let field_tys: Vec<&syn::Type> = origin_fields.iter().map(|field| &field.ty).collect();
   let with_default = quote! {
     #[allow(clippy::too_many_arguments)]
-    fn with_default(#(#field_names: #field_tys,)*) -> Self {
+    fn with_default(runtime_template: &::rspack_core::RuntimeTemplate, #(#field_names: #field_tys,)*) -> Self {
       Self {
         source_map_kind: ::rspack_util::source_map::SourceMapKind::empty(),
         custom_source: None,
         cached_generated_code: Default::default(),
+        chunk: None,
+        id: runtime_template.create_runtime_module_identifier(stringify!(#name)),
         #(#field_names,)*
       }
     }
   };
 
+  let with_name = quote! {
+    #[allow(clippy::too_many_arguments)]
+    fn with_name(runtime_template: &::rspack_core::RuntimeTemplate, name: &str, #(#field_names: #field_tys,)*) -> Self {
+      Self {
+        source_map_kind: ::rspack_util::source_map::SourceMapKind::empty(),
+        custom_source: None,
+        cached_generated_code: Default::default(),
+        chunk: None,
+        id: runtime_template.create_custom_runtime_module_identifier(name),
+        #(#field_names,)*
+      }
+    }
+  };
   quote! {
     #[rspack_cacheable::cacheable]
     #input
 
     impl #impl_generics #name #ty_generics #where_clause {
       #with_default
 
+      #with_name
+
       async fn get_generated_code(
         &self,
         compilation: &::rspack_core::Compilation,
@@ -96,8 +123,21 @@ pub fn impl_runtime_module(
       }
     }
 
+    impl #impl_generics ::rspack_core::AttachableRuntimeModule for #name #ty_generics #where_clause {
+      fn attach(&mut self, chunk: ::rspack_core::ChunkUkey) -> () {
+        self.chunk = Some(chunk);
+      }
+    }
+
+    impl #impl_generics ::rspack_core::NamedRuntimeModule for #name #ty_generics #where_clause {
+      fn name(&self) -> ::rspack_collections::Identifier {
+        self.id
+      }
+    }
+
     impl #impl_generics rspack_collections::Identifiable for #name #ty_generics #where_clause {
       fn identifier(&self) -> rspack_collections::Identifier {
+        use rspack_core::NamedRuntimeModule;
         self.name()
       }
     }
@@ -150,6 +190,7 @@ pub fn impl_runtime_module(
       }
 
       fn readable_identifier(&self, _context: &::rspack_core::Context) -> std::borrow::Cow<str> {
+        use rspack_core::NamedRuntimeModule;
         self.name().as_str().into()
       }
 
@@ -194,6 +235,7 @@ pub fn impl_runtime_module(
         runtime: Option<&::rspack_core::RuntimeSpec>,
       ) -> rspack_error::Result<::rspack_hash::RspackHashDigest> {
         use rspack_util::ext::DynHash;
+        use rspack_core::NamedRuntimeModule;
         let mut hasher = rspack_hash::RspackHash::from(&compilation.options.output);
         self.name().dyn_hash(&mut hasher);
         self.stage().dyn_hash(&mut hasher);

crates/rspack_macros/src/runtime_module.rs

@@ -23,7 +23,7 @@ tracing            = { workspace = true }
 trybuild           = { workspace = true }
 
 [package.metadata.cargo-shear]
-ignored = ["tracing", "rspack_cacheable", "async-trait", "rspack_hash", "tokio"]
+ignored = ["tracing", "rspack_cacheable", "rspack_collections", "async-trait", "rspack_hash", "tokio"]
 
 [lints]
 workspace = true

crates/rspack_macros_test/Cargo.toml

@@ -1,6 +1,5 @@
 use std::marker::PhantomData;
 
-use rspack_collections::Identifier;
 use rspack_core::{Compilation, RuntimeModule, rspack_sources::Source};
 use rspack_macros::impl_runtime_module;
 
@@ -15,10 +14,6 @@ fn with_generic() {
 
   #[async_trait::async_trait]
   impl<T: std::fmt::Debug + Send + Sync + Eq + 'static> RuntimeModule for Foo<T> {
-    fn name(&self) -> Identifier {
-      todo!()
-    }
-
     async fn generate(&self, _: &Compilation) -> rspack_error::Result<String> {
       todo!()
     }

crates/rspack_macros_test/tests/runtime_module.rs

@@ -1,8 +1,7 @@
 use std::{borrow::Cow, ptr::NonNull, sync::LazyLock};
 
-use rspack_collections::Identifier;
 use rspack_core::{
-  BooleanMatcher, ChunkGroupOrderKey, ChunkUkey, Compilation, CrossOriginLoading, RuntimeGlobals,
+  BooleanMatcher, ChunkGroupOrderKey, Compilation, CrossOriginLoading, RuntimeGlobals,
   RuntimeModule, RuntimeModuleStage, RuntimeTemplate, compile_boolean_matcher, impl_runtime_module,
 };
 use rspack_plugin_runtime::{
@@ -42,10 +41,7 @@ static CSS_LOADING_WITH_PRELOAD_RUNTIME_REQUIREMENTS: LazyLock<RuntimeGlobals> =
 
 #[impl_runtime_module]
 #[derive(Debug)]
-pub struct CssLoadingRuntimeModule {
-  id: Identifier,
-  chunk: Option<ChunkUkey>,
-}
+pub struct CssLoadingRuntimeModule {}
 
 impl CssLoadingRuntimeModule {
   pub fn get_runtime_requirements_basic() -> RuntimeGlobals {
@@ -67,13 +63,7 @@ impl CssLoadingRuntimeModule {
 
 impl CssLoadingRuntimeModule {
   pub fn new(runtime_template: &RuntimeTemplate) -> Self {
-    Self::with_default(
-      Identifier::from(format!(
-        "{}css_loading",
-        runtime_template.runtime_module_prefix()
-      )),
-      None,
-    )
+    Self::with_default(runtime_template)
   }
 
   fn template_id(&self, id: TemplateId) -> String {
@@ -105,10 +95,6 @@ enum TemplateId {
 
 #[async_trait::async_trait]
 impl RuntimeModule for CssLoadingRuntimeModule {
-  fn name(&self) -> Identifier {
-    self.id
-  }
-
   fn template(&self) -> Vec<(String, String)> {
     vec![
       (
@@ -395,10 +381,6 @@ installedChunks[chunkId] = 0;
     }
   }
 
-  fn attach(&mut self, chunk: ChunkUkey) {
-    self.chunk = Some(chunk);
-  }
-
   fn stage(&self) -> RuntimeModuleStage {
     RuntimeModuleStage::Attach
   }

crates/rspack_plugin_css/src/runtime/mod.rs

@@ -32,7 +32,7 @@ use tokio::sync::RwLock;
 use crate::{
   chunk_link::ChunkLinkContext, dependency::dyn_import::DynamicImportDependencyTemplate,
   ensure_entry_exports::ensure_entry_exports, esm_lib_parser_plugin::EsmLibParserPlugin,
-  preserve_modules::preserve_modules, runtime::RegisterModuleRuntime,
+  preserve_modules::preserve_modules, runtime::EsmRegisterModuleRuntimeModule,
 };
 
 pub static RSPACK_ESM_RUNTIME_CHUNK: &str = "RSPACK_ESM_RUNTIME";
@@ -362,15 +362,20 @@ async fn additional_chunk_runtime_requirements(
 #[plugin_hook(CompilationRuntimeRequirementInTree for EsmLibraryPlugin)]
 async fn runtime_requirements_in_tree(
   &self,
-  _compilation: &Compilation,
+  compilation: &Compilation,
   chunk_ukey: &ChunkUkey,
   _all_runtime_requirements: &RuntimeGlobals,
   runtime_requirements: &RuntimeGlobals,
   _runtime_requirements_mut: &mut RuntimeGlobals,
   runtime_modules_to_add: &mut Vec<(ChunkUkey, Box<dyn RuntimeModule>)>,
 ) -> Result<Option<()>> {
   if runtime_requirements.contains(RuntimeGlobals::REQUIRE) {
-    runtime_modules_to_add.push((*chunk_ukey, Box::new(RegisterModuleRuntime::default())));
+    runtime_modules_to_add.push((
+      *chunk_ukey,
+      Box::new(EsmRegisterModuleRuntimeModule::new(
+        &compilation.runtime_template,
+      )),
+    ));
   }
 
   Ok(None)