refactor: rm UNASSOCIATED_MODULE_INSTANCE_REFS (#8327)

crates/node_binding/src/plugins/interceptor.rs

@@ -33,7 +33,7 @@ use rspack_core::{
   CompilationBuildModule, CompilationBuildModuleHook, CompilationChunkAsset,
   CompilationChunkAssetHook, CompilationChunkHash, CompilationChunkHashHook,
   CompilationExecuteModule, CompilationExecuteModuleHook, CompilationFinishModules,
-  CompilationFinishModulesHook, CompilationOptimizeChunkModules,
+  CompilationFinishModulesHook, CompilationId, CompilationOptimizeChunkModules,
   CompilationOptimizeChunkModulesHook, CompilationOptimizeModules, CompilationOptimizeModulesHook,
   CompilationOptimizeTree, CompilationOptimizeTreeHook, CompilationParams,
   CompilationProcessAssets, CompilationProcessAssetsHook, CompilationRuntimeModule,
@@ -1011,10 +1011,14 @@ impl CompilerAssetEmitted for CompilerAssetEmittedTap {
 
 #[async_trait]
 impl CompilationBuildModule for CompilationBuildModuleTap {
-  async fn run(&self, module: &mut BoxModule) -> rspack_error::Result<()> {
+  async fn run(
+    &self,
+    compilation_id: CompilationId,
+    module: &mut BoxModule,
+  ) -> rspack_error::Result<()> {
     self
       .function
-      .call_with_sync(JsModuleWrapper::new(module.as_ref(), None))
+      .call_with_sync(JsModuleWrapper::new(module.as_ref(), compilation_id, None))
       .await
   }
 
@@ -1025,10 +1029,14 @@ impl CompilationBuildModule for CompilationBuildModuleTap {
 
 #[async_trait]
 impl CompilationStillValidModule for CompilationStillValidModuleTap {
-  async fn run(&self, module: &mut BoxModule) -> rspack_error::Result<()> {
+  async fn run(
+    &self,
+    compilation_id: CompilationId,
+    module: &mut BoxModule,
+  ) -> rspack_error::Result<()> {
     self
       .function
-      .call_with_sync(JsModuleWrapper::new(module.as_ref(), None))
+      .call_with_sync(JsModuleWrapper::new(module.as_ref(), compilation_id, None))
       .await
   }
 
@@ -1039,10 +1047,14 @@ impl CompilationStillValidModule for CompilationStillValidModuleTap {
 
 #[async_trait]
 impl CompilationSucceedModule for CompilationSucceedModuleTap {
-  async fn run(&self, module: &mut BoxModule) -> rspack_error::Result<()> {
+  async fn run(
+    &self,
+    compilation_id: CompilationId,
+    module: &mut BoxModule,
+  ) -> rspack_error::Result<()> {
     self
       .function
-      .call_with_sync(JsModuleWrapper::new(module.as_ref(), None))
+      .call_with_sync(JsModuleWrapper::new(module.as_ref(), compilation_id, None))
       .await
   }
 

crates/rspack_binding_options/src/options/raw_builtins/raw_lazy_compilation.rs

@@ -4,7 +4,7 @@ use napi::{
 };
 use napi_derive::napi;
 use rspack_binding_values::JsModuleWrapper;
-use rspack_core::ModuleIdentifier;
+use rspack_core::{CompilationId, ModuleIdentifier};
 use rspack_napi::threadsafe_function::ThreadsafeFunction;
 use rspack_plugin_lazy_compilation::{
   backend::{Backend, ModuleInfo},
@@ -41,10 +41,10 @@ pub struct LazyCompilationTestFn {
 }
 
 impl LazyCompilationTestCheck for LazyCompilationTestFn {
-  fn test(&self, m: &dyn rspack_core::Module) -> bool {
+  fn test(&self, compilation_id: CompilationId, m: &dyn rspack_core::Module) -> bool {
     let res = self
       .tsfn
-      .blocking_call_with_sync(JsModuleWrapper::new(m, None))
+      .blocking_call_with_sync(JsModuleWrapper::new(m, compilation_id, None))
       .expect("failed to invoke lazyCompilation.test");
 
     res.unwrap_or(false)

crates/rspack_binding_options/src/options/raw_split_chunks/raw_split_chunk_cache_group_test.rs

@@ -28,7 +28,11 @@ impl FromNapiValue for JsCacheGroupTestCtx {
 impl<'a> From<CacheGroupTestFnCtx<'a>> for JsCacheGroupTestCtx {
   fn from(value: CacheGroupTestFnCtx<'a>) -> Self {
     JsCacheGroupTestCtx {
-      module: JsModuleWrapper::new(value.module, Some(value.compilation)),
+      module: JsModuleWrapper::new(
+        value.module,
+        value.compilation.id(),
+        Some(value.compilation),
+      ),
     }
   }
 }

crates/rspack_binding_options/src/options/raw_split_chunks/raw_split_chunk_name.rs

@@ -35,7 +35,11 @@ impl FromNapiValue for RawChunkOptionNameCtx {
 impl<'a> From<ChunkNameGetterFnCtx<'a>> for RawChunkOptionNameCtx {
   fn from(value: ChunkNameGetterFnCtx<'a>) -> Self {
     RawChunkOptionNameCtx {
-      module: JsModuleWrapper::new(value.module, Some(value.compilation)),
+      module: JsModuleWrapper::new(
+        value.module,
+        value.compilation.id(),
+        Some(value.compilation),
+      ),
       chunks: value
         .chunks
         .iter()

crates/rspack_binding_options/src/plugins/js_loader/context.rs

@@ -96,7 +96,7 @@ impl TryFrom<&mut LoaderContext<RunnerContext>> for JsLoaderContext {
     Ok(JsLoaderContext {
       resource_data: cx.resource_data.as_ref().into(),
       module_identifier: module.identifier().to_string(),
-      module: JsModuleWrapper::new(module, None),
+      module: JsModuleWrapper::new(module, cx.context.compilation_id, None),
       hot: cx.hot,
       content: match cx.content() {
         Some(c) => Either::B(c.to_owned().into_bytes().into()),

crates/rspack_binding_values/src/chunk_graph.rs

@@ -21,7 +21,7 @@ pub fn get_chunk_modules(
 
   return modules
     .iter()
-    .map(|module| JsModuleWrapper::new(module.as_ref(), Some(compilation)))
+    .map(|module| JsModuleWrapper::new(module.as_ref(), compilation.id(), Some(compilation)))
     .collect::<Vec<_>>();
 }
 
@@ -42,7 +42,7 @@ pub fn get_chunk_entry_modules(
   return modules
     .iter()
     .filter_map(|module| module_graph.module_by_identifier(module))
-    .map(|module| JsModuleWrapper::new(module.as_ref(), Some(compilation)))
+    .map(|module| JsModuleWrapper::new(module.as_ref(), compilation.id(), Some(compilation)))
     .collect::<Vec<_>>();
 }
 
@@ -86,7 +86,7 @@ pub fn get_chunk_modules_iterable_by_source_type(
         SourceType::from(source_type.as_str()),
         &compilation.get_module_graph(),
       )
-      .map(|module| JsModuleWrapper::new(module, Some(compilation)))
+      .map(|module| JsModuleWrapper::new(module, compilation.id(), Some(compilation)))
       .collect(),
   )
 }

crates/rspack_binding_values/src/chunk_group.rs

@@ -57,7 +57,7 @@ impl JsChunkGroup {
             let module = compilation
               .module_by_identifier(&module_id)
               .unwrap_or_else(|| panic!("failed to retrieve module by id: {}", module_id));
-            JsModuleWrapper::new(module.as_ref(), Some(compilation))
+            JsModuleWrapper::new(module.as_ref(), compilation.id(), Some(compilation))
           }),
           request: origin.request.clone(),
         })

crates/rspack_binding_values/src/compilation/mod.rs

@@ -168,9 +168,9 @@ impl JsCompilation {
         .modules()
         .keys()
         .filter_map(|module_id| {
-          compilation
-            .module_by_identifier(module_id)
-            .map(|module| JsModuleWrapper::new(module.as_ref(), Some(compilation)))
+          compilation.module_by_identifier(module_id).map(|module| {
+            JsModuleWrapper::new(module.as_ref(), compilation.id(), Some(compilation))
+          })
         })
         .collect::<Vec<_>>(),
     )
@@ -185,9 +185,9 @@ impl JsCompilation {
         .built_modules
         .iter()
         .filter_map(|module_id| {
-          compilation
-            .module_by_identifier(module_id)
-            .map(|module| JsModuleWrapper::new(module.as_ref(), Some(compilation)))
+          compilation.module_by_identifier(module_id).map(|module| {
+            JsModuleWrapper::new(module.as_ref(), compilation.id(), Some(compilation))
+          })
         })
         .collect::<Vec<_>>(),
     )
@@ -592,13 +592,15 @@ impl JsCompilation {
     let compilation = self.as_mut()?;
 
     callbackify(env, f, async {
+      let compilation_id = compilation.id();
+
       let mut modules = compilation
         .rebuild_module(
           IdentifierSet::from_iter(module_identifiers.into_iter().map(ModuleIdentifier::from)),
           |modules| {
             modules
               .into_iter()
-              .map(|module| JsModuleWrapper::new(module.as_ref(), None))
+              .map(|module| JsModuleWrapper::new(module.as_ref(), compilation_id, None))
               .collect::<Vec<_>>()
           },
         )

crates/rspack_binding_values/src/module.rs

@@ -83,16 +83,14 @@ impl JsDependenciesBlock {
 pub struct JsModule {
   identifier: ModuleIdentifier,
   module: NonNull<dyn Module>,
+  compilation_id: CompilationId,
   compilation: Option<NonNull<Compilation>>,
 }
 
 impl JsModule {
-  fn attach(&mut self, compilation: *const Compilation) {
+  fn attach(&mut self, compilation: NonNull<Compilation>) {
     if self.compilation.is_none() {
-      self.compilation = Some(
-        #[allow(clippy::unwrap_used)]
-        NonNull::new(compilation as *mut Compilation).unwrap(),
-      );
+      self.compilation = Some(compilation);
     }
   }
 
@@ -323,7 +321,9 @@ impl JsModule {
             .filter_map(|inner_module_info| {
               compilation
                 .module_by_identifier(&inner_module_info.id)
-                .map(|module| JsModuleWrapper::new(module.as_ref(), Some(compilation)))
+                .map(|module| {
+                  JsModuleWrapper::new(module.as_ref(), compilation.id(), Some(compilation))
+                })
             })
             .collect::<Vec<_>>();
           Either::A(inner_modules)
@@ -341,14 +341,12 @@ impl JsModule {
   }
 }
 
-type ModuleInstanceRefs = IdentifierMap<OneShotRef<ClassInstance<JsModule>>>;
+type ModuleInstanceRefs = IdentifierMap<OneShotRef<JsModule>>;
 
 type ModuleInstanceRefsByCompilationId = RefCell<HashMap<CompilationId, ModuleInstanceRefs>>;
 
 thread_local! {
   static MODULE_INSTANCE_REFS: ModuleInstanceRefsByCompilationId = Default::default();
-
-  static UNASSOCIATED_MODULE_INSTANCE_REFS: RefCell<ModuleInstanceRefs> = Default::default();
 }
 
 // The difference between JsModuleWrapper and JsModule is:
@@ -358,20 +356,26 @@ thread_local! {
 pub struct JsModuleWrapper {
   identifier: ModuleIdentifier,
   module: NonNull<dyn Module>,
+  compilation_id: CompilationId,
   compilation: Option<NonNull<Compilation>>,
 }
 
 unsafe impl Send for JsModuleWrapper {}
 
 impl JsModuleWrapper {
-  pub fn new(module: &dyn Module, compilation: Option<&Compilation>) -> Self {
+  pub fn new(
+    module: &dyn Module,
+    compilation_id: CompilationId,
+    compilation: Option<&Compilation>,
+  ) -> Self {
     #[allow(clippy::not_unsafe_ptr_arg_deref)]
     let identifier = module.identifier();
 
     #[allow(clippy::unwrap_used)]
     Self {
       identifier,
       module: NonNull::new(module as *const dyn Module as *mut dyn Module).unwrap(),
+      compilation_id,
       compilation: compilation
         .map(|c| NonNull::new(c as *const Compilation as *mut Compilation).unwrap()),
     }
@@ -398,86 +402,52 @@ impl ToNapiValue for JsModuleWrapper {
   unsafe fn to_napi_value(env: sys::napi_env, val: Self) -> Result<sys::napi_value> {
     let module = unsafe { val.module.as_ref() };
 
-    match val.compilation {
-      Some(compilation_ptr) => MODULE_INSTANCE_REFS.with(|refs| {
-        let compilation = unsafe { compilation_ptr.as_ref() };
-
-        let mut refs_by_compilation_id = refs.borrow_mut();
-        let entry = refs_by_compilation_id.entry(compilation.id());
-        let refs = match entry {
-          std::collections::hash_map::Entry::Occupied(entry) => entry.into_mut(),
-          std::collections::hash_map::Entry::Vacant(entry) => {
-            let refs = IdentifierMap::default();
-            entry.insert(refs)
-          }
-        };
-
-        UNASSOCIATED_MODULE_INSTANCE_REFS.with(|ref_cell| {
-          let mut unassociated_refs = ref_cell.borrow_mut();
-          if let Some(unassociated_ref) = unassociated_refs.remove(&module.identifier()) {
-            let mut instance = unassociated_ref.from_napi_value()?;
-            instance.as_mut().attach(compilation_ptr.as_ptr());
-
-            let napi_value = ToNapiValue::to_napi_value(env, &unassociated_ref);
-            refs.insert(module.identifier(), unassociated_ref);
-            napi_value
+    MODULE_INSTANCE_REFS.with(|refs| {
+      let mut refs_by_compilation_id = refs.borrow_mut();
+      let entry = refs_by_compilation_id.entry(val.compilation_id);
+      let refs = match entry {
+        std::collections::hash_map::Entry::Occupied(entry) => entry.into_mut(),
+        std::collections::hash_map::Entry::Vacant(entry) => {
+          let refs = IdentifierMap::default();
+          entry.insert(refs)
+        }
+      };
+
+      match refs.entry(module.identifier()) {
+        std::collections::hash_map::Entry::Occupied(entry) => {
+          let r = entry.get();
+          let instance = r.from_napi_mut_ref()?;
+          if let Some(compilation) = val.compilation {
+            instance.attach(compilation);
           } else {
-            match refs.entry(module.identifier()) {
-              std::collections::hash_map::Entry::Occupied(entry) => {
-                let r = entry.get();
-                ToNapiValue::to_napi_value(env, r)
-              }
-              std::collections::hash_map::Entry::Vacant(entry) => {
-                let instance: ClassInstance<JsModule> = JsModule {
-                  identifier: val.identifier,
-                  module: val.module,
-                  compilation: Some(compilation_ptr),
-                }
-                .into_instance(Env::from_raw(env))?;
-                let r = entry.insert(OneShotRef::new(env, instance)?);
-                ToNapiValue::to_napi_value(env, r)
-              }
-            }
-          }
-        })
-      }),
-      None => UNASSOCIATED_MODULE_INSTANCE_REFS.with(|ref_cell| {
-        let mut refs = ref_cell.borrow_mut();
-        match refs.entry(module.identifier()) {
-          std::collections::hash_map::Entry::Occupied(entry) => {
-            let r = entry.get();
-
-            let mut instance: ClassInstance<JsModule> = r.from_napi_value()?;
-            if !std::ptr::addr_eq(instance.module.as_ptr(), val.module.as_ptr()) {
-              instance.module = val.module;
-            }
-            ToNapiValue::to_napi_value(env, r)
-          }
-          std::collections::hash_map::Entry::Vacant(entry) => {
-            let instance = JsModule {
-              identifier: val.identifier,
-              module: val.module,
-              compilation: None,
-            }
-            .into_instance(Env::from_raw(env))?;
-            let r = entry.insert(OneShotRef::new(env, instance)?);
-            ToNapiValue::to_napi_value(env, r)
+            instance.module = val.module;
           }
+          ToNapiValue::to_napi_value(env, r)
         }
-      }),
-    }
+        std::collections::hash_map::Entry::Vacant(entry) => {
+          let js_module = JsModule {
+            identifier: val.identifier,
+            module: val.module,
+            compilation_id: val.compilation_id,
+            compilation: val.compilation,
+          };
+          let r = entry.insert(OneShotRef::new(env, js_module)?);
+          ToNapiValue::to_napi_value(env, r)
+        }
+      }
+    })
   }
 }
 
 impl FromNapiValue for JsModuleWrapper {
   unsafe fn from_napi_value(env: sys::napi_env, napi_val: sys::napi_value) -> Result<Self> {
     let instance: ClassInstance<JsModule> = FromNapiValue::from_napi_value(env, napi_val)?;
-    let module = instance.module;
 
     Ok(JsModuleWrapper {
       identifier: instance.identifier,
       #[allow(clippy::unwrap_used)]
-      module: NonNull::new(module.as_ptr()).unwrap(),
+      module: instance.module,
+      compilation_id: instance.compilation_id,
       compilation: instance.compilation,
     })
   }