feat(tree-shaking): optimize import namespace used all exports to par…

…tial used of source modules (#1584)

* test: ✅ add failure test case

* feat: ✨ add collect prop explicit field

* feat: ✨ collect explicit member prop access to optizmize used export in dep module
replace USE ALL export to partial used exports

* fix: dedestructuring before collect ExplicitProps

* test: ✅ add back trace env var

* test: ✅ update test case

* fix: 🐛 add globals and helpers

---------

Co-authored-by: Jinbao1001 <nodewebli@gmail.com>

.github/workflows/ci.yml

@@ -115,6 +115,8 @@ jobs:
       - name: LS
         run: ls -l ./packages/mako
       - name: Test E2E
+        env:
+          RUST_BACKTRACE: full
         run: pnpm ${{ matrix.script }}
 
   lint:

crates/mako/src/plugins/tree_shaking.rs

@@ -9,6 +9,7 @@ use crate::compiler::Context;
 use crate::module_graph::ModuleGraph;
 use crate::plugin::{Plugin, PluginTransformJsParam};
 
+mod collect_explicit_prop;
 mod module;
 mod module_side_effects_flag;
 mod remove_useless_stmts;

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

@@ -0,0 +1,222 @@
+use std::collections::{HashMap, HashSet};
+
+use swc_core::ecma::ast::{ComputedPropName, Id, Ident, Lit, MemberExpr, MemberProp};
+use swc_core::ecma::visit::{Visit, VisitWith};
+
+#[derive(Debug)]
+pub struct IdExplicitPropAccessCollector {
+    to_detected: HashSet<Id>,
+    accessed_by_explicit_prop_count: HashMap<Id, usize>,
+    ident_accessed_count: HashMap<Id, usize>,
+    accessed_by: HashMap<Id, HashSet<String>>,
+}
+
+impl IdExplicitPropAccessCollector {
+    pub(crate) fn new(ids: HashSet<Id>) -> Self {
+        Self {
+            to_detected: ids,
+            accessed_by_explicit_prop_count: Default::default(),
+            ident_accessed_count: Default::default(),
+            accessed_by: Default::default(),
+        }
+    }
+    pub(crate) fn explicit_accessed_props(mut self) -> HashMap<String, Vec<String>> {
+        self.to_detected
+            .iter()
+            .filter_map(|id| {
+                let member_prop_accessed = self.accessed_by_explicit_prop_count.get(id);
+                let ident_accessed = self.ident_accessed_count.get(id);
+
+                match (member_prop_accessed, ident_accessed) {
+                    // all ident are accessed explicitly, so there is member expr there is a name
+                    // ident, and at last plus the extra ident in import decl, that's 1 comes from.
+                    (Some(m), Some(i)) if (i - m) == 1 => {
+                        let mut accessed_by = Vec::from_iter(self.accessed_by.remove(id).unwrap());
+                        accessed_by.sort();
+
+                        let str_key = format!("{}#{}", id.0, id.1.as_u32());
+
+                        Some((str_key, accessed_by))
+                    }
+                    // Some un-explicitly access e.g: obj[foo]
+                    _ => None,
+                }
+            })
+            .collect()
+    }
+
+    fn increase_explicit_prop_accessed_count(&mut self, id: Id) {
+        self.accessed_by_explicit_prop_count
+            .entry(id.clone())
+            .and_modify(|c| {
+                *c += 1;
+            })
+            .or_insert(1);
+    }
+
+    fn insert_member_accessed_by(&mut self, id: Id, prop: &str) {
+        self.increase_explicit_prop_accessed_count(id.clone());
+        self.accessed_by
+            .entry(id)
+            .and_modify(|accessed| {
+                accessed.insert(prop.to_string());
+            })
+            .or_insert(HashSet::from([prop.to_string()]));
+    }
+}
+
+impl Visit for IdExplicitPropAccessCollector {
+    fn visit_ident(&mut self, n: &Ident) {
+        let id = n.to_id();
+
+        if self.to_detected.contains(&id) {
+            self.ident_accessed_count
+                .entry(id)
+                .and_modify(|c| {
+                    *c += 1;
+                })
+                .or_insert(1);
+        }
+    }
+
+    fn visit_member_expr(&mut self, n: &MemberExpr) {
+        if let Some(obj_ident) = n.obj.as_ident() {
+            let id = obj_ident.to_id();
+
+            if self.to_detected.contains(&id) {
+                match &n.prop {
+                    MemberProp::Ident(prop_ident) => {
+                        self.insert_member_accessed_by(id, prop_ident.sym.as_ref());
+                    }
+                    MemberProp::PrivateName(_) => {}
+                    MemberProp::Computed(ComputedPropName { expr, .. }) => {
+                        if let Some(lit) = expr.as_lit()
+                            && let Lit::Str(str) = lit
+                        {
+                            let visited_by = str.value.to_string();
+                            self.insert_member_accessed_by(id, &visited_by)
+                        }
+                    }
+                }
+            }
+        }
+
+        n.visit_children_with(self);
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use maplit::hashset;
+
+    use super::*;
+    use crate::ast::tests::TestUtils;
+
+    #[test]
+    fn test_no_prop() {
+        let fields = extract_explicit_fields(
+            r#"
+        import * as foo from "./foo.js";
+        console.log(foo)
+        "#,
+        );
+
+        assert_eq!(fields, None);
+    }
+    #[test]
+    fn test_no_access() {
+        let fields = extract_explicit_fields(
+            r#"
+        import * as foo from "./foo.js";
+        "#,
+        );
+
+        assert_eq!(fields, None);
+    }
+
+    #[test]
+    fn test_computed_prop() {
+        let fields = extract_explicit_fields(
+            r#"
+        import * as foo from "./foo.js";
+        foo['f' + 'o' + 'o']
+        "#,
+        );
+
+        assert_eq!(fields, None);
+    }
+
+    #[test]
+    fn test_simple_explicit_prop() {
+        let fields = extract_explicit_fields(
+            r#"
+        import * as foo from "./foo.js";
+        foo.x;
+        foo.y;
+        "#,
+        );
+
+        assert_eq!(fields.unwrap(), vec!["x".to_string(), "y".to_string()]);
+    }
+
+    #[test]
+    fn test_nest_prop_explicit_prop() {
+        let fields = extract_explicit_fields(
+            r#"
+        import * as foo from "./foo.js";
+        foo.x.z[foo.y]
+        "#,
+        );
+
+        assert_eq!(fields.unwrap(), vec!["x".to_string(), "y".to_string()]);
+    }
+
+    #[test]
+    fn test_string_literal_prop_explicit() {
+        let fields = extract_explicit_fields(
+            r#"
+        import * as foo from "./foo.js";
+        foo['x']
+        "#,
+        );
+
+        assert_eq!(fields.unwrap(), vec!["x".to_string()]);
+    }
+
+    #[test]
+    fn test_num_literal_prop_not_explicit() {
+        let fields = extract_explicit_fields(
+            r#"
+        import * as foo from "./foo.js";
+        foo[1]
+        "#,
+        );
+
+        assert_eq!(fields, None);
+    }
+
+    fn extract_explicit_fields(code: &str) -> Option<Vec<String>> {
+        let tu = TestUtils::gen_js_ast(code);
+
+        let id = namespace_id(&tu);
+        let str = format!("{}#{}", id.0, id.1.as_u32());
+
+        let mut v = IdExplicitPropAccessCollector::new(hashset! { id });
+        tu.ast.js().ast.visit_with(&mut v);
+
+        v.explicit_accessed_props().remove(&str)
+    }
+
+    fn namespace_id(tu: &TestUtils) -> Id {
+        tu.ast.js().ast.body[0]
+            .as_module_decl()
+            .unwrap()
+            .as_import()
+            .unwrap()
+            .specifiers[0]
+            .as_namespace()
+            .unwrap()
+            .local
+            .to_id()
+    }
+}

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

@@ -1,9 +1,15 @@
+use std::collections::HashSet;
+
+use swc_core::common::util::take::Take;
+use swc_core::common::SyntaxContext;
 use swc_core::ecma::ast::{
-    Decl, ExportDecl, ExportSpecifier, ImportDecl, ImportSpecifier, Module as SwcModule,
-    ModuleExportName,
+    Decl, ExportDecl, ExportSpecifier, Id, ImportDecl, ImportSpecifier, Module as SwcModule,
+    Module, ModuleExportName,
 };
-use swc_core::ecma::visit::{VisitMut, VisitMutWith, VisitWith};
+use swc_core::ecma::transforms::compat::es2015::destructuring;
+use swc_core::ecma::visit::{Fold, VisitMut, VisitMutWith, VisitWith};
 
+use super::collect_explicit_prop::IdExplicitPropAccessCollector;
 use crate::plugins::tree_shaking::module::TreeShakeModule;
 use crate::plugins::tree_shaking::statement_graph::analyze_imports_and_exports::{
     analyze_imports_and_exports, StatementInfo,
@@ -105,10 +111,10 @@ pub fn remove_useless_stmts(
 
     // remove from the end to the start
     stmts_to_remove.reverse();
-
     for stmt in stmts_to_remove {
         swc_module.body.remove(stmt);
     }
+    optimize_import_namespace(&mut used_import_infos, swc_module);
 
     (used_import_infos, used_export_from_infos)
 }
@@ -231,6 +237,68 @@ impl VisitMut for UselessExportStmtRemover {
     }
 }
 
+fn optimize_import_namespace(import_infos: &mut [ImportInfo], module: &mut Module) {
+    let namespaces = import_infos
+        .iter()
+        .filter_map(|import_info| {
+            let ns = import_info
+                .specifiers
+                .iter()
+                .filter_map(|sp| match sp {
+                    ImportSpecifierInfo::Namespace(ns) => Some(ns.clone()),
+                    _ => None,
+                })
+                .collect::<Vec<String>>();
+            if ns.is_empty() {
+                None
+            } else {
+                Some(ns)
+            }
+        })
+        .flatten()
+        .collect::<Vec<String>>();
+
+    let ids = namespaces
+        .iter()
+        .map(|ns| {
+            let (sym, ctxt) = ns.rsplit_once('#').unwrap();
+            (sym.into(), SyntaxContext::from_u32(ctxt.parse().unwrap()))
+        })
+        .collect::<HashSet<Id>>();
+
+    if !ids.is_empty() {
+        let mut v = IdExplicitPropAccessCollector::new(ids);
+        let destucturing_module = destructuring(Default::default()).fold_module(module.clone());
+        destucturing_module.visit_with(&mut v);
+        let explicit_prop_accessed_ids = v.explicit_accessed_props();
+
+        import_infos.iter_mut().for_each(|ii| {
+            ii.specifiers = ii
+                .specifiers
+                .take()
+                .into_iter()
+                .flat_map(|specifier_info| {
+                    if let ImportSpecifierInfo::Namespace(ref ns) = specifier_info {
+                        if let Some(visited_fields) = explicit_prop_accessed_ids.get(ns) {
+                            return visited_fields
+                                .iter()
+                                .map(|v| {
+                                    let imported_name = format!("{v}#0");
+                                    ImportSpecifierInfo::Named {
+                                        imported: Some(imported_name.clone()),
+                                        local: imported_name,
+                                    }
+                                })
+                                .collect::<Vec<_>>();
+                        }
+                    }
+                    vec![specifier_info]
+                })
+                .collect::<Vec<_>>();
+        })
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

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

@@ -11,6 +11,7 @@ use anyhow::Result;
 use rayon::prelude::*;
 use swc_core::common::util::take::Take;
 use swc_core::common::GLOBALS;
+use swc_core::ecma::transforms::base::helpers::{Helpers, HELPERS};
 
 use self::skip_module::skip_module_optimize;
 use crate::compiler::Context;
@@ -134,23 +135,24 @@ pub fn optimize_modules(module_graph: &mut ModuleGraph, context: &Arc<Context>)
     let mut current_index: usize = 0;
     let len = tree_shake_modules_ids.len();
 
-    {
+    GLOBALS.set(&context.meta.script.globals, || {
         mako_profile_scope!("tree-shake");
 
         while current_index < len {
             mako_profile_scope!(
                 "tree-shake-module",
                 &tree_shake_modules_ids[current_index].id
             );
-
-            current_index = shake_module(
-                module_graph,
-                &tree_shake_modules_ids,
-                &tree_shake_modules_map,
-                current_index,
-            );
+            HELPERS.set(&Helpers::new(true), || {
+                current_index = shake_module(
+                    module_graph,
+                    &tree_shake_modules_ids,
+                    &tree_shake_modules_map,
+                    current_index,
+                );
+            });
         }
-    }
+    });
 
     {
         mako_profile_scope!("update ast");

e2e/fixtures/tree-shaking.import_namespace/expect.js

@@ -0,0 +1,14 @@
+const assert = require("assert");
+const {
+  parseBuildResult,
+  injectSimpleJest,
+  moduleReg,
+} = require("../../../scripts/test-utils");
+const { files } = parseBuildResult(__dirname);
+
+injectSimpleJest();
+const content = files["index.js"];
+
+expect(content).toContain("shouldKeep1");
+expect(content).toContain("shouldKeep2");
+expect(content).not.toContain("shouldNotKeep");