feat: store pointer in ModuleDTO

crates/node_binding/binding.d.ts

@@ -19,12 +19,6 @@ export class ExternalObject<T> {
     [K: symbol]: T
   }
 }
-export class DependenciesBlockDto {
-  get dependencies(): Array<JsDependency>
-  get blocks(): Array<DependenciesBlockDto>
-}
-export type DependenciesBlockDTO = DependenciesBlockDto
-
 export class DependenciesDto {
   get fileDependencies(): Array<string>
   get addedFileDependencies(): Array<string>
@@ -139,6 +133,11 @@ export class JsContextModuleFactoryBeforeResolveData {
   set recursive(recursive: boolean)
 }
 
+export class JsDependenciesBlock {
+  get dependencies(): Array<JsDependency>
+  get blocks(): Array<JsDependenciesBlock>
+}
+
 export class JsDependency {
   get type(): string
   get category(): string
@@ -194,7 +193,7 @@ export class ModuleDto {
   get factoryMeta(): JsFactoryMeta | undefined
   get type(): string
   get layer(): string | undefined
-  get blocks(): Array<DependenciesBlockDto>
+  get blocks(): Array<JsDependenciesBlock>
   size(ty?: string | undefined | null): number
   get modules(): ModuleDTO[] | undefined
 }

crates/node_binding/src/lib.rs

@@ -11,7 +11,7 @@ use std::sync::Mutex;
 use compiler::{Compiler, CompilerState, CompilerStateGuard};
 use napi::bindgen_prelude::*;
 use rspack_binding_options::BuiltinPlugin;
-use rspack_core::{Compilation, PluginExt};
+use rspack_core::PluginExt;
 use rspack_error::Diagnostic;
 use rspack_fs_node::{AsyncNodeWritableFileSystem, ThreadsafeNodeFS};
 
@@ -171,7 +171,7 @@ impl Rspack {
       Ok(unsafe { s.compiler.as_mut().get_unchecked_mut() })
     })?;
 
-    self.cleanup_last_compilation(&compiler.compilation);
+    self.cleanup_last_compilation();
 
     // SAFETY:
     // 1. `Compiler` is pinned and stored on the heap.
@@ -182,8 +182,8 @@ impl Rspack {
     )
   }
 
-  fn cleanup_last_compilation(&self, compilation: &Compilation) {
-    JsCompilationWrapper::cleanup(compilation.id());
+  fn cleanup_last_compilation(&self) {
+    JsCompilationWrapper::cleanup_last_compilation();
   }
 }
 

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

@@ -29,7 +29,7 @@ impl FromNapiValue for JsCacheGroupTestCtx {
 impl<'a> From<CacheGroupTestFnCtx<'a>> for JsCacheGroupTestCtx {
   fn from(value: CacheGroupTestFnCtx<'a>) -> Self {
     JsCacheGroupTestCtx {
-      module: ModuleDTOWrapper::new(value.module.identifier(), value.compilation),
+      module: ModuleDTOWrapper::new(value.module, value.compilation),
     }
   }
 }

crates/rspack_binding_values/src/chunk_graph.rs

@@ -17,7 +17,7 @@ pub fn get_chunk_modules(js_chunk_ukey: u32, compilation: &JsCompilation) -> Vec
 
   return modules
     .iter()
-    .map(|module| ModuleDTOWrapper::new(module.identifier(), compilation))
+    .map(|module| ModuleDTOWrapper::new(module.as_ref(), compilation))
     .collect::<Vec<_>>();
 }
 

crates/rspack_binding_values/src/chunk_group.rs

