wip: poc of dev tree-shaking

crates/mako/Cargo.toml

@@ -81,6 +81,7 @@ hyper-staticfile      = "0.9.6"
 hyper-tungstenite     = "0.10.0"
 indexmap              = "2.0.0"
 indicatif             = "0.17.8"
+log                   = "0.4.22"
 md5                   = "0.7.0"
 mdxjs                 = "0.2.6"
 mime_guess            = "2.0.4"

crates/mako/src/build.rs

@@ -88,16 +88,18 @@ impl Compiler {
             }
 
             let mut module_graph = self.context.module_graph.write().unwrap();
+            let mut module_registry = self.context.module_registry.write().unwrap();
 
             // handle current module
             let info = module.info.as_ref().unwrap();
             let resolved_deps = info.deps.resolved_deps.clone();
-            let m = module_graph.get_module_mut(&module.id);
+            let m = module_registry.get_module_mut(&module.id);
             if let Some(m) = m {
                 m.set_info(module.info);
             } else {
                 module_ids.insert(module.id.clone());
-                module_graph.add_module(module);
+                module_graph.add_module(&module);
+                module_registry.add_module(module);
             }
 
             // handle deps
@@ -132,7 +134,8 @@ impl Compiler {
                     // 拿到依赖之后需要直接添加 module 到 module_graph 里，不能等依赖 build 完再添加
                     // 是因为由于是异步处理各个模块，后者会导致大量重复任务的 build_module 任务（3 倍左右）
                     module_ids.insert(module.id.clone());
-                    module_graph.add_module(module);
+                    module_graph.add_module(&module);
+                    module_registry.add_module(module);
                 }
                 module_graph.add_dependency(&module_id, &dep_module_id, dep.dependency);
             }

crates/mako/src/compiler.rs

@@ -18,7 +18,7 @@
 use crate::config::{Config, Mode, ModuleIdStrategy, OutputMode};
 use crate::generate::chunk_graph::ChunkGraph;
 use crate::generate::optimize_chunk::OptimizeChunksInfo;
-use crate::module_graph::ModuleGraph;
+use crate::module_graph::{ModuleGraph, ModuleRegistry};
 use crate::plugin::{Plugin, PluginDriver, PluginGenerateEndParams};
 use crate::plugins;
 use crate::resolve::{get_resolvers, Resolvers};
@@ -29,7 +29,9 @@
 
 pub struct Context {
     pub module_graph: RwLock<ModuleGraph>,
+    pub optimized_module_graph: RwLock<ModuleGraph>,
     pub chunk_graph: RwLock<ChunkGraph>,
+    pub module_registry: RwLock<ModuleRegistry>,
     pub assets_info: Mutex<HashMap<String, String>>,
     pub modules_with_missing_deps: RwLock<Vec<String>>,
     pub config: Config,
@@ -123,6 +125,8 @@
             args: Args { watch: false },
             root: PathBuf::from(""),
             module_graph: RwLock::new(ModuleGraph::new()),
+            module_registry: Default::default(),
+            optimized_module_graph: RwLock::new(ModuleGraph::new()),
             chunk_graph: RwLock::new(ChunkGraph::new()),
             assets_info: Mutex::new(HashMap::new()),
             modules_with_missing_deps: RwLock::new(Vec::new()),
@@ -366,6 +370,8 @@
                 args,
                 root,
                 module_graph: RwLock::new(ModuleGraph::new()),
+                optimized_module_graph: RwLock::new(ModuleGraph::new()),
+                module_registry: Default::default(),
                 chunk_graph: RwLock::new(ChunkGraph::new()),
                 assets_info: Mutex::new(HashMap::new()),
                 modules_with_missing_deps: RwLock::new(Vec::new()),
@@ -432,11 +438,11 @@
             let module_graph = self.context.module_graph.read().unwrap();
             assign_numeric_ids(
                 module_graph.modules(),
-                |a, b| compare_modules_by_incoming_edges(&module_graph, &a.id, &b.id),
+                |a, b| compare_modules_by_incoming_edges(&module_graph, a, b),
                 |module, id| {
                     let mut numeric_ids_map = self.context.numeric_ids_map.write().unwrap();
                     // reserved ten indexes for swc helper and others runtime module
-                    numeric_ids_map.insert(module.id.id.clone(), id + 10);
+                    numeric_ids_map.insert(module.id.clone(), id + 10);
                 },
             )
         }
