feat: move ensure runtime to entry

crates/mako/src/compiler.rs

@@ -298,6 +298,10 @@
             plugins.push(Arc::new(plugins::ssu::SUPlus::new()));
         }
 
+        if args.watch && config.experimental.central_ensure {
+            plugins.push(Arc::new(plugins::central_ensure::CentralChunkEnsure {}));
+        }
+
         if let Some(minifish_config) = &config._minifish {
             let inject = if let Some(inject) = &minifish_config.inject {
                 let mut map = HashMap::new();

crates/mako/src/config/experimental.rs

@@ -10,6 +10,7 @@ pub struct ExperimentalConfig {
     pub magic_comment: bool,
     #[serde(deserialize_with = "deserialize_detect_loop")]
     pub detect_circular_dependence: Option<DetectCircularDependence>,
+    pub central_ensure: bool,
 }
 
 #[derive(Deserialize, Serialize, Debug)]

crates/mako/src/config/mako.config.default.json

@@ -71,7 +71,8 @@
     "detectCircularDependence": {
       "ignores": ["node_modules"],
       "graphviz": false
-    }
+    },
+    "centralEnsure": true
   },
   "useDefineForClassFields": true,
   "emitDecoratorMetadata": false,

crates/mako/src/generate/chunk_graph.rs

@@ -63,6 +63,12 @@
         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/hmr.rs

@@ -10,6 +10,7 @@
 use crate::generate::chunk::Chunk;
 use crate::generate::generate_chunks::modules_to_js_stmts;
 use crate::module::ModuleId;
+use crate::plugins::central_ensure::module_ensure_map;
 
 impl Compiler {
     pub fn generate_hmr_chunk(
@@ -22,10 +23,22 @@
         let module_graph = &self.context.module_graph.read().unwrap();
         let (js_stmts, _) = modules_to_js_stmts(module_ids, module_graph, &self.context).unwrap();
         let content = include_str!("../runtime/runtime_hmr.js").to_string();
-        let content = content.replace("__CHUNK_ID__", &chunk.id.id).replace(
-            "__runtime_code__",
-            &format!("runtime._h='{}';", current_hash),
-        );
+
+        let mut runtime_code_snippets = vec![format!("runtime._h='{}';\n", current_hash)];
+
+        if self.context.config.experimental.central_ensure {
+            if let Ok(map) = module_ensure_map(&self.context) {
+                runtime_code_snippets.push(format!(
+                    "runtime.updateEnsure2Map({})",
+                    serde_json::to_string(&map)?
+                ));
+            }
+        }
+
+        let content = content
+            .replace("__CHUNK_ID__", &chunk.id.id)
+            .replace("__runtime_code__", &runtime_code_snippets.join("\n"));
+
         let mut js_ast = JsAst::build(filename, content.as_str(), self.context.clone())
             /* safe */
             .unwrap();

crates/mako/src/generate/transform.rs

@@ -27,7 +27,7 @@
 use crate::visitors::async_module::{mark_async, AsyncModule};
 use crate::visitors::common_js::common_js;
 use crate::visitors::css_imports::CSSImports;
-use crate::visitors::dep_replacer::{DepReplacer, DependenciesToReplace};
+use crate::visitors::dep_replacer::{DepReplacer, DependenciesToReplace, ResolvedReplaceInfo};
 use crate::visitors::dynamic_import::DynamicImport;
 use crate::visitors::mako_require::MakoRequire;
 use crate::visitors::meta_url_replacer::MetaUrlReplacer;
@@ -84,36 +84,45 @@
         thread_pool::spawn(move || {
             let module_graph = context.module_graph.read().unwrap();
             let deps = module_graph.get_dependencies(&module_id);
-            let mut resolved_deps: HashMap<String, (String, String)> = deps
+            let mut resolved_deps: HashMap<String, ResolvedReplaceInfo> = deps
                 .into_iter()
                 .map(|(id, dep)| {
-                    (
-                        dep.source.clone(),
-                        (
-                            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()
-                                }
-                                ResolveType::DynamicImport(import_options) => {
-                                    match import_options.get_chunk_name() {
-                                        Some(chunk_name) => {
-                                            generate_module_id(chunk_name, &context)
-                                        }
-                                        None => id.generate(&context),
-                                    }
-                                }
-                                _ => id.generate(&context),
-                            },
-                            id.id.clone(),
-                        ),
-                    )
+                    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();
+                            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(),
+                            }
+                        }
+                        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),
+                            });
+
+                            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(), replace_info)
                 })
                 .collect();
             insert_swc_helper_replace(&mut resolved_deps, &context);
