refactor: optimize modules and after optimize modules hook

crates/node_binding/binding.d.ts

@@ -316,8 +316,6 @@ export interface JsExecuteModuleResult {
 }
 
 export interface JsHooks {
-  optimizeModules: (compilation: JsCompilation) => void
-  afterOptimizeModules: (compilation: JsCompilation) => void
   optimizeTree: () => void
   optimizeChunkModules: (compilation: JsCompilation) => void
   afterResolve: (data: AfterResolveData) => Promise<(boolean | void | AfterResolveCreateData)[]>
@@ -1296,6 +1294,8 @@ export interface RegisterJsTaps {
   registerCompilationExecuteModuleTaps: (stages: Array<number>) => Array<{ function: ((arg: JsExecuteModuleArg) => void); stage: number; }>
   registerCompilationRuntimeModuleTaps: (stages: Array<number>) => Array<{ function: ((arg: JsRuntimeModuleArg) => JsRuntimeModule | undefined); stage: number; }>
   registerCompilationFinishModulesTaps: (stages: Array<number>) => Array<{ function: ((arg: JsCompilation) => Promise<void>); stage: number; }>
+  registerCompilationOptimizeModulesTaps: (stages: Array<number>) => Array<{ function: (() => boolean | undefined); stage: number; }>
+  registerCompilationAfterOptimizeModulesTaps: (stages: Array<number>) => Array<{ function: (() => void); stage: number; }>
   registerCompilationChunkAssetTaps: (stages: Array<number>) => Array<{ function: ((arg: JsChunkAssetArgs) => void); stage: number; }>
   registerCompilationProcessAssetsTaps: (stages: Array<number>) => Array<{ function: ((arg: JsCompilation) => Promise<void>); stage: number; }>
   registerCompilationAfterProcessAssetsTaps: (stages: Array<number>) => Array<{ function: ((arg: JsCompilation) => void); stage: number; }>

crates/node_binding/src/hook.rs

@@ -4,8 +4,6 @@ use std::sync::RwLock;
 #[derive(PartialEq)]
 pub enum Hook {
   OptimizeChunkModules,
-  OptimizeModules,
-  AfterOptimizeModules,
   OptimizeTree,
   ContextModuleFactoryBeforeResolve,
   ContextModuleFactoryAfterResolve,
@@ -18,8 +16,6 @@ impl From<String> for Hook {
   fn from(s: String) -> Self {
     match s.as_str() {
       "optimizeChunkModules" => Hook::OptimizeChunkModules,
-      "optimizeModules" => Hook::OptimizeModules,
-      "afterOptimizeModules" => Hook::AfterOptimizeModules,
       "optimizeTree" => Hook::OptimizeTree,
       "contextModuleFactoryBeforeResolve" => Hook::ContextModuleFactoryBeforeResolve,
       "contextModuleFactoryAfterResolve" => Hook::ContextModuleFactoryAfterResolve,

crates/node_binding/src/plugins/interceptor.rs

@@ -12,13 +12,14 @@ use rspack_binding_values::{
 };
 use rspack_core::{
   rspack_sources::SourceExt, AssetEmittedInfo, BeforeResolveArgs, BoxModule, Chunk, ChunkUkey,
-  CodeGenerationResults, Compilation, CompilationAfterProcessAssetsHook,
-  CompilationBuildModuleHook, CompilationChunkAssetHook, CompilationExecuteModuleHook,
-  CompilationFinishModulesHook, CompilationParams, CompilationProcessAssetsHook,
-  CompilationRuntimeModuleHook, CompilationStillValidModuleHook, CompilationSucceedModuleHook,
-  CompilerAfterEmitHook, CompilerAssetEmittedHook, CompilerCompilationHook, CompilerEmitHook,
-  CompilerFinishMakeHook, CompilerMakeHook, CompilerShouldEmitHook, CompilerThisCompilationHook,
-  ExecuteModuleId, MakeParam, ModuleIdentifier, NormalModuleFactoryBeforeResolveHook,
+  CodeGenerationResults, Compilation, CompilationAfterOptimizeModulesHook,
+  CompilationAfterProcessAssetsHook, CompilationBuildModuleHook, CompilationChunkAssetHook,
+  CompilationExecuteModuleHook, CompilationFinishModulesHook, CompilationOptimizeModulesHook,
+  CompilationParams, CompilationProcessAssetsHook, CompilationRuntimeModuleHook,
+  CompilationStillValidModuleHook, CompilationSucceedModuleHook, CompilerAfterEmitHook,
+  CompilerAssetEmittedHook, CompilerCompilationHook, CompilerEmitHook, CompilerFinishMakeHook,
+  CompilerMakeHook, CompilerShouldEmitHook, CompilerThisCompilationHook, ExecuteModuleId,
+  MakeParam, ModuleIdentifier, NormalModuleFactoryBeforeResolveHook,
 };
 use rspack_hook::{
   AsyncParallel3, AsyncSeries, AsyncSeries2, AsyncSeries3, AsyncSeriesBail, Hook, Interceptor,
@@ -304,6 +305,12 @@ pub struct RegisterJsTaps {
     ts_type = "(stages: Array<number>) => Array<{ function: ((arg: JsCompilation) => Promise<void>); stage: number; }>"
   )]
   pub register_compilation_finish_modules_taps: RegisterFunction<JsCompilation, Promise<()>>,
+  #[napi(
+    ts_type = "(stages: Array<number>) => Array<{ function: (() => boolean | undefined); stage: number; }>"
+  )]
+  pub register_compilation_optimize_modules_taps: RegisterFunction<(), Option<bool>>,
+  #[napi(ts_type = "(stages: Array<number>) => Array<{ function: (() => void); stage: number; }>")]
+  pub register_compilation_after_optimize_modules_taps: RegisterFunction<(), ()>,
   #[napi(
     ts_type = "(stages: Array<number>) => Array<{ function: ((arg: JsChunkAssetArgs) => void); stage: number; }>"
   )]
@@ -404,6 +411,18 @@ define_register!(
   cache = false,
   sync = false,
 );
+define_register!(
+  RegisterCompilationOptimizeModulesTaps,
+  tap = CompilationOptimizeModulesTap<(), Option<bool>> @ CompilationOptimizeModulesHook,
+  cache = true,
+  sync = false,
+);
+define_register!(
+  RegisterCompilationAfterOptimizeModulesTaps,
+  tap = CompilationAfterOptimizeModulesTap<(), ()> @ CompilationAfterOptimizeModulesHook,
+  cache = false,
+  sync = false,
+);
 define_register!(
   RegisterCompilationRuntimeModuleTaps,
   tap = CompilationRuntimeModuleTap<JsRuntimeModuleArg, Option<JsRuntimeModule>> @ CompilationRuntimeModuleHook,
@@ -660,6 +679,28 @@ impl AsyncSeries<Compilation> for CompilationFinishModulesTap {
   }
 }
 
+#[async_trait]
+impl AsyncSeriesBail<Compilation, bool> for CompilationOptimizeModulesTap {
+  async fn run(&self, _compilation: &mut Compilation) -> rspack_error::Result<Option<bool>> {
+    self.function.call(()).await
+  }
+
+  fn stage(&self) -> i32 {
+    self.stage
+  }
+}
+
+#[async_trait]
+impl AsyncSeries<Compilation> for CompilationAfterOptimizeModulesTap {
+  async fn run(&self, _compilation: &mut Compilation) -> rspack_error::Result<()> {
+    self.function.call(()).await
+  }
+
+  fn stage(&self) -> i32 {
+    self.stage
+  }
+}
+
 #[async_trait]
 impl AsyncSeries3<Compilation, ModuleIdentifier, ChunkUkey> for CompilationRuntimeModuleTap {
   async fn run(

crates/node_binding/src/plugins/mod.rs

@@ -19,11 +19,14 @@ use rspack_core::{
 };
 use rspack_hook::Hook as _;
 
-use self::interceptor::RegisterCompilationChunkAssetTaps;
 use self::interceptor::RegisterCompilationFinishModulesTaps;
 use self::interceptor::RegisterCompilationStillValidModuleTaps;
 use self::interceptor::RegisterCompilationSucceedModuleTaps;
 use self::interceptor::RegisterCompilerFinishMakeTaps;
+use self::interceptor::{
+  RegisterCompilationAfterOptimizeModulesTaps, RegisterCompilationChunkAssetTaps,
+  RegisterCompilationOptimizeModulesTaps,
+};
 use self::interceptor::{
   RegisterCompilationAfterProcessAssetsTaps, RegisterCompilerAssetEmittedTaps,
 };
@@ -59,6 +62,8 @@ pub struct JsHooksAdapterPlugin {
   register_compilation_succeed_module_taps: RegisterCompilationSucceedModuleTaps,
   register_compilation_execute_module_taps: RegisterCompilationExecuteModuleTaps,
   register_compilation_finish_modules_taps: RegisterCompilationFinishModulesTaps,
+  register_compilation_optimize_modules_taps: RegisterCompilationOptimizeModulesTaps,
+  register_compilation_after_optimize_modules_taps: RegisterCompilationAfterOptimizeModulesTaps,
   register_compilation_runtime_module_taps: RegisterCompilationRuntimeModuleTaps,
   register_compilation_chunk_asset_taps: RegisterCompilationChunkAssetTaps,
   register_compilation_process_assets_taps: RegisterCompilationProcessAssetsTaps,
@@ -158,6 +163,20 @@ impl rspack_core::Plugin for JsHooksAdapterPlugin {
       .compilation_hooks
       .finish_modules
       .intercept(self.register_compilation_finish_modules_taps.clone());
+    ctx
+      .context
+      .compilation_hooks
+      .optimize_modules
+      .intercept(self.register_compilation_optimize_modules_taps.clone());
+    ctx
+      .context
+      .compilation_hooks
+      .after_optimize_modules
+      .intercept(
+        self
+          .register_compilation_after_optimize_modules_taps
+          .clone(),
+      );
     ctx
       .context
       .compilation_hooks
@@ -304,36 +323,6 @@ impl rspack_core::Plugin for JsHooksAdapterPlugin {
     })
   }
 
-  async fn optimize_modules(
-    &self,
-    compilation: &mut rspack_core::Compilation,
-  ) -> rspack_error::Result<()> {
-    if self.is_hook_disabled(&Hook::OptimizeModules) {
-      return Ok(());
-    }
-    // SAFETY:
-    // 1. `Compiler` is stored on the heap and pinned in binding crate.
-    // 2. `Compilation` outlives `JsCompilation` and `Compiler` outlives `Compilation`.
-    // 3. `JsCompilation` was replaced everytime a new `Compilation` was created before getting accessed.
-    let compilation = unsafe { JsCompilation::from_compilation(compilation) };
-    self.hooks.optimize_modules.call(compilation).await
-  }
-
-  async fn after_optimize_modules(
-    &self,
-    compilation: &mut rspack_core::Compilation,
-  ) -> rspack_error::Result<()> {
-    if self.is_hook_disabled(&Hook::AfterOptimizeModules) {
-      return Ok(());
-    }
-    // SAFETY:
-    // 1. `Compiler` is stored on the heap and pinned in binding crate.
-    // 2. `Compilation` outlives `JsCompilation` and `Compiler` outlives `Compilation`.
-    // 3. `JsCompilation` was replaced everytime a new `Compilation` was created before getting accessed.
-    let compilation = unsafe { JsCompilation::from_compilation(compilation) };
-    self.hooks.after_optimize_modules.call(compilation).await
-  }
-
   async fn optimize_tree(
     &self,
     _compilation: &mut rspack_core::Compilation,
@@ -409,6 +398,13 @@ impl JsHooksAdapterPlugin {
       register_compilation_finish_modules_taps: RegisterCompilationFinishModulesTaps::new(
         register_js_taps.register_compilation_finish_modules_taps,
       ),
+      register_compilation_optimize_modules_taps: RegisterCompilationOptimizeModulesTaps::new(
+        register_js_taps.register_compilation_optimize_modules_taps,
+      ),
+      register_compilation_after_optimize_modules_taps:
+        RegisterCompilationAfterOptimizeModulesTaps::new(
+          register_js_taps.register_compilation_after_optimize_modules_taps,
+        ),
       register_compilation_runtime_module_taps: RegisterCompilationRuntimeModuleTaps::new(
         register_js_taps.register_compilation_runtime_module_taps,
       ),

crates/rspack_binding_values/src/hooks.rs

@@ -8,10 +8,6 @@ use crate::{
 
 #[napi(object, object_to_js = false)]
 pub struct JsHooks {
-  #[napi(ts_type = "(compilation: JsCompilation) => void")]
-  pub optimize_modules: ThreadsafeFunction<JsCompilation, ()>,
-  #[napi(ts_type = "(compilation: JsCompilation) => void")]
-  pub after_optimize_modules: ThreadsafeFunction<JsCompilation, ()>,
   #[napi(ts_type = "() => void")]
   pub optimize_tree: ThreadsafeFunction<(), ()>,
   #[napi(ts_type = "(compilation: JsCompilation) => void")]

crates/rspack_core/src/compiler/compilation.rs

@@ -13,7 +13,9 @@ use rayon::prelude::*;
 use rspack_error::{error, Diagnostic, Result, Severity, TWithDiagnosticArray};
 use rspack_futures::FuturesResults;
 use rspack_hash::{RspackHash, RspackHashDigest};
-use rspack_hook::{AsyncSeries2Hook, AsyncSeries3Hook, AsyncSeriesHook, SyncSeries4Hook};
+use rspack_hook::{
+  AsyncSeries2Hook, AsyncSeries3Hook, AsyncSeriesBailHook, AsyncSeriesHook, SyncSeries4Hook,
+};
 use rspack_identifier::{Identifiable, Identifier, IdentifierMap, IdentifierSet};
 use rspack_sources::{BoxSource, CachedSource, SourceExt};
 use rustc_hash::{FxHashMap as HashMap, FxHashSet as HashSet, FxHasher};
@@ -52,6 +54,8 @@ pub type CompilationSucceedModuleHook = AsyncSeriesHook<BoxModule>;
 pub type CompilationExecuteModuleHook =
   SyncSeries4Hook<ModuleIdentifier, IdentifierSet, CodeGenerationResults, ExecuteModuleId>;
 pub type CompilationFinishModulesHook = AsyncSeriesHook<Compilation>;
+pub type CompilationOptimizeModulesHook = AsyncSeriesBailHook<Compilation, bool>;
+pub type CompilationAfterOptimizeModulesHook = AsyncSeriesHook<Compilation>;
 pub type CompilationRuntimeModuleHook = AsyncSeries3Hook<Compilation, ModuleIdentifier, ChunkUkey>;
 pub type CompilationChunkAssetHook = AsyncSeries2Hook<Chunk, String>;
 pub type CompilationProcessAssetsHook = AsyncSeriesHook<Compilation>;
@@ -64,6 +68,8 @@ pub struct CompilationHooks {
   pub succeed_module: CompilationSucceedModuleHook,
   pub execute_module: CompilationExecuteModuleHook,
   pub finish_modules: CompilationFinishModulesHook,
+  pub optimize_modules: CompilationOptimizeModulesHook,
+  pub after_optimize_modules: CompilationAfterOptimizeModulesHook,
   pub runtime_module: CompilationRuntimeModuleHook,
   pub chunk_asset: CompilationChunkAssetHook,
   pub process_assets: CompilationProcessAssetsHook,
@@ -844,8 +850,19 @@ impl Compilation {
     let start = logger.time("create chunks");
     use_code_splitting_cache(self, |compilation| async {
       build_chunk_graph(compilation)?;
-      plugin_driver.optimize_modules(compilation).await?;
-      plugin_driver.after_optimize_modules(compilation).await?;
+      while matches!(
+        plugin_driver
+          .compilation_hooks
+          .optimize_modules
+          .call(compilation)
+          .await?,
+        Some(true)
+      ) {}
+      plugin_driver
+        .compilation_hooks
+        .after_optimize_modules
+        .call(compilation)
+        .await?;
       plugin_driver.optimize_chunks(compilation).await?;
       Ok(compilation)
     })

crates/rspack_core/src/plugin/api.rs

@@ -252,14 +252,6 @@ pub trait Plugin: Debug + Send + Sync {
     Ok(())
   }
 
-  async fn optimize_modules(&self, _compilation: &mut Compilation) -> Result<()> {
-    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

@@ -408,23 +408,6 @@ impl PluginDriver {
     Ok(())
   }
 
-  #[instrument(name = "plugin:optimize_modules", skip_all)]
-  pub async fn optimize_modules(&self, compilation: &mut Compilation) -> Result<()> {
-    for plugin in &self.plugins {
-      plugin.optimize_modules(compilation).await?;
-    }
-    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 {

crates/rspack_plugin_progress/src/lib.rs

@@ -13,7 +13,7 @@ use rspack_core::{
   PluginOptimizeChunksOutput,
 };
 use rspack_error::Result;
-use rspack_hook::{plugin, plugin_hook, AsyncSeries, AsyncSeries2};
+use rspack_hook::{plugin, plugin_hook, AsyncSeries, AsyncSeries2, AsyncSeriesBail};
 
 #[derive(Debug, Clone, Default)]
 pub struct ProgressPluginOptions {
@@ -308,6 +308,18 @@ async fn finish_modules(&self, _compilation: &mut Compilation) -> Result<()> {
   Ok(())
 }
 
+#[plugin_hook(AsyncSeriesBail<Compilation, bool> for ProgressPlugin)]
+async fn optimize_modules(&self, _compilation: &mut Compilation) -> Result<Option<bool>> {
+  self.sealing_hooks_report("module optimization", 7);
+  Ok(None)
+}
+
+#[plugin_hook(AsyncSeries<Compilation> for ProgressPlugin)]
+async fn after_optimize_modules(&self, _compilation: &mut Compilation) -> Result<()> {
+  self.sealing_hooks_report("after module optimization", 8);
+  Ok(())
+}
+
 #[plugin_hook(AsyncSeries<Compilation> for ProgressPlugin, stage = Compilation::PROCESS_ASSETS_STAGE_ADDITIONAL)]
 async fn process_assets(&self, _compilation: &mut Compilation) -> Result<()> {
   self.sealing_hooks_report("asset processing", 35);
@@ -379,6 +391,16 @@ impl Plugin for ProgressPlugin {
       .compilation_hooks
       .finish_modules
       .tap(finish_modules::new(self));
+    ctx
+      .context
+      .compilation_hooks
+      .optimize_modules
+      .tap(optimize_modules::new(self));
+    ctx
+      .context
+      .compilation_hooks
+      .after_optimize_modules
+      .tap(after_optimize_modules::new(self));
     ctx
       .context
       .compilation_hooks
@@ -408,16 +430,6 @@ impl Plugin for ProgressPlugin {
     Ok(None)
   }
 
-  async fn optimize_modules(&self, _compilation: &mut Compilation) -> Result<()> {
-    self.sealing_hooks_report("module optimization", 7);
-    Ok(())
-  }
-
-  async fn after_optimize_modules(&self, _compilation: &mut Compilation) -> Result<()> {
-    self.sealing_hooks_report("after module optimization", 8);
-    Ok(())
-  }
-
   async fn optimize_chunks(
     &self,
     _ctx: PluginContext,