@@ -482,7 +488,8 @@
         crate::mako_profile_function!();
         let cg = self.context.chunk_graph.read().unwrap();
         let mg = self.context.module_graph.read().unwrap();
-        cg.full_hash(&mg)
+        let mr = self.context.module_registry.read().unwrap();
+        cg.full_hash(&mg, &mr)
     }
 
     fn clean_dist(&self) -> Result<()> {

crates/mako/src/dev/update.rs

@@ -128,10 +128,10 @@
             clear_resolver_cache(&self.context.resolvers);
             let mut modules_with_missing_deps =
                 self.context.modules_with_missing_deps.write().unwrap();
-            let mut module_graph = self.context.module_graph.write().unwrap();
+            let mut module_registry = self.context.module_registry.write().unwrap();
             for module_id in modules_with_missing_deps.clone().iter() {
                 let id = ModuleId::new(module_id.clone());
-                let module = module_graph.get_module_mut(&id).unwrap();
+                let module = module_registry.get_module_mut(&id).unwrap();
                 let missing_deps = module.info.clone().unwrap().deps.missing_deps;
                 for (_source, dep) in missing_deps {
                     let resolved =
@@ -365,6 +365,8 @@
         let mut dep_changed_module_ids = HashSet::new();
 
         let mut module_graph = self.context.module_graph.write().unwrap();
+
+        let mut module_registry = self.context.module_registry.write().unwrap();
         for (modified_module, diff, mut dependence_modules, dependencies) in modified_results {
             if diff.dependence_changed(&modified_module.id, &module_graph, &dependencies) {
                 dep_changed_module_ids.insert(modified_module.id.clone());
@@ -397,7 +399,8 @@
                     added.push(add_module_id.to_path());
                 }
 
-                module_graph.add_module(add_module);
+                module_graph.add_module(&add_module);
+                module_registry.add_module(add_module);
 
                 deps.iter().for_each(|&dep| {
                     module_graph.add_dependency(&modified_module.id, add_module_id, dep.clone());
@@ -421,7 +424,8 @@
             modified_module_ids.insert(modified_module.id.clone());
 
             // replace module
-            module_graph.replace_module(modified_module);
+            // module_graph.replace_module(modified_module);
+            module_registry.add_module(modified_module);
         }
 
         Result::Ok((modified_module_ids, dep_changed_module_ids, added))

crates/mako/src/dev/watch.rs

@@ -53,7 +53,7 @@
 
         self.watch_dir_recursive(self.root.into(), &self.get_ignore_list(true))?;
 
-        let module_graph = self.compiler.context.module_graph.read().unwrap();
+        let module_graph = self.compiler.context.module_registry.read().unwrap();
         let mut dirs = HashSet::new();
         module_graph.modules().iter().for_each(|module| {
             if let Some(ResolverResource::Resolved(resource)) = module

crates/mako/src/generate.rs

@@ -84,12 +84,17 @@ impl Compiler {
 
         let async_dep_map = self.mark_async();
 
+        {
+            let m = self.context.module_graph.read().unwrap();
+            *self.context.optimized_module_graph.write().unwrap() = m.clone();
+        };
+
         // Disable tree shaking in watch mode temporarily
         // ref: https://github.com/umijs/mako/issues/396
         if !self.context.args.watch {
             match self.context.config._tree_shaking {
                 Some(TreeShakingStrategy::Basic) => {
-                    let mut module_graph = self.context.module_graph.write().unwrap();
+                    let mut module_graph = self.context.optimized_module_graph.write().unwrap();
 
                     crate::mako_profile_scope!("tree shake");
                     self.context

crates/mako/src/generate/chunk.rs

@@ -7,7 +7,7 @@ use twox_hash::XxHash64;
 
 use crate::ast::file::parse_path;
 use crate::module::ModuleId;
-use crate::module_graph::ModuleGraph;
+use crate::module_graph::{ModuleGraph, ModuleRegistry};
 use crate::utils::url_safe_base64_encode;
 
 // TODO: Refact ChunkId
@@ -133,14 +133,14 @@ impl Chunk {
         self.modules.contains(module_id)
     }
 
-    pub fn hash(&self, mg: &ModuleGraph) -> u64 {
+    pub fn hash(&self, _mg: &ModuleGraph, registry: &ModuleRegistry) -> u64 {
         let mut sorted_module_ids = self.modules.iter().cloned().collect::<Vec<ModuleId>>();
         sorted_module_ids.sort_by_key(|m| m.id.clone());
 
         let mut hash: XxHash64 = Default::default();
 
         for id in sorted_module_ids {
-            let m = mg.get_module(&id).unwrap();
+            let m = registry.module(&id).unwrap();
 
             if let Some(info) = &m.info {
                 hash.write_u64(info.raw_hash);

crates/mako/src/generate/chunk_graph.rs

@@ -9,7 +9,7 @@ use twox_hash::XxHash64;
 
 use crate::generate::chunk::{Chunk, ChunkId, ChunkType};
 use crate::module::ModuleId;
-use crate::module_graph::ModuleGraph;
+use crate::module_graph::{ModuleGraph, ModuleRegistry};
 
 pub struct ChunkGraph {
     pub(crate) graph: StableDiGraph<Chunk, ()>,
@@ -104,13 +104,13 @@ impl ChunkGraph {
         self.graph.node_weights().map(|c| c.filename()).collect()
     }
 
-    pub fn full_hash(&self, module_graph: &ModuleGraph) -> u64 {
+    pub fn full_hash(&self, module_graph: &ModuleGraph, module_registry: &ModuleRegistry) -> u64 {
         let mut chunks = self.get_all_chunks();
         chunks.sort_by_key(|c| c.id.id.clone());
 
         let mut hasher: XxHash64 = Default::default();
         for c in chunks {
-            hasher.write_u64(c.hash(module_graph))
+            hasher.write_u64(c.hash(module_graph, module_registry))
         }
         hasher.finish()
     }

crates/mako/src/generate/chunk_pot.rs

@@ -17,7 +17,7 @@ pub use crate::generate::chunk_pot::util::CHUNK_FILE_NAME_HASH_LENGTH;
 use crate::generate::chunk_pot::util::{hash_hashmap, hash_vec};
 use crate::generate::generate_chunks::ChunkFile;
 use crate::module::{Module, ModuleAst, ModuleId};
-use crate::module_graph::ModuleGraph;
+use crate::module_graph::ModuleRegistry;
 use crate::ternary;
 
 pub struct ChunkPot<'a> {
@@ -31,7 +31,7 @@ pub struct ChunkPot<'a> {
 impl<'cp> ChunkPot<'cp> {
     pub fn from<'a: 'cp>(
         chunk: &'a Chunk,
-        mg: &'a ModuleGraph,
+        mg: &'a ModuleRegistry,
         context: &'cp Arc<Context>,
     ) -> Self {
         let (js_modules, stylesheet) = ChunkPot::split_modules(chunk.get_modules(), mg, context);
@@ -143,7 +143,7 @@ impl<'cp> ChunkPot<'cp> {
 
     fn split_modules<'a>(
         module_ids: &LinkedHashSet<ModuleId>,
-        module_graph: &'a ModuleGraph,
+        registry: &'a ModuleRegistry,
         context: &'a Arc<Context>,
     ) -> (JsModules<'a>, Option<CssModules<'a>>) {
         crate::mako_profile_function!(module_ids.len().to_string());
@@ -154,7 +154,7 @@ impl<'cp> ChunkPot<'cp> {
         let mut css_raw_hashes = vec![];
 
         module_ids.iter().for_each(|module_id| {
-            let module = module_graph.get_module(module_id).unwrap();
+            let module = registry.module(module_id).unwrap();
 
             if module.info.is_none() {
                 return;

crates/mako/src/generate/chunk_pot/ast_impl.rs

@@ -84,14 +84,13 @@
         None => None,
         _ => {
             mako_profile_scope!("build_source_map");
-
-            let module_graph = context.module_graph.read().unwrap();
+            let registry = context.module_registry.read().unwrap();
             let chunk_source_map = build_source_map(&source_map, cm);
 
             let mut chain_map = HashMap::<String, Vec<swc_sourcemap::SourceMap>>::new();
 
             chunk.get_modules().iter().for_each(|module_id| {
-                if let Some(module) = module_graph.get_module(module_id) {
+                if let Some(module) = registry.module(module_id) {
                     if let Some(info) = module.info.as_ref()
                         && matches!(info.ast, crate::module::ModuleAst::Css(_))
                     {

crates/mako/src/generate/generate_chunks.rs

@@ -19,6 +19,7 @@
 use crate::generate::chunk_pot::{get_css_chunk_filename, ChunkPot, CHUNK_FILE_NAME_HASH_LENGTH};
 use crate::generate::transform::transform_css_generate;
 use crate::module::{ModuleAst, ModuleId};
+use crate::module_graph::ModuleRegistry;
 use crate::utils::thread_pool;
 
 #[derive(Clone)]
@@ -141,7 +142,7 @@
             .par_iter()
             .map(|chunk| {
                 let context = self.context.clone();
-                let module_graph = context.module_graph.read().unwrap();
+                let module_registry = context.module_registry.read().unwrap();
                 let chunk_graph = self.context.chunk_graph.read().unwrap();
 
                 let (js_chunks_hash_placeholder, css_chunks_hash_placeholder) = chunk_graph
@@ -154,7 +155,7 @@
                             // TODO: maybe we can split chunks to chunk pots before generate, because normal chunks will be
                             // split here and fn generate_normal_chunk_files twice
                             let chunk_pot =
-                                ChunkPot::from(descendant_chunk, &module_graph, &context);
+                                ChunkPot::from(descendant_chunk, &module_registry, &context);
 
                             if self.context.config.hash {
                                 let placeholder = nanoid!(CHUNK_FILE_NAME_HASH_LENGTH);
@@ -188,7 +189,7 @@
                     );
 
                 let chunk_files = {
-                    let chunk_pot = ChunkPot::from(chunk, &module_graph, &context);
+                    let chunk_pot = ChunkPot::from(chunk, &module_registry, &context);
                     chunk_pot
                         .to_entry_chunk_files(
                             &context,
@@ -239,10 +240,10 @@
                 let context = self.context.clone();
                 let chunk_id = chunk.id.clone();
                 let chunk_graph = context.chunk_graph.read().unwrap();
-                let module_graph = context.module_graph.read().unwrap();
+                let module_registry = context.module_registry.read().unwrap();
                 let chunk = chunk_graph.chunk(&chunk_id).unwrap();
 
-                let chunk_files = ChunkPot::from(chunk, &module_graph, &context)
+                let chunk_files = ChunkPot::from(chunk, &module_registry, &context)
                     .to_normal_chunk_files(chunk, &context);
 
                 chunk_files
@@ -335,7 +336,7 @@
 
 pub fn modules_to_js_stmts(
     module_ids: &IndexSet<ModuleId>,
-    module_graph: &std::sync::RwLockReadGuard<crate::module_graph::ModuleGraph>,
+    module_registry: &ModuleRegistry,
     context: &Arc<Context>,
 ) -> Result<(Vec<PropOrSpread>, Option<Stylesheet>)> {
     let mut js_stmts = vec![];
@@ -344,7 +345,7 @@
     let module_ids: Vec<_> = module_ids.iter().collect();
 
     for module_id in module_ids {
-        let module = module_graph.get_module(module_id).unwrap();
+        let module = module_registry.get_module(module_id).unwrap();
         let ast = module.info.as_ref().unwrap();
         let ast = &ast.ast;
 

crates/mako/src/generate/group_chunk.rs

@@ -20,7 +20,7 @@ impl Compiler {
 
         let mut visited = HashSet::new();
         let mut edges = vec![];
-        let module_graph = self.context.module_graph.read().unwrap();
+        let module_graph = self.context.optimized_module_graph.read().unwrap();
         let mut chunk_graph = self.context.chunk_graph.write().unwrap();
         chunk_graph.clear();
 
@@ -282,6 +282,7 @@ impl Compiler {
                 self.hot_update_module_chunks(first_modified_module, &mut chunk_graph);
 
             // collect added async chunks modules from module_graph
+            // FIXME: it should be optimized module graph
             let module_graph = self.context.module_graph.read().unwrap();
             let async_chunk_modules = update_result
                 .added

crates/mako/src/generate/hmr.rs

@@ -19,8 +19,9 @@
         module_ids: &IndexSet<ModuleId>,
         current_hash: u64,
     ) -> Result<(String, String)> {
-        let module_graph = &self.context.module_graph.read().unwrap();
-        let (js_stmts, _) = modules_to_js_stmts(module_ids, module_graph, &self.context).unwrap();
+        let module_registry = &self.context.module_registry.read().unwrap();
+        let (js_stmts, _) =
+            modules_to_js_stmts(module_ids, module_registry, &self.context).unwrap();
         let content = include_str!("../runtime/runtime_hmr.js").to_string();
 
         let runtime_code_snippets = [