Add shake exports transform to next-swc (#32253)

* Add shake exports transform to next-swc

* Less cloning

* Use JsWord instead of String

packages/next-swc/crates/core/src/lib.rs

@@ -54,6 +54,7 @@ pub mod next_ssg;
 pub mod page_config;
 pub mod react_remove_properties;
 pub mod remove_console;
+pub mod shake_exports;
 pub mod styled_jsx;
 mod top_level_binding_collector;
 
@@ -86,6 +87,9 @@ pub struct TransformOptions {
 
     #[serde(default)]
     pub react_remove_properties: Option<react_remove_properties::Config>,
+
+    #[serde(default)]
+    pub shake_exports: Option<shake_exports::Config>,
 }
 
 pub fn custom_before_pass(file: Arc<SourceFile>, opts: &TransformOptions) -> impl Fold {
@@ -124,6 +128,10 @@ pub fn custom_before_pass(file: Arc<SourceFile>, opts: &TransformOptions) -> imp
                 Either::Left(react_remove_properties::remove_properties(config.clone())),
             _ => Either::Right(noop()),
         },
+        match &opts.shake_exports {
+            Some(config) => Either::Left(shake_exports::shake_exports(config.clone())),
+            None => Either::Right(noop()),
+        }
     )
 }
 

packages/next-swc/crates/core/src/shake_exports.rs

@@ -0,0 +1,116 @@
+use serde::Deserialize;
+use swc_atoms::js_word;
+use swc_atoms::JsWord;
+use swc_ecmascript::ast::*;
+use swc_ecmascript::transforms::optimization::simplify::dce::{dce, Config as DCEConfig};
+use swc_ecmascript::visit::{Fold, FoldWith};
+
+#[derive(Clone, Debug, Deserialize)]
+pub struct Config {
+    pub ignore: Vec<JsWord>,
+}
+
+pub fn shake_exports(config: Config) -> impl Fold {
+    ExportShaker {
+        ignore: config.ignore,
+        ..Default::default()
+    }
+}
+
+#[derive(Debug, Default)]
+struct ExportShaker {
+    ignore: Vec<JsWord>,
+    remove_export: bool,
+}
+
+impl Fold for ExportShaker {
+    fn fold_module(&mut self, module: Module) -> Module {
+        let module = module.fold_children_with(self);
+        let module = module.fold_with(&mut dce(DCEConfig::default()));
+
+        module
+    }
+
+    fn fold_module_items(&mut self, items: Vec<ModuleItem>) -> Vec<ModuleItem> {
+        let mut new_items = vec![];
+        for item in items {
+            let item = item.fold_children_with(self);
+            if !self.remove_export {
+                new_items.push(item)
+            }
+            self.remove_export = false;
+        }
+        new_items
+    }
+
+    fn fold_export_decl(&mut self, mut decl: ExportDecl) -> ExportDecl {
+        match &mut decl.decl {
+            Decl::Fn(fn_decl) => {
+                if !self.ignore.contains(&fn_decl.ident.sym) {
+                    self.remove_export = true;
+                }
+            }
+            Decl::Class(class_decl) => {
+                if !self.ignore.contains(&class_decl.ident.sym) {
+                    self.remove_export = true;
+                }
+            }
+            Decl::Var(var_decl) => {
+                var_decl.decls = var_decl
+                    .decls
+                    .iter()
+                    .filter_map(|var_decl| {
+                        if let Pat::Ident(BindingIdent { id, .. }) = &var_decl.name {
+                            if self.ignore.contains(&id.sym) {
+                                return Some(var_decl.to_owned());
+                            }
+                        }
+                        None
+                    })
+                    .collect();
+                if var_decl.decls.is_empty() {
+                    self.remove_export = true;
+                }
+            }
+            _ => {}
+        }
+        decl
+    }
+
+    fn fold_named_export(&mut self, mut export: NamedExport) -> NamedExport {
+        export.specifiers = export
+            .specifiers
+            .into_iter()
+            .filter_map(|spec| {
+                if let ExportSpecifier::Named(named_spec) = spec {
+                    if let Some(ident) = &named_spec.exported {
+                        if self.ignore.contains(&ident.sym) {
+                            return Some(ExportSpecifier::Named(named_spec));
+                        }
+                    } else if self.ignore.contains(&named_spec.orig.sym) {
+                        return Some(ExportSpecifier::Named(named_spec));
+                    }
+                }
+                None
+            })
+            .collect();
+        if export.specifiers.is_empty() {
+            self.remove_export = true
+        }
+        export
+    }
+
+    fn fold_export_default_decl(&mut self, decl: ExportDefaultDecl) -> ExportDefaultDecl {
+        if !self.ignore.contains(&js_word!("default")) {
+            self.remove_export = true
+        }
+        decl
+    }
+
+    fn fold_export_default_expr(&mut self, expr: ExportDefaultExpr) -> ExportDefaultExpr {
+        if !self.ignore.contains(&js_word!("default")) {
+            self.remove_export = true
+        }
+        expr
+    }
+}

