feat: support hooks.afterOptimizeModules (web-infra-dev#5102)

crates/node_binding/binding.d.ts

@@ -324,6 +324,7 @@ export interface JsHooks {
   afterEmit: (...args: any[]) => any
   make: (...args: any[]) => any
   optimizeModules: (...args: any[]) => any
+  afterOptimizeModules: (...args: any[]) => any
   optimizeTree: (...args: any[]) => any
   optimizeChunkModules: (...args: any[]) => any
   beforeCompile: (...args: any[]) => any

crates/node_binding/src/hook.rs

@@ -33,6 +33,7 @@ pub enum Hook {
   AfterCompile,
   FinishModules,
   OptimizeModules,
+  AfterOptimizeModules,
   OptimizeTree,
   /// webpack `compilation.hooks.chunkAsset`
   ChunkAsset,
@@ -77,6 +78,7 @@ impl From<String> for Hook {
       "afterCompile" => Hook::AfterCompile,
       "finishModules" => Hook::FinishModules,
       "optimizeModules" => Hook::OptimizeModules,
+      "afterOptimizeModules" => Hook::AfterOptimizeModules,
       "optimizeTree" => Hook::OptimizeTree,
       "chunkAsset" => Hook::ChunkAsset,
       "normalModuleFactoryResolveForScheme" => Hook::NormalModuleFactoryResolveForScheme,

crates/node_binding/src/plugins/mod.rs

@@ -45,6 +45,7 @@ pub struct JsHooksAdapter {
   pub should_emit_tsfn: ThreadsafeFunction<JsCompilation, Option<bool>>,
   pub after_emit_tsfn: ThreadsafeFunction<(), ()>,
   pub optimize_modules_tsfn: ThreadsafeFunction<JsCompilation, ()>,
+  pub after_optimize_modules_tsfn: ThreadsafeFunction<JsCompilation, ()>,
   pub optimize_tree_tsfn: ThreadsafeFunction<(), ()>,
   pub optimize_chunk_modules_tsfn: ThreadsafeFunction<JsCompilation, ()>,
   pub before_compile_tsfn: ThreadsafeFunction<(), ()>,
@@ -531,6 +532,26 @@ impl rspack_core::Plugin for JsHooksAdapter {
       .unwrap_or_else(|err| panic!("Failed to call optimize modules: {err}"))
   }
 
+  async fn after_optimize_modules(
+    &self,
+    compilation: &mut rspack_core::Compilation,
+  ) -> rspack_error::Result<()> {
+    if self.is_hook_disabled(&Hook::AfterOptimizeModules) {
+      return Ok(());
+    }
+    let compilation = JsCompilation::from_compilation(unsafe {
+      std::mem::transmute::<&'_ mut rspack_core::Compilation, &'static mut rspack_core::Compilation>(
+        compilation,
+      )
+    });
+    self
+      .after_optimize_modules_tsfn
+      .call(compilation, ThreadsafeFunctionCallMode::Blocking)
+      .into_rspack_result()?
+      .await
+      .unwrap_or_else(|err| panic!("Failed to call optimize modules: {err}"))
+  }
+
   async fn optimize_tree(
     &self,
     _compilation: &mut rspack_core::Compilation,
@@ -815,6 +836,7 @@ impl JsHooksAdapter {
       asset_emitted,
       after_emit,
       optimize_modules,
+      after_optimize_modules,
       optimize_tree,
       optimize_chunk_modules,
       before_resolve,
@@ -877,6 +899,8 @@ impl JsHooksAdapter {
     let make_tsfn: ThreadsafeFunction<(), ()> = js_fn_into_threadsafe_fn!(make, env);
     let optimize_modules_tsfn: ThreadsafeFunction<JsCompilation, ()> =
       js_fn_into_threadsafe_fn!(optimize_modules, env);
+    let after_optimize_modules_tsfn: ThreadsafeFunction<JsCompilation, ()> =
+      js_fn_into_threadsafe_fn!(after_optimize_modules, env);
     let optimize_tree_tsfn: ThreadsafeFunction<(), ()> =
       js_fn_into_threadsafe_fn!(optimize_tree, env);
     let optimize_chunk_modules_tsfn: ThreadsafeFunction<JsCompilation, ()> =
@@ -936,6 +960,7 @@ impl JsHooksAdapter {
       asset_emitted_tsfn,
       after_emit_tsfn,
       optimize_modules_tsfn,
+      after_optimize_modules_tsfn,
       optimize_tree_tsfn,
       optimize_chunk_modules_tsfn,
       before_compile_tsfn,

crates/rspack_binding_values/src/hooks.rs

@@ -27,6 +27,7 @@ pub struct JsHooks {
   pub after_emit: JsFunction,
   pub make: JsFunction,
   pub optimize_modules: JsFunction,
+  pub after_optimize_modules: JsFunction,
   pub optimize_tree: JsFunction,
   pub optimize_chunk_modules: JsFunction,
   pub before_compile: JsFunction,

crates/rspack_core/src/compiler/compilation.rs

@@ -1372,6 +1372,7 @@ impl Compilation {
     use_code_splitting_cache(self, |compilation| async {
       build_chunk_graph(compilation)?;
       plugin_driver.optimize_modules(compilation).await?;
+      plugin_driver.after_optimize_modules(compilation).await?;
       plugin_driver.optimize_chunks(compilation).await?;
       Ok(compilation)
     })

crates/rspack_core/src/plugin/api.rs

@@ -404,6 +404,10 @@ pub trait Plugin: Debug + Send + Sync {
     Ok(())
   }
 
+  async fn after_optimize_modules(&self, _compilation: &mut Compilation) -> Result<()> {
+    Ok(())
+  }
+
   async fn optimize_dependencies(&self, _compilation: &mut Compilation) -> Result<Option<()>> {
     Ok(None)
   }

crates/rspack_core/src/plugin/plugin_driver.rs

@@ -544,6 +544,15 @@ impl PluginDriver {
     Ok(())
   }
 
+  #[instrument(name = "plugin:after_optimize_modules", skip_all)]
+  pub async fn after_optimize_modules(&self, compilation: &mut Compilation) -> Result<()> {
+    for plugin in &self.plugins {
+      // `SyncHook`
+      plugin.after_optimize_modules(compilation).await?;
+    }
+    Ok(())
+  }
+
   #[instrument(name = "plugin:optimize_dependencies", skip_all)]
   pub async fn optimize_dependencies(&self, compilation: &mut Compilation) -> Result<Option<()>> {
     for plugin in &self.plugins {

packages/rspack/src/Compilation.ts

@@ -106,6 +106,7 @@ export class Compilation {
 		log: tapable.SyncBailHook<[string, LogEntry], true>;
 		additionalAssets: any;
 		optimizeModules: tapable.SyncBailHook<Iterable<JsModule>, undefined>;
+		afterOptimizeModules: tapable.SyncHook<Iterable<JsModule>, undefined>;
 		optimizeTree: tapable.AsyncSeriesBailHook<
 			[Iterable<Chunk>, Iterable<JsModule>],
 			undefined
@@ -169,6 +170,7 @@ export class Compilation {
 			]),
 			log: new tapable.SyncBailHook(["origin", "logEntry"]),
 			optimizeModules: new tapable.SyncBailHook(["modules"]),
+			afterOptimizeModules: new tapable.SyncBailHook(["modules"]),
 			optimizeTree: new tapable.AsyncSeriesBailHook(["chunks", "modules"]),
 			optimizeChunkModules: new tapable.AsyncSeriesBailHook([
 				"chunks",

packages/rspack/src/Compiler.ts

@@ -331,6 +331,7 @@ class Compiler {
 				// No matter how it will be implemented, it will be copied to the child compiler.
 				compilation: this.#compilation.bind(this),
 				optimizeModules: this.#optimizeModules.bind(this),
+				afterOptimizeModules: this.#afterOptimizeModules.bind(this),
 				optimizeTree: this.#optimizeTree.bind(this),
 				optimizeChunkModules: this.#optimizeChunkModules.bind(this),
 				finishModules: this.#finishModules.bind(this),
@@ -647,6 +648,7 @@ class Compiler {
 			optimizeTree: this.compilation.hooks.optimizeTree,
 			finishModules: this.compilation.hooks.finishModules,
 			optimizeModules: this.compilation.hooks.optimizeModules,
+			afterOptimizeModules: this.compilation.hooks.afterOptimizeModules,
 			chunkAsset: this.compilation.hooks.chunkAsset,
 			beforeResolve: this.compilation.normalModuleFactory?.hooks.beforeResolve,
 			afterResolve: this.compilation.normalModuleFactory?.hooks.afterResolve,
@@ -791,6 +793,13 @@ class Compiler {
 		this.#updateDisabledHooks();
 	}
 
+	async #afterOptimizeModules() {
+		await this.compilation.hooks.afterOptimizeModules.promise(
+			this.compilation.modules
+		);
+		this.#updateDisabledHooks();
+	}
+
 	#chunkAsset(assetArg: binding.JsChunkAssetArgs) {
 		this.compilation.hooks.chunkAsset.call(assetArg.chunk, assetArg.filename);
 		this.#updateDisabledHooks();

packages/rspack/tests/Compiler.test.ts

@@ -1332,6 +1332,34 @@ describe("Compiler", () => {
 			});
 		});
 
+		it("should call afterOptimizeModules hook correctly", done => {
+			class MyPlugin {
+				apply(compiler: Compiler) {
+					let a = 1;
+					compiler.hooks.compilation.tap("MyPlugin", compilation => {
+						compilation.hooks.optimizeModules.tap("MyPlugin", () => {
+							a += 1;
+						});
+
+						compilation.hooks.afterOptimizeModules.tap("MyPlugin", modules => {
+							expect(a).toBeGreaterThan(1);
+							expect(modules.length).toEqual(1);
+							expect(modules[0].resource.includes("d.js")).toBeTruthy();
+						});
+					});
+				}
+			}
+			const compiler = rspack({
+				entry: "./d",
+				context: path.join(__dirname, "fixtures"),
+				plugins: [new MyPlugin()]
+			});
+
+			compiler.build(err => {
+				done(err);
+			});
+		});
+
 		it("should call getCache function correctly", done => {
 			class MyPlugin {
 				apply(compiler: Compiler) {