[WIP] fix: treeshaking & skipmodule should jump the async module

.gitignore

@@ -5,6 +5,7 @@
 /examples/*/.umi-production
 /client/node_modules
 /e2e/*/*/dist
+/e2e/**/node_modules
 /e2e/fixtures.umi/*/node_modules/.cache
 /e2e/fixtures.umi/stable-hash/node_modules
 /e2e/fixtures.umi/*/.umi

crates/mako/src/compiler.rs

@@ -18,6 +18,7 @@ use crate::ast::file::win_path;
 use crate::config::{Config, ModuleIdStrategy, OutputMode};
 use crate::generate::chunk_graph::ChunkGraph;
 use crate::generate::optimize_chunk::OptimizeChunksInfo;
+use crate::module::{Dependency, ModuleId};
 use crate::module_graph::ModuleGraph;
 use crate::plugin::{Plugin, PluginDriver, PluginGenerateEndParams};
 use crate::plugins;
@@ -32,6 +33,7 @@ pub struct Context {
     pub chunk_graph: RwLock<ChunkGraph>,
     pub assets_info: Mutex<HashMap<String, String>>,
     pub modules_with_missing_deps: RwLock<Vec<String>>,
+    pub async_deps_map: RwLock<HashMap<ModuleId, Vec<Dependency>>>,
     pub config: Config,
     pub numeric_ids_map: RwLock<HashMap<String, usize>>,
     pub args: Args,
@@ -125,6 +127,7 @@ impl Default for Context {
             module_graph: RwLock::new(ModuleGraph::new()),
             chunk_graph: RwLock::new(ChunkGraph::new()),
             assets_info: Mutex::new(HashMap::new()),
+            async_deps_map: RwLock::new(HashMap::new()),
             modules_with_missing_deps: RwLock::new(Vec::new()),
             meta: Meta::new(),
             plugin_driver: Default::default(),
@@ -360,6 +363,7 @@ impl Compiler {
                 module_graph: RwLock::new(ModuleGraph::new()),
                 chunk_graph: RwLock::new(ChunkGraph::new()),
                 assets_info: Mutex::new(HashMap::new()),
+                async_deps_map: RwLock::new(HashMap::new()),
                 modules_with_missing_deps: RwLock::new(Vec::new()),
                 meta: Meta::new(),
                 plugin_driver,

crates/mako/src/generate.rs

@@ -28,11 +28,10 @@ use crate::compiler::{Compiler, Context};
 use crate::config::{DevtoolConfig, OutputMode, TreeShakingStrategy};
 use crate::dev::update::UpdateResult;
 use crate::generate::generate_chunks::{ChunkFile, ChunkFileType};
-use crate::module::{Dependency, ModuleId};
+use crate::module::ModuleId;
 use crate::plugins::bundless_compiler::BundlessCompiler;
 use crate::stats::StatsJsonMap;
 use crate::utils::base64_encode;
-use crate::visitors::async_module::mark_async;
 
 #[derive(Serialize)]
 struct ChunksUrlMap {
@@ -53,18 +52,6 @@ impl Compiler {
         Ok(stats)
     }
 
-    fn mark_async(&self) -> HashMap<ModuleId, Vec<Dependency>> {
-        let module_ids = {
-            let module_graph = self.context.module_graph.read().unwrap();
-            let (mut module_ids, _) = module_graph.toposort();
-            // start from the leaf nodes, so reverser the sort
-            module_ids.reverse();
-            drop(module_graph);
-            module_ids
-        };
-        mark_async(&module_ids, &self.context)
-    }
-
     pub fn generate(&self) -> Result<StatsJsonMap> {
         debug!("generate");
         let t_generate = Instant::now();
@@ -82,8 +69,6 @@ impl Compiler {
         debug!("tree_shaking");
         let t_tree_shaking = Instant::now();
 
-        let async_dep_map = self.mark_async();
-
         // Disable tree shaking in watch mode temporarily
         // ref: https://github.com/umijs/mako/issues/396
         if !self.context.args.watch {
@@ -139,7 +124,8 @@ impl Compiler {
         // 因为放 chunks 的循环里，一个 module 可能存在于多个 chunk 里，可能会被编译多遍
         let t_transform_modules = Instant::now();
         debug!("transform all modules");
-        self.transform_all(async_dep_map)?;
+        self.transform_all()?;
+
         let t_transform_modules = t_transform_modules.elapsed();
 
         // ensure output dir exists

crates/mako/src/generate/transform.rs

@@ -1,4 +1,5 @@
 use std::collections::HashMap;
+use std::ops::DerefMut;
 use std::sync::mpsc::channel;
 use std::sync::Arc;
 use std::time::Instant;
@@ -34,7 +35,7 @@
 use crate::visitors::optimize_define_utils::OptimizeDefineUtils;
 
 impl Compiler {
-    pub fn transform_all(&self, async_deps_map: HashMap<ModuleId, Vec<Dependency>>) -> Result<()> {
+    pub fn transform_all(&self) -> Result<()> {
         let t = Instant::now();
         let context = &self.context;
         let module_ids = {
@@ -47,39 +48,40 @@
                 .collect::<Vec<_>>()
         };
 
-        transform_modules_in_thread(&module_ids, context, async_deps_map)?;
+        transform_modules_in_thread(&module_ids, context)?;
         debug!(">> transform modules in {}ms", t.elapsed().as_millis());
         Ok(())
     }
 }
 
 pub fn transform_modules(module_ids: Vec<ModuleId>, context: &Arc<Context>) -> Result<()> {
     let t = Instant::now();
-    let async_deps_by_module_id = mark_async(&module_ids, context);
+    let mut module_graph = context.module_graph.write().unwrap();
+    mark_async(module_graph.deref_mut(), &module_ids, context);
     debug!(">> mark async in {}ms", t.elapsed().as_millis());
     let t = Instant::now();
-    transform_modules_in_thread(&module_ids, context, async_deps_by_module_id)?;
+    transform_modules_in_thread(&module_ids, context)?;
     debug!(">> transform modules in {}ms", t.elapsed().as_millis());
     Ok(())
 }
 
 pub fn transform_modules_in_thread(
     module_ids: &Vec<ModuleId>,
     context: &Arc<Context>,
-    async_deps_by_module_id: HashMap<ModuleId, Vec<Dependency>>,
 ) -> Result<()> {
     crate::mako_profile_function!();
 
     let (rs, rr) = channel::<Result<(ModuleId, ModuleAst)>>();
-
     for module_id in module_ids {
+        let async_deps_by_module_id = context.async_deps_map.read().unwrap();
         let context = context.clone();
         let rs = rs.clone();
         let module_id = module_id.clone();
         let async_deps = async_deps_by_module_id
             .get(&module_id)
-            .expect(&module_id.id)
-            .clone();
+            .cloned()
+            .unwrap_or(vec![]);
+
         thread_pool::spawn(move || {
             let module_graph = context.module_graph.read().unwrap();
             let deps = module_graph.get_dependencies(&module_id);

crates/mako/src/module.rs

@@ -491,7 +491,13 @@
 
 impl Debug for Module {
     fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
-        write!(f, "id={}({:?})", self.id.id, self.get_module_type())
+        write!(
+            f,
+            "id={}({:?}) is_async={:?}",
+            self.id.id,
+            self.get_module_type(),
+            self.info.as_ref().unwrap().is_async
+        )
     }
 }
 

crates/mako/src/plugins/tree_shaking/shake.rs

@@ -21,25 +21,33 @@
 use crate::plugins::tree_shaking::shake::module_concatenate::optimize_module_graph;
 use crate::plugins::tree_shaking::statement_graph::{ExportInfo, ExportSpecifierInfo, ImportInfo};
 use crate::plugins::tree_shaking::{module, remove_useless_stmts, statement_graph};
+use crate::visitors::async_module::mark_async;
 use crate::{mako_profile_function, mako_profile_scope};
 
 type TreeShakingModuleMap = HashMap<ModuleId, RefCell<TreeShakeModule>>;
 
 pub fn optimize_modules(module_graph: &mut ModuleGraph, context: &Arc<Context>) -> Result<()> {
+    let module_ids = {
+        let (mut module_ids, _) = module_graph.toposort();
+        // start from the leaf nodes, so reverser the sort
+        module_ids.reverse();
+        module_ids
+    };
+    mark_async(module_graph, &module_ids, context);
     let (topo_sorted_modules, _cyclic_modules) = {
         mako_profile_scope!("tree shake topo-sort");
         module_graph.toposort()
     };
-
     let (_skipped, tree_shake_modules_ids): (Vec<ModuleId>, Vec<ModuleId>) =
         topo_sorted_modules.into_iter().partition(|module_id| {
             let module = module_graph.get_module(module_id).unwrap();
 
             let module_type = module.get_module_type();
 
             // skip non script modules and external modules
-            if module_type != ModuleType::Script || module.is_external() {
-                if module_type != ModuleType::Script && !module.is_external() {
+            let is_async = &module.info.as_ref().unwrap().is_async;
+            if module_type != ModuleType::Script || module.is_external() || *is_async {
+                if module_type != ModuleType::Script && !module.is_external() || *is_async {
                     // mark all non script modules' script dependencies as side_effects
                     for dep_id in module_graph.dependence_module_ids(module_id) {
                         let dep_module = module_graph.get_module_mut(&dep_id).unwrap();

crates/mako/src/plugins/tree_shaking/shake/skip_module.rs

@@ -493,7 +493,6 @@ pub(super) fn skip_module_optimize(
         if module_graph.has_module(module_id) {
             // stmt_id is reversed order
             for to_replace in replaces.iter() {
-                // println!("{} apply with {:?}", module_id.id, to_replace.1);
                 apply_replace(to_replace, module_id, module_graph, context);
             }
 

crates/mako/src/visitors/async_module.rs

@@ -1,4 +1,4 @@
-use std::collections::HashMap;
+use std::collections::{HashSet, VecDeque};
 use std::sync::Arc;
 
 use swc_core::common::util::take::Take;
@@ -14,6 +14,7 @@
 use crate::ast::DUMMY_CTXT;
 use crate::compiler::Context;
 use crate::module::{Dependency, ModuleId};
+use crate::module_graph::ModuleGraph;
 
 const ASYNC_IMPORTED_MODULE: &str = "_async__mako_imported_module_";
 
@@ -205,30 +206,61 @@
     }
 }
 
-pub fn mark_async(
-    module_ids: &[ModuleId],
-    context: &Arc<Context>,
-) -> HashMap<ModuleId, Vec<Dependency>> {
-    let mut async_deps_by_module_id = HashMap::new();
-    let mut module_graph = context.module_graph.write().unwrap();
-    // TODO: 考虑成环的场景
+pub fn mark_async(module_graph: &mut ModuleGraph, module_ids: &[ModuleId], context: &Arc<Context>) {
+    let mut to_visit_queue = module_graph
+        .modules()
+        .iter()
+        .filter_map(|m| {
+            m.info
+                .as_ref()
+                .and_then(|i| if i.is_async { Some(m.id.clone()) } else { None })
+        })
+        .collect::<VecDeque<_>>();
+    let mut visited = HashSet::new();
+
+    // polluted async to dependants
+    while let Some(module_id) = to_visit_queue.pop_front() {
+        if visited.contains(&module_id) {
+            continue;
+        }
+
+        module_graph
+            .get_dependents(&module_id)
+            .iter()
+            .filter_map(|(dependant, dependency)| {
+                if !dependency.resolve_type.is_sync_esm() {
+                    return None;
+                }
+                if !visited.contains(*dependant) {
+                    Some((*dependant).clone())
+                } else {
+                    None
+                }
+            })
+            .collect::<Vec<_>>()
+            .iter()
+            .for_each(|module_id| {
+                let m = module_graph.get_module_mut(module_id).unwrap();
+                m.info.as_mut().unwrap().is_async = true;
+
+                to_visit_queue.push_back(module_id.clone());
+            });
+
+        visited.insert(module_id.clone());
+    }
+
     module_ids.iter().for_each(|module_id| {
         let deps = module_graph.get_dependencies_info(module_id);
+        let mut async_deps_map = context.async_deps_map.write().unwrap();
         let async_deps: Vec<Dependency> = deps
             .into_iter()
             .filter(|(_, dep, is_async)| dep.resolve_type.is_sync_esm() && *is_async)
             .map(|(_, dep, _)| dep.clone())
             .collect();
-        let module = module_graph.get_module_mut(module_id).unwrap();
-        if let Some(info) = module.info.as_mut() {
-            // a module with async deps need to be polluted into async module
-            if !info.is_async && !async_deps.is_empty() {
-                info.is_async = true;
-            }
-            async_deps_by_module_id.insert(module_id.clone(), async_deps);
+        if !async_deps.is_empty() {
+            async_deps_map.insert(module_id.clone(), async_deps);
         }
     });
-    async_deps_by_module_id
 }
 
 #[cfg(test)]
@@ -255,8 +287,8 @@
         "#
         .trim());
         assert_eq!(
-            code,
-            r#"
+      code,
+      r#"
 __mako_require__._async(module, async (handleAsyncDeps, asyncResult)=>{
     "use strict";
     Object.defineProperty(exports, "__esModule", {
@@ -273,7 +305,7 @@
     asyncResult();
 }, true);
             "#.trim()
-        );
+    );
     }
 
     #[test]
@@ -286,8 +318,8 @@
         "#
         .trim());
         assert_eq!(
-            code,
-            r#"
+      code,
+      r#"
 __mako_require__._async(module, async (handleAsyncDeps, asyncResult)=>{
     "use strict";
     Object.defineProperty(exports, "__esModule", {
@@ -308,8 +340,8 @@
     asyncResult();
 }, true);
 "#
-                .trim()
-        );
+        .trim()
+    );
     }
 
     #[test]
@@ -321,8 +353,8 @@
                 "#
         .trim());
         assert_eq!(
-            code,
-            r#"
+      code,
+      r#"
 __mako_require__._async(module, async (handleAsyncDeps, asyncResult)=>{
     "use strict";
     Object.defineProperty(exports, "__esModule", {
@@ -340,8 +372,8 @@
     asyncResult();
 }, true);
 "#
-                .trim()
-        );
+        .trim()
+    );
     }
 
     #[test]
@@ -374,8 +406,8 @@
                 "#
         .trim());
         assert_eq!(
-            code,
-            r#"
+      code,
+      r#"
 __mako_require__._async(module, async (handleAsyncDeps, asyncResult)=>{
     "use strict";
     Object.defineProperty(exports, "__esModule", {
@@ -393,7 +425,7 @@
     asyncResult();
 }, true);
 "#.trim()
-        );
+    );
     }
 
     fn run(js_code: &str) -> String {