feat: add buildModule hook (#3466)

* feat: add buildModule hook

* feat: use object assign

---------

Co-authored-by: Hana <[email protected]>

crates/node_binding/binding.d.ts

@@ -199,6 +199,7 @@ export interface JsHooks {
   afterCompile: (...args: any[]) => any
   finishModules: (...args: any[]) => any
   finishMake: (...args: any[]) => any
+  buildModule: (...args: any[]) => any
   beforeResolve: (...args: any[]) => any
   afterResolve: (...args: any[]) => any
   contextModuleBeforeResolve: (...args: any[]) => any

crates/node_binding/src/hook.rs

@@ -5,6 +5,7 @@ use std::sync::{Arc, RwLock};
 pub enum Hook {
   Make,
   FinishMake,
+  BuildModule,
   Compilation,
   ThisCompilation,
   ProcessAssetsStageAdditional,
@@ -37,6 +38,7 @@ impl From<String> for Hook {
     match s.as_str() {
       "make" => Hook::Make,
       "finishMake" => Hook::FinishMake,
+      "buildModule" => Hook::BuildModule,
       "compilation" => Hook::Compilation,
       "thisCompilation" => Hook::ThisCompilation,
       "processAssetsStageAdditional" => Hook::ProcessAssetsStageAdditional,

crates/node_binding/src/js_values/hooks.rs

@@ -22,6 +22,7 @@ pub struct JsHooks {
   pub after_compile: JsFunction,
   pub finish_modules: JsFunction,
   pub finish_make: JsFunction,
+  pub build_module: JsFunction,
   pub before_resolve: JsFunction,
   pub after_resolve: JsFunction,
   pub context_module_before_resolve: JsFunction,

crates/node_binding/src/plugins/mod.rs

@@ -43,6 +43,7 @@ pub struct JsHooksAdapter {
   pub after_compile_tsfn: ThreadsafeFunction<JsCompilation, ()>,
   pub finish_modules_tsfn: ThreadsafeFunction<JsCompilation, ()>,
   pub finish_make_tsfn: ThreadsafeFunction<JsCompilation, ()>,
+  pub build_module_tsfn: ThreadsafeFunction<JsModule, ()>, // TODO
   pub chunk_asset_tsfn: ThreadsafeFunction<JsChunkAssetArgs, ()>,
   pub before_resolve: ThreadsafeFunction<BeforeResolveData, (Option<bool>, BeforeResolveData)>,
   pub after_resolve: ThreadsafeFunction<AfterResolveData, Option<bool>>,
@@ -464,6 +465,22 @@ impl rspack_core::Plugin for JsHooksAdapter {
       .map_err(|err| internal_error!("Failed to call finish make: {err}"))?
   }
 
+  async fn build_module(&self, module: &mut dyn rspack_core::Module) -> rspack_error::Result<()> {
+    if self.is_hook_disabled(&Hook::BuildModule) {
+      return Ok(());
+    }
+
+    self
+      .build_module_tsfn
+      .call(
+        module.to_js_module().expect("Convert to js_module failed."),
+        ThreadsafeFunctionCallMode::NonBlocking,
+      )
+      .into_rspack_result()?
+      .await
+      .map_err(|err| internal_error!("Failed to call build module: {err}"))?
+  }
+
   async fn finish_modules(
     &self,
     compilation: &mut rspack_core::Compilation,
@@ -586,6 +603,7 @@ impl JsHooksAdapter {
       after_compile,
       finish_modules,
       finish_make,
+      build_module,
       chunk_asset,
       succeed_module,
       still_valid_module,
@@ -626,6 +644,8 @@ impl JsHooksAdapter {
       js_fn_into_theadsafe_fn!(after_compile, env);
     let finish_make_tsfn: ThreadsafeFunction<JsCompilation, ()> =
       js_fn_into_theadsafe_fn!(finish_make, env);
+    let build_module_tsfn: ThreadsafeFunction<JsModule, ()> =
+      js_fn_into_theadsafe_fn!(build_module, env);
     let finish_modules_tsfn: ThreadsafeFunction<JsCompilation, ()> =
       js_fn_into_theadsafe_fn!(finish_modules, env);
     let context_module_before_resolve: ThreadsafeFunction<BeforeResolveData, Option<bool>> =
@@ -670,6 +690,7 @@ impl JsHooksAdapter {
       normal_module_factory_resolve_for_scheme,
       finish_modules_tsfn,
       finish_make_tsfn,
+      build_module_tsfn,
       chunk_asset_tsfn,
       after_resolve,
       succeed_module_tsfn,

crates/rspack_core/src/plugin/plugin_driver.rs

@@ -526,6 +526,7 @@ impl PluginDriver {
     Ok(())
   }
 
+  #[instrument(name = "plugin:build_module", skip_all)]
   pub async fn build_module(&self, module: &mut dyn Module) -> Result<()> {
     for plugin in &self.plugins {
       plugin.build_module(module).await?;

packages/rspack/src/compilation.ts

@@ -43,6 +43,7 @@ import {
 	createFakeCompilationDependencies,
 	createFakeProcessAssetsHook
 } from "./util/fake";
+import { NormalizedJsModule, normalizeJsModule } from "./util/normalization";
 import MergeCaller from "./util/MergeCaller";
 
 export type AssetInfo = Partial<JsAssetInfo> & Record<string, any>;
@@ -90,6 +91,7 @@ export class Compilation {
 		processWarnings: tapable.SyncWaterfallHook<[Error[]]>;
 		succeedModule: tapable.SyncHook<[JsModule], undefined>;
 		stillValidModule: tapable.SyncHook<[JsModule], undefined>;
+		buildModule: tapable.SyncHook<[NormalizedJsModule]>;
 	};
 	options: RspackOptionsNormalized;
 	outputOptions: OutputNormalized;
@@ -125,7 +127,8 @@ export class Compilation {
 			chunkAsset: new tapable.SyncHook(["chunk", "filename"]),
 			processWarnings: new tapable.SyncWaterfallHook(["warnings"]),
 			succeedModule: new tapable.SyncHook(["module"]),
-			stillValidModule: new tapable.SyncHook(["module"])
+			stillValidModule: new tapable.SyncHook(["module"]),
+			buildModule: new tapable.SyncHook(["module"])
 		};
 		this.compiler = compiler;
 		this.resolverFactory = compiler.resolverFactory;
@@ -608,12 +611,7 @@ export class Compilation {
 	);
 
 	get modules() {
-		return this.getModules().map(item => {
-			return {
-				identifier: () => item.moduleIdentifier,
-				...item
-			};
-		});
+		return this.getModules().map(item => normalizeJsModule(item));
 	}
 
 	get chunks() {

packages/rspack/src/compiler.ts

@@ -37,6 +37,7 @@ import { WatchFileSystem } from "./util/fs";
 import { getScheme } from "./util/scheme";
 import Watching from "./watching";
 import { NormalModule } from "./normalModule";
+import { normalizeJsModule } from "./util/normalization";
 
 class EntryPlugin {
 	constructor(
@@ -343,7 +344,8 @@ class Compiler {
 					contextModuleBeforeResolve:
 						this.#contextModuleBeforeResolve.bind(this),
 					succeedModule: this.#succeedModule.bind(this),
-					stillValidModule: this.#stillValidModule.bind(this)
+					stillValidModule: this.#stillValidModule.bind(this),
+					buildModule: this.#buildModule.bind(this)
 				},
 				createThreadsafeNodeFSFromRaw(this.outputFileSystem),
 				loaderContext => runLoader(loaderContext, this)
@@ -606,7 +608,8 @@ class Compiler {
 			beforeResolve: this.compilation.normalModuleFactory?.hooks.beforeResolve,
 			afterResolve: this.compilation.normalModuleFactory?.hooks.afterResolve,
 			succeedModule: this.compilation.hooks.succeedModule,
-			stillValidModule: this.compilation.hooks.stillValidModule
+			stillValidModule: this.compilation.hooks.stillValidModule,
+			buildModule: this.compilation.hooks.buildModule
 		};
 		for (const [name, hook] of Object.entries(hookMap)) {
 			if (hook?.taps.length === 0) {
@@ -637,6 +640,12 @@ class Compiler {
 		this.#updateDisabledHooks();
 	}
 
+	async #buildModule(module: binding.JsModule) {
+		const normalized = normalizeJsModule(module);
+		this.compilation.hooks.buildModule.call(normalized);
+		this.#updateDisabledHooks();
+	}
+
 	async #processAssets(stage: number) {
 		await this.compilation
 			.__internal_getProcessAssetsHookByStage(stage)

packages/rspack/src/util/normalization.ts

@@ -0,0 +1,11 @@
+import { JsModule } from "@rspack/binding";
+
+export interface NormalizedJsModule extends JsModule {
+	identifier: () => string;
+}
+
+export function normalizeJsModule(m: JsModule): NormalizedJsModule {
+	return Object.assign(m, {
+		identifier: () => m.moduleIdentifier
+	});
+}

packages/rspack/tests/configCases/hooks/build-module/a.js

@@ -0,0 +1 @@
+export const a = 3;

packages/rspack/tests/configCases/hooks/build-module/index.js

@@ -0,0 +1,5 @@
+import { a } from "./a.js";
+
+it("should compile successfully with build-module", () => {
+	expect(a).toBe(3);
+});

packages/rspack/tests/configCases/hooks/build-module/webpack.config.js

@@ -0,0 +1,23 @@
+const { strict } = require("assert");
+const pluginName = "plugin";
+
+class Plugin {
+	apply(compiler) {
+		let identifiers = [];
+		compiler.hooks.compilation.tap(pluginName, compilation => {
+			compilation.hooks.buildModule.tap(pluginName, m => {
+				identifiers.push(m.identifier());
+			});
+		});
+		compiler.hooks.done.tap(pluginName, () => {
+			strict(identifiers.some(i => i.endsWith("index.js")));
+			strict(identifiers.some(i => i.endsWith("a.js")));
+		});
+	}
+}
+
+/**@type {import('@rspack/cli').Configuration}*/
+module.exports = {
+	context: __dirname,
+	plugins: [new Plugin()]
+};