packages/next-swc/crates/core/tests/fixture.rs

@@ -1,7 +1,12 @@
 use next_swc::{
-    amp_attributes::amp_attributes, next_dynamic::next_dynamic, next_ssg::next_ssg,
-    page_config::page_config_test, react_remove_properties::remove_properties,
-    remove_console::remove_console, styled_jsx::styled_jsx,
+    amp_attributes::amp_attributes,
+    next_dynamic::next_dynamic,
+    next_ssg::next_ssg,
+    page_config::page_config_test,
+    react_remove_properties::remove_properties,
+    remove_console::remove_console,
+    shake_exports::{shake_exports, Config as ShakeExportsConfig},
+    styled_jsx::styled_jsx,
 };
 use std::path::PathBuf;
 use swc_common::{chain, comments::SingleThreadedComments, FileName, Mark, Span, DUMMY_SP};
@@ -158,3 +163,39 @@ fn react_remove_properties_custom_fixture(input: PathBuf) {
         &output,
     );
 }
+
+#[fixture("tests/fixture/shake-exports/most-usecases/input.js")]
+fn shake_exports_fixture(input: PathBuf) {
+    let output = input.parent().unwrap().join("output.js");
+    test_fixture(
+        syntax(),
+        &|_tr| {
+            shake_exports(ShakeExportsConfig {
+                ignore: vec![
+                    String::from("keep").into(),
+                    String::from("keep1").into(),
+                    String::from("keep2").into(),
+                    String::from("keep3").into(),
+                    String::from("keep4").into(),
+                ],
+            })
+        },
+        &input,
+        &output,
+    );
+}
+
+#[fixture("tests/fixture/shake-exports/keep-default/input.js")]
+fn shake_exports_fixture_default(input: PathBuf) {
+    let output = input.parent().unwrap().join("output.js");
+    test_fixture(
+        syntax(),
+        &|_tr| {
+            shake_exports(ShakeExportsConfig {
+                ignore: vec![String::from("default").into()],
+            })
+        },
+        &input,
+        &output,
+    );
+}

packages/next-swc/crates/core/tests/fixture/shake-exports/keep-default/input.js

@@ -0,0 +1,9 @@
+let shouldBeRemoved = 'should be removed'
+export function removeFunction() {
+  console.log(shouldBeRemoved);
+}
+
+let shouldBeKept = 'should be kept'
+export default function shouldBeKept() {
+  console.log(shouldBeKept)
+}

packages/next-swc/crates/core/tests/fixture/shake-exports/keep-default/output.js

@@ -0,0 +1,4 @@
+let shouldBeKept = "should be kept";
+export default function shouldBeKept() {
+    console.log(shouldBeKept);
+}

packages/next-swc/crates/core/tests/fixture/shake-exports/most-usecases/input.js

@@ -0,0 +1,36 @@
+let shouldBeKept = 'should be kept'
+export async function keep() {
+  console.log(shouldBeKept)
+}
+
+let shouldBeRemoved = 'should be removed'
+export function removeFunction() {
+  console.log(shouldBeRemoved);
+}
+
+export let removeVarDeclaration = 'should be removed'
+export let removeVarDeclarationUndefined // should also be removed
+export let multipleDecl = 'should be removed', keep1 = 'should be kept'
+export let keep2 = 'should be kept'
+
+export class RemoveClass {
+  remove() {
+    console.log('should be removed')
+  }
+}
+
+let x = 'x'
+let y = 'y'
+
+// This should be removed
+export {x, y as z}
+
+let keep3 = 'should be kept'
+let asKeep = 'should be kept'
+let removeNamed = 'should be removed'
+
+export {keep3, asKeep as keep4, removeNamed}
+
+export default function removeDefault() {
+  console.log('should be removed')
+}

packages/next-swc/crates/core/tests/fixture/shake-exports/most-usecases/output.js

@@ -0,0 +1,9 @@
+let shouldBeKept = "should be kept";
+export async function keep() {
+    console.log(shouldBeKept);
+}
+export let keep1 = "should be kept";
+export let keep2 = "should be kept";
+let keep3 = "should be kept";
+let asKeep = "should be kept";
+export { keep3, asKeep as keep4 };

packages/next-swc/crates/core/tests/full.rs

@@ -58,6 +58,7 @@ fn test(input: &Path, minify: bool) {
                 styled_components: Some(assert_json("{}")),
                 remove_console: None,
                 react_remove_properties: None,
+                shake_exports: None,
             };
 
             let options = options.patch(&fm);