@@ -162,10 +171,20 @@
     Ok(())
 }
 
-fn insert_swc_helper_replace(map: &mut HashMap<String, (String, String)>, context: &Arc<Context>) {
+fn insert_swc_helper_replace(
+    map: &mut HashMap<String, ResolvedReplaceInfo>,
+    context: &Arc<Context>,
+) {
     SWC_HELPERS.into_iter().for_each(|h| {
         let m_id: ModuleId = h.to_string().into();
-        map.insert(m_id.id.clone(), (m_id.generate(context), h.to_string()));
+        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.rs

@@ -1,5 +1,6 @@
 pub mod async_runtime;
 pub mod bundless_compiler;
+pub mod central_ensure;
 pub mod context_module;
 pub mod copy;
 pub mod detect_circular_dependence;

crates/mako/src/plugins/bundless_compiler.rs

@@ -24,7 +24,7 @@
 use crate::config::Config;
 use crate::module::{ModuleAst, ModuleId};
 use crate::plugin::{Plugin, PluginTransformJsParam};
-use crate::visitors::dep_replacer::{DepReplacer, DependenciesToReplace};
+use crate::visitors::dep_replacer::{DepReplacer, DependenciesToReplace, ResolvedReplaceInfo};
 use crate::visitors::dynamic_import::DynamicImport;
 
 pub struct BundlessCompiler {
@@ -80,11 +80,18 @@
                             }
                         };
 
-                        Ok((dep.source.clone(), (replacement.clone(), replacement)))
+                        Ok((
+                            dep.source.clone(),
+                            ResolvedReplaceInfo {
+                                chunk_id: None,
+                                to_replace_source: replacement,
+                                resolved_module_id: id.clone(),
+                            },
+                        ))
                     })
                     .collect::<Result<Vec<_>>>();
 
-                let resolved_deps: HashMap<String, (String, String)> =
+                let resolved_deps: HashMap<String, ResolvedReplaceInfo> =
                     resolved_deps?.into_iter().collect();
 
                 drop(module_graph);

crates/mako/src/plugins/central_ensure.rs

@@ -0,0 +1,79 @@
+use std::collections::BTreeMap;
+use std::sync::Arc;
+
+use anyhow::anyhow;
+
+use crate::compiler::Context;
+use crate::module::generate_module_id;
+use crate::plugin::Plugin;
+
+pub struct CentralChunkEnsure {}
+
+pub fn module_ensure_map(context: &Arc<Context>) -> anyhow::Result<BTreeMap<String, Vec<String>>> {
+    let mg = context
+        .module_graph
+        .read()
+        .map_err(|e| anyhow!("Read_Module_Graph_error:\n{:?}", e))?;
+    let cg = context
+        .chunk_graph
+        .read()
+        .map_err(|e| anyhow!("Read_Chunk_Graph_error:\n{:?}", e))?;
+
+    let mut chunk_async_map: BTreeMap<String, Vec<String>> = Default::default();
+
+    mg.modules().iter().for_each(|module| {
+        let be_dynamic_imported = mg
+            .get_dependents(&module.id)
+            .iter()
+            .any(|(_, dep)| dep.resolve_type.is_dynamic_esm());
+
+        if be_dynamic_imported {
+            cg.get_async_chunk_for_module(&module.id)
+                .iter()
+                .for_each(|chunk| {
+                    let deps_chunks = cg
+                        .installable_descendants_chunk(&chunk.id)
+                        .iter()
+                        .map(|chunk_id| chunk_id.generate(context))
+                        .collect::<Vec<_>>();
+
+                    chunk_async_map.insert(generate_module_id(&module.id.id, context), deps_chunks);
+                });
+        }
+    });
+
+    Ok(chunk_async_map)
+}
+
+impl Plugin for CentralChunkEnsure {
+    fn name(&self) -> &str {
+        "dev_ensure2"
+    }
+    fn runtime_plugins(&self, context: &Arc<Context>) -> anyhow::Result<Vec<String>> {
+        let chunk_async_map = module_ensure_map(context)?;
+
+        // TODO: compress the map to reduce duplicated chunk ids
+        let ensure_map = serde_json::to_string(&chunk_async_map)?;
+
+        let runtime = format!(
+            r#"
+(function(){{
+  let map = {ensure_map};
+  requireModule.updateEnsure2Map = function(newMapping) {{
+    map = newMapping;
+  }}
+  requireModule.ensure2 = function(chunkId){{
+    let toEnsure = map[chunkId];
+    if (toEnsure) {{
+      return Promise.all(toEnsure.map(function(c){{ return requireModule.ensure(c); }}))
+    }}else{{
+      return Promise.resolve();
+    }}
+  }};
+}})();
+"#
+        );
+
+        Ok(vec![runtime])
+    }
+}

crates/mako/src/visitors/dep_replacer.rs

@@ -23,14 +23,18 @@
     pub unresolved_mark: Mark,
 }
 
