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

@@ -22,10 +22,18 @@
         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 runtime_code_snippets = [
+            format!("runtime._h='{}';", current_hash),
+            self.context
+                .plugin_driver
+                .hmr_runtime_update_code(&self.context)?,
+        ];
+
+        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,38 +84,53 @@
         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
-                .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(),
-                        ),
-                    )
-                })
-                .collect();
+            let mut resolved_deps: HashMap<String, ResolvedReplaceInfo> = Default::default();
+
+            deps.into_iter().for_each(|(id, dep)| {
+                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(),
+                    },
+                };
+
+                resolved_deps
+                    .entry(dep.source.clone())
+                    .and_modify(|info: &mut ResolvedReplaceInfo| {
+                        match (&replace_info.chunk_id, &info.chunk_id) {
+                            (None, _) => {}
+                            (Some(id), _) => info.chunk_id = Some(id.clone()),
+                        }
+                    })
+                    .or_insert(replace_info);
+            });
             insert_swc_helper_replace(&mut resolved_deps, &context);
             let module = module_graph.get_module(&module_id).unwrap();
             let info = module.info.as_ref().unwrap();
@@ -162,10 +177,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/plugin.rs

@@ -159,6 +159,10 @@
         Ok(Vec::new())
     }
 
+    fn hmr_runtime_updates(&self, _context: &Arc<Context>) -> Result<Vec<String>> {
+        Ok(Vec::new())
+    }
+
     fn optimize_module_graph(
         &self,
         _module_graph: &mut ModuleGraph,
@@ -368,6 +372,14 @@
         Ok(plugins.join("\n"))
     }
 
+    pub fn hmr_runtime_update_code(&self, context: &Arc<Context>) -> Result<String> {
+        let mut plugins = Vec::new();
+        for plugin in &self.plugins {
+            plugins.extend(plugin.hmr_runtime_updates(context)?);
+        }
+        Ok(plugins.join("\n"))
+    }
+
     pub fn optimize_module_graph(
         &self,
         module_graph: &mut ModuleGraph,

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,90 @@
+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])
+    }
+
+    fn hmr_runtime_updates(&self, _context: &Arc<Context>) -> anyhow::Result<Vec<String>> {
+        let map = module_ensure_map(_context)?;
+
+        let update_mapping = format!(
+            "runtime.updateEnsure2Map({});",
+            serde_json::to_string(&map)?
+        );
+
+        Ok(vec![update_mapping])
+    }
+}