refactor(ecma_preset_env): visit_mut for polyfills

crates/swc_ecma_preset_env/src/corejs2/entry.rs

@@ -1,9 +1,9 @@
 use indexmap::IndexSet;
 use preset_env_base::{version::should_enable, Versions};
 use swc_atoms::js_word;
-use swc_common::{util::move_map::MoveMap, DUMMY_SP};
+use swc_common::DUMMY_SP;
 use swc_ecma_ast::*;
-use swc_ecma_visit::{Fold, FoldWith};
+use swc_ecma_visit::{noop_visit_mut_type, VisitMut, VisitMutWith};
 
 use super::builtin::BUILTINS;
 
@@ -64,31 +64,21 @@ impl Entry {
     }
 }
 
-impl Fold for Entry {
-    fn fold_import_decl(&mut self, i: ImportDecl) -> ImportDecl {
-        let i: ImportDecl = i.fold_children_with(self);
+impl VisitMut for Entry {
+    noop_visit_mut_type!();
 
+    fn visit_mut_import_decl(&mut self, i: &mut ImportDecl) {
         let remove = i.specifiers.is_empty() && self.add_all(&i.src.value);
 
         if remove {
-            ImportDecl {
-                src: Str {
-                    span: DUMMY_SP,
-                    value: js_word!(""),
-                    ..*i.src
-                }
-                .into(),
-                ..i
-            }
-        } else {
-            i
+            i.src.value = js_word!("");
+            i.src.span = DUMMY_SP;
         }
     }
 
-    fn fold_module_items(&mut self, items: Vec<ModuleItem>) -> Vec<ModuleItem> {
-        items.move_flat_map(|item| {
-            let item: ModuleItem = item.fold_with(self);
-
+    fn visit_mut_module_items(&mut self, items: &mut Vec<ModuleItem>) {
+        items.retain_mut(|item| {
+            item.visit_mut_children_with(self);
             if let ModuleItem::Stmt(Stmt::Expr(ExprStmt { expr, .. })) = &item {
                 if let Expr::Call(CallExpr {
                     callee: Callee::Expr(callee),
@@ -115,13 +105,12 @@ impl Fold for Entry {
                             }
                             && self.add_all("@swc/polyfill")
                         {
-                            return None;
+                            return false;
                         }
                     }
                 }
             }
-
-            Some(item)
+            true
         })
     }
 }

crates/swc_ecma_preset_env/src/corejs3/entry.rs

@@ -7,7 +7,7 @@ use preset_env_base::{
 use swc_atoms::js_word;
 use swc_common::{collections::AHashMap, DUMMY_SP};
 use swc_ecma_ast::*;
-use swc_ecma_visit::{Fold, FoldWith};
+use swc_ecma_visit::VisitMut;
 
 use super::{compat::DATA as CORE_JS_COMPAT_DATA, data::MODULES_BY_VERSION};
 
@@ -91,24 +91,13 @@ impl Entry {
     }
 }
 
-impl Fold for Entry {
-    fn fold_import_decl(&mut self, i: ImportDecl) -> ImportDecl {
-        let i: ImportDecl = i.fold_children_with(self);
-
+impl VisitMut for Entry {
+    fn visit_mut_import_decl(&mut self, i: &mut ImportDecl) {
         let remove = i.specifiers.is_empty() && self.add(&i.src.value);
 
         if remove {
-            ImportDecl {
-                src: Str {
-                    span: DUMMY_SP,
-                    value: js_word!(""),
-                    ..*i.src
-                }
-                .into(),
-                ..i
-            }
-        } else {
-            i
+            i.src.span = DUMMY_SP;
+            i.src.value = js_word!("");
         }
     }
 }

crates/swc_ecma_preset_env/src/lib.rs

@@ -22,7 +22,7 @@ use swc_ecma_transforms::{
     Assumptions,
 };
 use swc_ecma_utils::prepend_stmts;
-use swc_ecma_visit::{Fold, FoldWith, VisitWith};
+use swc_ecma_visit::{as_folder, Fold, VisitMut, VisitMutWith, VisitWith};
 
 pub use self::transform_data::Feature;
 
@@ -311,7 +311,7 @@ where
 
     chain!(
         pass,
-        Polyfills {
+        as_folder(Polyfills {
             mode: c.mode,
             regenerator: should_enable!(Regenerator, true),
             corejs: c.core_js.unwrap_or(Version {
@@ -323,7 +323,7 @@ where
             targets,
             includes: included_modules,
             excludes: excluded_modules,
-        }
+        })
     )
 }
 
@@ -338,8 +338,8 @@ struct Polyfills {
     excludes: AHashSet<String>,
 }
 
-impl Fold for Polyfills {
-    fn fold_module(&mut self, mut m: Module) -> Module {
+impl VisitMut for Polyfills {
+    fn visit_mut_module(&mut self, m: &mut Module) {
         let span = m.span;
 
         let required = match self.mode {
@@ -365,7 +365,7 @@ impl Fold for Polyfills {
                     _ => unimplemented!("corejs version other than 2 / 3"),
                 };
 
-                if regenerator::is_required(&m) {
+                if regenerator::is_required(m) {
                     r.insert("regenerator-runtime/runtime.js");
                 }
 
@@ -374,13 +374,13 @@ impl Fold for Polyfills {
             Some(Mode::Entry) => match self.corejs {
                 Version { major: 2, .. } => {
                     let mut v = corejs2::Entry::new(self.targets, self.regenerator);
-                    m = m.fold_with(&mut v);
+                    m.visit_mut_with(&mut v);
                     v.imports
                 }
 
                 Version { major: 3, .. } => {
                     let mut v = corejs3::Entry::new(self.targets, self.corejs, !self.regenerator);
-                    m = m.fold_with(&mut v);
+                    m.visit_mut_with(&mut v);
                     v.imports
                 }
 
@@ -450,11 +450,9 @@ impl Fold for Polyfills {
         }
 
         m.body.retain(|item| !matches!(item, ModuleItem::ModuleDecl(ModuleDecl::Import(ImportDecl { src, .. })) if src.span == DUMMY_SP && src.value == js_word!("")));
-
-        m
     }
 
-    fn fold_script(&mut self, _: Script) -> Script {
+    fn visit_mut_script(&mut self, _: &mut Script) {
         unimplemented!("automatic polyfill for scripts")
     }
 }