-type ResolvedModuleId = String;
-type ResolvedModulePath = String;
+#[derive(Debug, Clone)]
+pub struct ResolvedReplaceInfo {
+    pub chunk_id: Option<String>,
+    pub to_replace_source: String,
+    pub resolved_module_id: ModuleId,
+}
 
 #[derive(Debug, Clone)]
 pub struct DependenciesToReplace {
     // resolved stores the "source" maps to (generate_id, raw_module_id)
     // e.g. "react" => ("hashed_id", "/abs/to/react/index.js")
-    pub resolved: HashMap<String, (ResolvedModuleId, ResolvedModulePath)>,
+    pub resolved: HashMap<String, ResolvedReplaceInfo>,
     pub missing: HashMap<String, Dependency>,
 }
 
@@ -124,8 +128,9 @@
                     // css
                     // TODO: add testcases for this
                     let is_replaceable_css =
-                        if let Some((_, raw_id)) = self.to_replace.resolved.get(&source_string) {
-                            let (path, _search, query, _) = parse_path(raw_id).unwrap();
+                        if let Some(replace_info) = self.to_replace.resolved.get(&source_string) {
+                            let (path, _search, query, _) =
+                                parse_path(&replace_info.resolved_module_id.id).unwrap();
                             // when inline_css is enabled
                             // css is parsed as js modules
                             self.context.config.inline_css.is_none()
@@ -194,7 +199,7 @@
 impl DepReplacer<'_> {
     fn replace_source(&mut self, source: &mut Str) {
         if let Some(replacement) = self.to_replace.resolved.get(&source.value.to_string()) {
-            let module_id = replacement.0.clone();
+            let module_id = replacement.to_replace_source.clone();
             let span = source.span;
             *source = Str::from(module_id);
             source.span = span;
@@ -247,7 +252,7 @@
     use swc_core::common::GLOBALS;
     use swc_core::ecma::visit::VisitMutWith;
 
-    use super::{DepReplacer, DependenciesToReplace, ResolvedModuleId, ResolvedModulePath};
+    use super::{DepReplacer, DependenciesToReplace, ResolvedReplaceInfo};
     use crate::ast::tests::TestUtils;
     use crate::module::{Dependency, ImportType, ModuleId, ResolveType};
 
@@ -339,17 +344,14 @@
         );
     }
 
-    fn build_resolved(
-        key: &str,
-        module_id: &str,
-    ) -> HashMap<String, (ResolvedModuleId, ResolvedModulePath)> {
+    fn build_resolved(key: &str, module_id: &str) -> HashMap<String, ResolvedReplaceInfo> {
         hashmap! {
             key.to_string() =>
-            (
-
-             module_id.to_string(),
-             "".to_string()
-            )
+            ResolvedReplaceInfo {
+                chunk_id: None,
+                to_replace_source: module_id.into(),
+                resolved_module_id: "".into(),
+            }
         }
     }
 
@@ -367,7 +369,7 @@
 
     fn run(
         js_code: &str,
-        resolved: HashMap<String, (ResolvedModuleId, ResolvedModulePath)>,
+        resolved: HashMap<String, ResolvedReplaceInfo>,
         missing: HashMap<String, Dependency>,
     ) -> String {
         let mut test_utils = TestUtils::gen_js_ast(js_code);