refactor: 🎨 add back chunk id in replace info

crates/mako/src/generate/chunk_graph.rs

@@ -63,6 +63,12 @@ impl ChunkGraph {
         self.graph.node_weights().find(|c| c.has_module(module_id))
     }
 
+    pub fn get_async_chunk_for_module(&self, module_id: &ModuleId) -> Option<&Chunk> {
+        self.graph
+            .node_weights()
+            .find(|c| c.has_module(module_id) && matches!(c.chunk_type, ChunkType::Async))
+    }
+
     // pub fn get_chunk_by_id(&self, id: &String) -> Option<&Chunk> {
     //     self.graph.node_weights().find(|c| c.id.id.eq(id))
     // }

crates/mako/src/generate/transform.rs

@@ -87,24 +87,42 @@ pub fn transform_modules_in_thread(
             let mut resolved_deps: HashMap<String, ResolvedReplaceInfo> = deps
                 .into_iter()
                 .map(|(id, dep)| {
-                    let to_replace_source = match &dep.resolve_type {
+                    let replace_info = match &dep.resolve_type {
                         ResolveType::Worker(import_options) => {
                             let chunk_id = match import_options.get_chunk_name() {
                                 Some(chunk_name) => generate_module_id(chunk_name, &context),
                                 None => id.generate(&context),
                             };
                             let chunk_graph = context.chunk_graph.read().unwrap();
-                            chunk_graph.chunk(&chunk_id.into()).unwrap().filename()
+                            let chunk_name =
+                                chunk_graph.chunk(&chunk_id.into()).unwrap().filename();
+
+                            ResolvedReplaceInfo {
+                                chunk_id: None,
+                                to_replace_source: chunk_name,
+                                resolved_module_id: id.clone(),
+                            }
                         }
-                        _ => id.generate(&context),
-                    };
+                        ResolveType::DynamicImport(import_options) => {
+                            let chunk_id = Some(match import_options.get_chunk_name() {
+                                Some(chunk_name) => generate_module_id(chunk_name, &context),
+                                None => id.generate(&context),
+                            });
 
-                    let resolve_replace_info = ResolvedReplaceInfo {
-                        to_replace_source,
-                        resolved_module_id: id.clone(),
+                            ResolvedReplaceInfo {
+                                chunk_id,
+                                to_replace_source: id.generate(&context),
+                                resolved_module_id: id.clone(),
+                            }
+                        }
+                        _ => ResolvedReplaceInfo {
+                            chunk_id: None,
+                            to_replace_source: id.generate(&context),
+                            resolved_module_id: id.clone(),
+                        },
                     };
 
-                    (dep.source.clone(), resolve_replace_info)
+                    (dep.source.clone(), replace_info)
                 })
                 .collect();
             insert_swc_helper_replace(&mut resolved_deps, &context);
@@ -162,6 +180,7 @@ fn insert_swc_helper_replace(
         map.insert(
             m_id.id.clone(),
             ResolvedReplaceInfo {
+                chunk_id: None,
                 to_replace_source: m_id.generate(context),
                 resolved_module_id: m_id,
             },

crates/mako/src/plugins/bundless_compiler.rs

@@ -83,6 +83,7 @@ impl BundlessCompiler {
                         Ok((
                             dep.source.clone(),
                             ResolvedReplaceInfo {
+                                chunk_id: None,
                                 to_replace_source: replacement,
                                 resolved_module_id: id.clone(),
                             },

crates/mako/src/plugins/central_ensure.rs

@@ -28,7 +28,7 @@ pub fn module_ensure_map(context: &Arc<Context>) -> anyhow::Result<BTreeMap<Stri
             .any(|(_, dep)| dep.resolve_type.is_dynamic_esm());
 
         if be_dynamic_imported {
-            cg.get_chunk_for_module(&module.id)
+            cg.get_async_chunk_for_module(&module.id)
                 .iter()
                 .for_each(|chunk| {
                     let deps_chunks = cg

crates/mako/src/visitors/dep_replacer.rs

@@ -25,6 +25,7 @@ pub struct DepReplacer<'a> {
 
 #[derive(Debug, Clone)]
 pub struct ResolvedReplaceInfo {
+    pub chunk_id: Option<String>,
     pub to_replace_source: String,
     pub resolved_module_id: ModuleId,
 }
@@ -347,6 +348,7 @@ try {
         hashmap! {
             key.to_string() =>
             ResolvedReplaceInfo {
+                chunk_id: None,
                 to_replace_source: module_id.into(),
                 resolved_module_id: "".into(),
             }

crates/mako/src/visitors/dynamic_import.rs

@@ -9,11 +9,11 @@ use swc_core::ecma::utils::{
 };
 use swc_core::ecma::visit::{VisitMut, VisitMutWith};
 
-use super::dep_replacer::miss_throw_stmt;
+use super::dep_replacer::{miss_throw_stmt, ResolvedReplaceInfo};
 use crate::ast::utils::{is_dynamic_import, promise_all, require_ensure};
 use crate::ast::DUMMY_CTXT;
 use crate::compiler::Context;
-use crate::module::ModuleId;
+use crate::generate::chunk::ChunkId;
 use crate::visitors::dep_replacer::DependenciesToReplace;
 
 pub struct DynamicImport<'a> {
@@ -118,7 +118,7 @@ impl<'a> VisitMut for DynamicImport<'a> {
                         {
                             self.central_ensure(&generated_module_id)
                         } else {
-                            self.inline_ensure(&resolved_info.resolved_module_id, &self.context)
+                            self.inline_ensure(resolved_info, &self.context)
                         };
 
                         let lazy_require_call =
@@ -155,11 +155,12 @@ impl DynamicImport<'_> {
     // build the Promise.all([...]) part
     // Promise.all([ require.ensure("id") ]).then(require.bind(require, "id"))
     // Promise.all([ require.ensure("d1"), require.ensure("id)]).then(require.bind(require, "id"))
-    fn inline_ensure(&self, module_id: &ModuleId, context: &Arc<Context>) -> Expr {
+    fn inline_ensure(&self, replace_info: &ResolvedReplaceInfo, context: &Arc<Context>) -> Expr {
         let chunk_graph = context.chunk_graph.read().unwrap();
 
+        let init_chunk_id: ChunkId = replace_info.chunk_id.as_ref().unwrap().clone().into();
         let chunk_ids = {
-            let chunk = chunk_graph.get_chunk_for_module(module_id);
+            let chunk = chunk_graph.chunk(&init_chunk_id);
             let chunk_ids = match chunk {
                 Some(chunk) => {
                     [
@@ -247,10 +248,12 @@ Promise.all([
         let dep_to_replace = DependenciesToReplace {
             resolved: maplit::hashmap! {
                 "@swc/helpers/_/_interop_require_wildcard".to_string() => ResolvedReplaceInfo {
+                    chunk_id: None,
                     to_replace_source: "hashed_helper".to_string(),
                     resolved_module_id:"dummy".into()
                 },
                 "foo".to_string() => ResolvedReplaceInfo {
+                    chunk_id: Some("foo".into()),
                     to_replace_source: "foo".into(),
                     resolved_module_id: "foo".into()
                 }