feat: add beforeModuleIds hook (#12809)

feat: add beforeModuleIds hook

Co-authored-by: hardfist <yangjianzju@gmail.com>

Author: hamlim

crates/node_binding/napi-binding.d.ts

@@ -717,6 +717,14 @@ export interface JsBeforeEmitData {
   uid?: number
 }
 
+export interface JsBeforeModuleIdsArg {
+  modules: Array<JsModuleForIds>
+}
+
+export interface JsBeforeModuleIdsResult {
+  assignments: Record<string, string>
+}
+
 export interface JsBuildMeta {
   strictEsmModule?: boolean
   hasTopLevelAwait?: boolean
@@ -980,6 +988,10 @@ export interface JsModuleDescriptor {
   id?: string | number | null
 }
 
+export interface JsModuleForIds {
+  identifier: string
+}
+
 export interface JsNormalModuleFactoryCreateModuleArgs {
   dependencyType: string
   rawRequest: string
@@ -2966,39 +2978,40 @@ export declare enum RegisterJsTapKind {
   CompilationAfterOptimizeModules = 14,
   CompilationOptimizeTree = 15,
   CompilationOptimizeChunkModules = 16,
-  CompilationAdditionalTreeRuntimeRequirements = 17,
-  CompilationRuntimeRequirementInTree = 18,
-  CompilationRuntimeModule = 19,
-  CompilationChunkHash = 20,
-  CompilationChunkAsset = 21,
-  CompilationProcessAssets = 22,
-  CompilationAfterProcessAssets = 23,
-  CompilationSeal = 24,
-  CompilationAfterSeal = 25,
-  NormalModuleFactoryBeforeResolve = 26,
-  NormalModuleFactoryFactorize = 27,
-  NormalModuleFactoryResolve = 28,
-  NormalModuleFactoryAfterResolve = 29,
-  NormalModuleFactoryCreateModule = 30,
-  NormalModuleFactoryResolveForScheme = 31,
-  ContextModuleFactoryBeforeResolve = 32,
-  ContextModuleFactoryAfterResolve = 33,
-  JavascriptModulesChunkHash = 34,
-  HtmlPluginBeforeAssetTagGeneration = 35,
-  HtmlPluginAlterAssetTags = 36,
-  HtmlPluginAlterAssetTagGroups = 37,
-  HtmlPluginAfterTemplateExecution = 38,
-  HtmlPluginBeforeEmit = 39,
-  HtmlPluginAfterEmit = 40,
-  RuntimePluginCreateScript = 41,
-  RuntimePluginCreateLink = 42,
-  RuntimePluginLinkPreload = 43,
-  RuntimePluginLinkPrefetch = 44,
-  RsdoctorPluginModuleGraph = 45,
-  RsdoctorPluginChunkGraph = 46,
-  RsdoctorPluginModuleIds = 47,
-  RsdoctorPluginModuleSources = 48,
-  RsdoctorPluginAssets = 49
+  CompilationBeforeModuleIds = 17,
+  CompilationAdditionalTreeRuntimeRequirements = 18,
+  CompilationRuntimeRequirementInTree = 19,
+  CompilationRuntimeModule = 20,
+  CompilationChunkHash = 21,
+  CompilationChunkAsset = 22,
+  CompilationProcessAssets = 23,
+  CompilationAfterProcessAssets = 24,
+  CompilationSeal = 25,
+  CompilationAfterSeal = 26,
+  NormalModuleFactoryBeforeResolve = 27,
+  NormalModuleFactoryFactorize = 28,
+  NormalModuleFactoryResolve = 29,
+  NormalModuleFactoryAfterResolve = 30,
+  NormalModuleFactoryCreateModule = 31,
+  NormalModuleFactoryResolveForScheme = 32,
+  ContextModuleFactoryBeforeResolve = 33,
+  ContextModuleFactoryAfterResolve = 34,
+  JavascriptModulesChunkHash = 35,
+  HtmlPluginBeforeAssetTagGeneration = 36,
+  HtmlPluginAlterAssetTags = 37,
+  HtmlPluginAlterAssetTagGroups = 38,
+  HtmlPluginAfterTemplateExecution = 39,
+  HtmlPluginBeforeEmit = 40,
+  HtmlPluginAfterEmit = 41,
+  RuntimePluginCreateScript = 42,
+  RuntimePluginCreateLink = 43,
+  RuntimePluginLinkPreload = 44,
+  RuntimePluginLinkPrefetch = 45,
+  RsdoctorPluginModuleGraph = 46,
+  RsdoctorPluginChunkGraph = 47,
+  RsdoctorPluginModuleIds = 48,
+  RsdoctorPluginModuleSources = 49,
+  RsdoctorPluginAssets = 50
 }
 
 export interface RegisterJsTaps {
@@ -3022,6 +3035,7 @@ export interface RegisterJsTaps {
   registerCompilationAfterOptimizeModulesTaps: (stages: Array<number>) => Array<{ function: (() => void); stage: number; }>
   registerCompilationOptimizeTreeTaps: (stages: Array<number>) => Array<{ function: (() => Promise<void>); stage: number; }>
   registerCompilationOptimizeChunkModulesTaps: (stages: Array<number>) => Array<{ function: (() => Promise<boolean | undefined>); stage: number; }>
+  registerCompilationBeforeModuleIdsTaps: (stages: Array<number>) => Array<{ function: ((arg: JsBeforeModuleIdsArg) => JsBeforeModuleIdsResult); stage: number; }>
   registerCompilationChunkHashTaps: (stages: Array<number>) => Array<{ function: ((arg: Chunk) => Buffer); 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; }>

crates/rspack_binding_api/src/plugins/interceptor.rs

@@ -1,4 +1,5 @@
 use std::{
+  collections::HashMap,
   hash::Hash,
   ptr::NonNull,
   sync::{Arc, RwLock},
@@ -13,32 +14,32 @@ use napi::{
 use rspack_collections::{IdentifierMap, IdentifierSet};
 use rspack_core::{
   AfterResolveResult, AssetEmittedInfo, AsyncModulesArtifact, BeforeResolveResult, BindingCell,
-  BoxModule, ChunkUkey, Compilation, CompilationAdditionalTreeRuntimeRequirements,
+  BoxModule, ChunkGraph, ChunkUkey, Compilation, CompilationAdditionalTreeRuntimeRequirements,
   CompilationAdditionalTreeRuntimeRequirementsHook, CompilationAfterOptimizeModules,
   CompilationAfterOptimizeModulesHook, CompilationAfterProcessAssets,
   CompilationAfterProcessAssetsHook, CompilationAfterSeal, CompilationAfterSealHook,
-  CompilationBuildModule, CompilationBuildModuleHook, CompilationChunkAsset,
-  CompilationChunkAssetHook, CompilationChunkHash, CompilationChunkHashHook,
-  CompilationExecuteModule, CompilationExecuteModuleHook, CompilationFinishModules,
-  CompilationFinishModulesHook, CompilationId, CompilationOptimizeChunkModules,
-  CompilationOptimizeChunkModulesHook, CompilationOptimizeModules, CompilationOptimizeModulesHook,
-  CompilationOptimizeTree, CompilationOptimizeTreeHook, CompilationParams,
-  CompilationProcessAssets, CompilationProcessAssetsHook, CompilationRuntimeModule,
-  CompilationRuntimeModuleHook, CompilationRuntimeRequirementInTree,
-  CompilationRuntimeRequirementInTreeHook, CompilationSeal, CompilationSealHook,
-  CompilationStillValidModule, CompilationStillValidModuleHook, CompilationSucceedModule,
-  CompilationSucceedModuleHook, CompilerAfterEmit, CompilerAfterEmitHook, CompilerAssetEmitted,
-  CompilerAssetEmittedHook, CompilerCompilation, CompilerCompilationHook, CompilerEmit,
-  CompilerEmitHook, CompilerFinishMake, CompilerFinishMakeHook, CompilerId, CompilerMake,
-  CompilerMakeHook, CompilerShouldEmit, CompilerShouldEmitHook, CompilerThisCompilation,
-  CompilerThisCompilationHook, ContextModuleFactoryAfterResolve,
+  CompilationBeforeModuleIds, CompilationBeforeModuleIdsHook, CompilationBuildModule,
+  CompilationBuildModuleHook, CompilationChunkAsset, CompilationChunkAssetHook,
+  CompilationChunkHash, CompilationChunkHashHook, CompilationExecuteModule,
+  CompilationExecuteModuleHook, CompilationFinishModules, CompilationFinishModulesHook,
+  CompilationId, CompilationOptimizeChunkModules, CompilationOptimizeChunkModulesHook,
+  CompilationOptimizeModules, CompilationOptimizeModulesHook, CompilationOptimizeTree,
+  CompilationOptimizeTreeHook, CompilationParams, CompilationProcessAssets,
+  CompilationProcessAssetsHook, CompilationRuntimeModule, CompilationRuntimeModuleHook,
+  CompilationRuntimeRequirementInTree, CompilationRuntimeRequirementInTreeHook, CompilationSeal,
+  CompilationSealHook, CompilationStillValidModule, CompilationStillValidModuleHook,
+  CompilationSucceedModule, CompilationSucceedModuleHook, CompilerAfterEmit, CompilerAfterEmitHook,
+  CompilerAssetEmitted, CompilerAssetEmittedHook, CompilerCompilation, CompilerCompilationHook,
+  CompilerEmit, CompilerEmitHook, CompilerFinishMake, CompilerFinishMakeHook, CompilerId,
+  CompilerMake, CompilerMakeHook, CompilerShouldEmit, CompilerShouldEmitHook,
+  CompilerThisCompilation, CompilerThisCompilationHook, ContextModuleFactoryAfterResolve,
   ContextModuleFactoryAfterResolveHook, ContextModuleFactoryBeforeResolve,
   ContextModuleFactoryBeforeResolveHook, ExecuteModuleId, Module, ModuleFactoryCreateData,
-  ModuleIdentifier, NormalModuleCreateData, NormalModuleFactoryAfterResolve,
-  NormalModuleFactoryAfterResolveHook, NormalModuleFactoryBeforeResolve,
-  NormalModuleFactoryBeforeResolveHook, NormalModuleFactoryCreateModule,
-  NormalModuleFactoryCreateModuleHook, NormalModuleFactoryFactorize,
-  NormalModuleFactoryFactorizeHook, NormalModuleFactoryResolve,
+  ModuleId, ModuleIdentifier, ModuleIdsArtifact, NormalModuleCreateData,
+  NormalModuleFactoryAfterResolve, NormalModuleFactoryAfterResolveHook,
+  NormalModuleFactoryBeforeResolve, NormalModuleFactoryBeforeResolveHook,
+  NormalModuleFactoryCreateModule, NormalModuleFactoryCreateModuleHook,
+  NormalModuleFactoryFactorize, NormalModuleFactoryFactorizeHook, NormalModuleFactoryResolve,
   NormalModuleFactoryResolveForScheme, NormalModuleFactoryResolveForSchemeHook,
   NormalModuleFactoryResolveHook, NormalModuleFactoryResolveResult, ResourceData, RuntimeGlobals,
   RuntimeModule, Scheme, build_module_graph::BuildModuleGraphArtifact, parse_resource,
@@ -101,6 +102,34 @@ use crate::{
   source::JsSourceToJs,
 };
 
+#[napi(object)]
+pub struct JsModuleForIds {
+  pub identifier: String,
+}
+
+#[napi(object)]
+pub struct JsBeforeModuleIdsArg {
+  pub modules: Vec<JsModuleForIds>,
+}
+
+impl JsBeforeModuleIdsArg {
+  pub fn new(_compilation: &Compilation, modules: &IdentifierSet) -> Self {
+    Self {
+      modules: modules
+        .iter()
+        .map(|id| JsModuleForIds {
+          identifier: id.to_string(),
+        })
+        .collect(),
+    }
+  }
+}
+
+#[napi(object)]
+pub struct JsBeforeModuleIdsResult {
+  pub assignments: HashMap<String, String>,
+}
+
 #[napi(object)]
 pub struct JsTap<'f> {
   #[napi(ts_type = "(...args: any[]) => any")]
@@ -343,6 +372,7 @@ pub enum RegisterJsTapKind {
   CompilationAfterOptimizeModules,
   CompilationOptimizeTree,
   CompilationOptimizeChunkModules,
+  CompilationBeforeModuleIds,
   CompilationAdditionalTreeRuntimeRequirements,
   CompilationRuntimeRequirementInTree,
   CompilationRuntimeModule,
@@ -479,6 +509,11 @@ pub struct RegisterJsTaps {
     ts_type = "(stages: Array<number>) => Array<{ function: (() => Promise<boolean | undefined>); stage: number; }>"
   )]
   pub register_compilation_optimize_chunk_modules_taps: RegisterFunction<(), Promise<Option<bool>>>,
+  #[napi(
+    ts_type = "(stages: Array<number>) => Array<{ function: ((arg: JsBeforeModuleIdsArg) => JsBeforeModuleIdsResult); stage: number; }>"
+  )]
+  pub register_compilation_before_module_ids_taps:
+    RegisterFunction<JsBeforeModuleIdsArg, JsBeforeModuleIdsResult>,
   #[napi(
     ts_type = "(stages: Array<number>) => Array<{ function: ((arg: Chunk) => Buffer); stage: number; }>"
   )]
@@ -750,6 +785,13 @@ define_register!(
   kind = RegisterJsTapKind::CompilationOptimizeChunkModules,
   skip = true,
 );
+define_register!(
+  RegisterCompilationBeforeModuleIdsTaps,
+  tap = CompilationBeforeModuleIdsTap<JsBeforeModuleIdsArg, JsBeforeModuleIdsResult> @ CompilationBeforeModuleIdsHook,
+  cache = false,
+  kind = RegisterJsTapKind::CompilationBeforeModuleIds,
+  skip = true,
+);
 define_register!(
   RegisterCompilationAdditionalTreeRuntimeRequirementsTaps,
   tap = CompilationAdditionalTreeRuntimeRequirementsTap<JsAdditionalTreeRuntimeRequirementsArg, Option<JsAdditionalTreeRuntimeRequirementsResult>> @ CompilationAdditionalTreeRuntimeRequirementsHook,
@@ -1275,6 +1317,30 @@ impl CompilationOptimizeChunkModules for CompilationOptimizeChunkModulesTap {
   }
 }
 
+#[async_trait]
+impl CompilationBeforeModuleIds for CompilationBeforeModuleIdsTap {
+  async fn run(
+    &self,
+    compilation: &Compilation,
+    modules: &IdentifierSet,
+    module_ids: &mut ModuleIdsArtifact,
+  ) -> rspack_error::Result<()> {
+    let arg = JsBeforeModuleIdsArg::new(compilation, modules);
+    let result: JsBeforeModuleIdsResult = self.function.call_with_sync(arg).await?;
+
+    for (identifier_str, id) in result.assignments {
+      let identifier = ModuleIdentifier::from(identifier_str.as_str());
+      ChunkGraph::set_module_id(module_ids, identifier, ModuleId::from(id));
+    }
+
+    Ok(())
+  }
+
+  fn stage(&self) -> i32 {
+    self.stage
+  }
+}
+
 #[async_trait]
 impl CompilationAdditionalTreeRuntimeRequirements
   for CompilationAdditionalTreeRuntimeRequirementsTap

crates/rspack_binding_api/src/plugins/js_hooks_plugin.rs

@@ -35,6 +35,7 @@ pub struct JsHooksAdapterPlugin {
   register_compilation_after_optimize_modules_taps: RegisterCompilationAfterOptimizeModulesTaps,
   register_compilation_optimize_tree_taps: RegisterCompilationOptimizeTreeTaps,
   register_compilation_optimize_chunk_modules_taps: RegisterCompilationOptimizeChunkModulesTaps,
+  register_compilation_before_module_ids_taps: RegisterCompilationBeforeModuleIdsTaps,
   register_compilation_additional_tree_runtime_requirements_taps:
     RegisterCompilationAdditionalTreeRuntimeRequirementsTaps,
   register_compilation_runtime_requirement_in_tree_taps:
@@ -158,6 +159,10 @@ impl Plugin for JsHooksAdapterPlugin {
         .register_compilation_optimize_chunk_modules_taps
         .clone(),
     );
+    ctx
+      .compilation_hooks
+      .before_module_ids
+      .intercept(self.register_compilation_before_module_ids_taps.clone());
     ctx
       .compilation_hooks
       .additional_tree_runtime_requirements
@@ -291,6 +296,9 @@ impl Plugin for JsHooksAdapterPlugin {
     self
       .register_compilation_optimize_chunk_modules_taps
       .clear_cache();
+    self
+      .register_compilation_before_module_ids_taps
+      .clear_cache();
     self
       .register_compilation_additional_tree_runtime_requirements_taps
       .clear_cache();
@@ -542,6 +550,10 @@ impl JsHooksAdapterPlugin {
             register_js_taps.register_compilation_optimize_chunk_modules_taps,
             non_skippable_registers.clone(),
           ),
+        register_compilation_before_module_ids_taps: RegisterCompilationBeforeModuleIdsTaps::new(
+          register_js_taps.register_compilation_before_module_ids_taps,
+          non_skippable_registers.clone(),
+        ),
         register_compilation_additional_tree_runtime_requirements_taps:
           RegisterCompilationAdditionalTreeRuntimeRequirementsTaps::new(
             register_js_taps.register_compilation_additional_tree_runtime_requirements_taps,

crates/rspack_core/src/compilation/mod.rs

@@ -105,6 +105,7 @@ define_hook!(CompilationAfterOptimizeModules: Series(compilation: &Compilation))
 define_hook!(CompilationOptimizeChunks: SeriesBail(compilation: &mut Compilation) -> bool);
 define_hook!(CompilationOptimizeTree: Series(compilation: &Compilation));
 define_hook!(CompilationOptimizeChunkModules: SeriesBail(compilation: &mut Compilation) -> bool);
+define_hook!(CompilationBeforeModuleIds: Series(compilation: &Compilation, modules: &IdentifierSet, module_ids: &mut ModuleIdsArtifact));
 define_hook!(CompilationModuleIds: Series(compilation: &Compilation, module_ids: &mut ModuleIdsArtifact, diagnostics: &mut Vec<Diagnostic>));
 define_hook!(CompilationChunkIds: Series(compilation: &Compilation, chunk_by_ukey: &mut ChunkByUkey, named_chunk_ids_artifact: &mut ChunkNamedIdArtifact, diagnostics: &mut Vec<Diagnostic>));
 define_hook!(CompilationRuntimeModule: Series(compilation: &Compilation, module: &ModuleIdentifier, chunk: &ChunkUkey, runtime_modules: &mut IdentifierMap<Box<dyn RuntimeModule>>));
@@ -142,6 +143,7 @@ pub struct CompilationHooks {
   pub optimize_chunks: CompilationOptimizeChunksHook,
   pub optimize_tree: CompilationOptimizeTreeHook,
   pub optimize_chunk_modules: CompilationOptimizeChunkModulesHook,
+  pub before_module_ids: CompilationBeforeModuleIdsHook,
   pub module_ids: CompilationModuleIdsHook,
   pub chunk_ids: CompilationChunkIdsHook,
   pub runtime_module: CompilationRuntimeModuleHook,

crates/rspack_core/src/compilation/module_ids/mod.rs

@@ -3,6 +3,29 @@ use async_trait::async_trait;
 use super::*;
 use crate::{cache::Cache, compilation::pass::PassExt};
 
+/// Collects module identifiers that need ID assignment.
+/// A module needs an ID if:
+/// - It doesn't already have one assigned
+/// - It needs an ID (need_id() returns true)
+/// - It's part of at least one chunk
+fn get_modules_needing_ids(
+  compilation: &Compilation,
+  module_ids_artifact: &ModuleIdsArtifact,
+) -> IdentifierSet {
+  let chunk_graph = &compilation.chunk_graph;
+  compilation
+    .get_module_graph()
+    .modules()
+    .values()
+    .filter(|m| {
+      m.need_id()
+        && ChunkGraph::get_module_id(module_ids_artifact, m.identifier()).is_none()
+        && chunk_graph.get_number_of_module_chunks(m.identifier()) != 0
+    })
+    .map(|m| m.identifier())
+    .collect()
+}
+
 pub struct ModuleIdsPass;
 
 #[async_trait]
@@ -24,6 +47,23 @@ impl PassExt for ModuleIdsPass {
       compilation.module_ids_artifact.clear();
     }
 
+    let mut module_ids_artifact = mem::take(&mut compilation.module_ids_artifact);
+
+    // Call beforeModuleIds hook - allows plugins to assign custom IDs
+    let modules_needing_ids = get_modules_needing_ids(compilation, &module_ids_artifact);
+    compilation
+      .plugin_driver
+      .clone()
+      .compilation_hooks
+      .before_module_ids
+      .call(compilation, &modules_needing_ids, &mut module_ids_artifact)
+      .await
+      .map_err(|e| e.wrap_err("caused by plugins in Compilation.hooks.beforeModuleIds"))?;
+
+    // Put artifact back so moduleIds plugins can see custom IDs from beforeModuleIds
+    // when they call get_used_module_ids_and_modules
+    compilation.module_ids_artifact = module_ids_artifact;
+
     let mut diagnostics = vec![];
     let mut module_ids_artifact = mem::take(&mut compilation.module_ids_artifact);
     compilation

crates/rspack_ids/src/deterministic_module_ids_plugin.rs

@@ -9,7 +9,7 @@ use rspack_hook::{plugin, plugin_hook};
 
 use crate::id_helpers::{
   assign_deterministic_ids, compare_modules_by_pre_order_index_or_identifier, get_full_module_name,
-  get_used_module_ids_and_modules,
+  get_used_module_ids_and_modules_with_artifact,
 };
 
 #[plugin]
@@ -34,7 +34,8 @@ async fn module_ids(
     module_ids.clear();
   }
 
-  let (mut used_ids, modules) = get_used_module_ids_and_modules(compilation, None);
+  let (mut used_ids, modules) =
+    get_used_module_ids_and_modules_with_artifact(compilation, module_ids, None);
 
   let mut module_ids_map = std::mem::take(module_ids);
   let context = compilation.options.context.as_ref();