@@ -53,9 +53,12 @@ impl JsChunkGroup {
         .origins()
         .iter()
         .map(|origin| JsChunkGroupOrigin {
-          module: origin
-            .module_id
-            .map(|module_id| ModuleDTOWrapper::new(module_id, compilation)),
+          module: origin.module_id.map(|module_id| {
+            let module = compilation
+              .module_by_identifier(&module_id)
+              .unwrap_or_else(|| panic!("failed to retrieve module by id: {}", module_id));
+            ModuleDTOWrapper::new(module.as_ref(), compilation)
+          }),
           request: origin.request.clone(),
         })
         .collect::<Vec<_>>(),

crates/rspack_binding_values/src/compilation/mod.rs

@@ -15,7 +15,6 @@ use rspack_core::get_chunk_group_from_ukey;
 use rspack_core::rspack_sources::BoxSource;
 use rspack_core::AssetInfo;
 use rspack_core::ChunkUkey;
-use rspack_core::CompilationId;
 use rspack_core::ModuleIdentifier;
 use rspack_error::Diagnostic;
 use rspack_napi::napi::bindgen_prelude::*;
@@ -142,8 +141,12 @@ impl JsCompilation {
       .get_module_graph()
       .modules()
       .keys()
-      .cloned()
-      .map(|module_id| ModuleDTOWrapper::new(module_id, self.0))
+      .filter_map(|module_id| {
+        self
+          .0
+          .module_by_identifier(module_id)
+          .map(|module| ModuleDTOWrapper::new(module.as_ref(), self.0))
+      })
       .collect::<Vec<_>>()
   }
 
@@ -153,7 +156,12 @@ impl JsCompilation {
       .0
       .built_modules
       .iter()
-      .map(|module_id| ModuleDTOWrapper::new(*module_id, self.0))
+      .filter_map(|module_id| {
+        self
+          .0
+          .module_by_identifier(module_id)
+          .map(|module| ModuleDTOWrapper::new(module.as_ref(), self.0))
+      })
       .collect::<Vec<_>>()
   }
 
@@ -567,19 +575,8 @@ impl JsCompilation {
   }
 }
 
-#[derive(Default)]
-struct CompilationInstanceRefs(RefCell<HashMap<CompilationId, OneShotRef>>);
-
-impl Drop for CompilationInstanceRefs {
-  fn drop(&mut self) {
-    // cleanup references to be executed in cases of panic or unexpected termination
-    let mut refs = self.0.borrow_mut();
-    refs.drain();
-  }
-}
-
 thread_local! {
-  static COMPILATION_INSTANCE_REFS: CompilationInstanceRefs = Default::default();
+  static COMPILATION_INSTANCE_REF: RefCell<Option<OneShotRef>> = Default::default();
 }
 
 // The difference between JsCompilationWrapper and JsCompilation is:
@@ -601,42 +598,29 @@ impl JsCompilationWrapper {
     })
   }
 
-  pub fn cleanup(compilation_id: CompilationId) {
-    COMPILATION_INSTANCE_REFS.with(|ref_cell| {
-      let mut refs = ref_cell.0.borrow_mut();
-      refs.remove(&compilation_id);
+  pub fn cleanup_last_compilation() {
+    COMPILATION_INSTANCE_REF.with(|ref_cell| {
+      let _ = ref_cell.borrow_mut().take();
     });
-    ModuleDTOWrapper::cleanup(compilation_id);
+    ModuleDTOWrapper::cleanup_last_compilation();
   }
 }
 
 impl ToNapiValue for JsCompilationWrapper {
   unsafe fn to_napi_value(env: sys::napi_env, val: Self) -> Result<sys::napi_value> {
-    COMPILATION_INSTANCE_REFS.with(|ref_cell| {
-      let mut env_wrapper = Env::from_raw(env);
-      let mut refs = ref_cell.0.borrow_mut();
-      let compilation_id = val.0.id();
-      let mut vacant = false;
-      let napi_value = match refs.entry(compilation_id) {
-        std::collections::hash_map::Entry::Occupied(entry) => {
-          let r = entry.get();
-          ToNapiValue::to_napi_value(env, r)
-        }
-        std::collections::hash_map::Entry::Vacant(entry) => {
-          vacant = true;
+    COMPILATION_INSTANCE_REF.with(|ref_cell| {
+      let mut one_shot_ref = ref_cell.borrow_mut();
+      match &*one_shot_ref {
+        Some(r) => ToNapiValue::to_napi_value(env, r),
+        None => {
+          let env_wrapper = Env::from_raw(env);
           let instance = JsCompilation(val.0).into_instance(env_wrapper)?;
           let napi_value = ToNapiValue::to_napi_value(env, instance)?;
           let r = OneShotRef::new(env, napi_value)?;
-          let r = entry.insert(r);
-          ToNapiValue::to_napi_value(env, r)
+          *one_shot_ref = Some(r);
+          Ok(napi_value)
         }
-      };
-      if vacant {
-        // cleanup references to be executed when the JS thread exits normally
-        let _ = env_wrapper
-          .add_env_cleanup_hook((), move |_| JsCompilationWrapper::cleanup(compilation_id));
       }
-      napi_value
     })
   }
 }