refactor(minifier): replace VisitMut with Traverse for inject and define plugins (#5705)

closes #5704

Author: Boshen

crates/oxc_minifier/src/plugins/inject_global_variables.rs

@@ -1,9 +1,12 @@
+use std::sync::Arc;
+
 use cow_utils::CowUtils;
+
 use oxc_allocator::Allocator;
-use oxc_ast::{ast::*, visit::walk_mut, AstBuilder, VisitMut};
+use oxc_ast::{ast::*, AstBuilder};
 use oxc_semantic::{ScopeTree, SymbolTable};
 use oxc_span::{CompactStr, SPAN};
-use std::sync::Arc;
+use oxc_traverse::{traverse_mut, Traverse, TraverseCtx};
 
 use super::replace_global_defines::{DotDefine, ReplaceGlobalDefines};
 
@@ -100,12 +103,18 @@ impl<'a> From<&InjectImport> for DotDefineState<'a> {
     }
 }
 
+#[must_use]
+pub struct InjectGlobalVariablesReturn {
+    pub symbols: SymbolTable,
+    pub scopes: ScopeTree,
+}
+
 /// Injects import statements for global variables.
 ///
 /// References:
 ///
 /// * <https://www.npmjs.com/package/@rollup/plugin-inject>
-pub struct InjectGlobalVariables<'a, 'b> {
+pub struct InjectGlobalVariables<'a> {
     ast: AstBuilder<'a>,
     config: InjectGlobalVariablesConfig,
 
@@ -116,36 +125,32 @@ pub struct InjectGlobalVariables<'a, 'b> {
     /// Identifiers for which dot define replaced a member expression.
     replaced_dot_defines:
         Vec<(/* identifier of member expression */ CompactStr, /* local */ CompactStr)>,
-
-    symbols: &'b mut SymbolTable, // will be used to keep symbols in sync
-    scopes: &'b mut ScopeTree,
 }
 
-impl<'a, 'b> VisitMut<'a> for InjectGlobalVariables<'a, 'b> {
-    fn visit_expression(&mut self, expr: &mut Expression<'a>) {
-        self.replace_dot_defines(expr);
-        walk_mut::walk_expression(self, expr);
+impl<'a> Traverse<'a> for InjectGlobalVariables<'a> {
+    fn enter_expression(&mut self, expr: &mut Expression<'a>, ctx: &mut TraverseCtx<'a>) {
+        self.replace_dot_defines(expr, ctx);
     }
 }
 
-impl<'a, 'b> InjectGlobalVariables<'a, 'b> {
-    pub fn new(
-        allocator: &'a Allocator,
-        symbols: &'b mut SymbolTable,
-        scopes: &'b mut ScopeTree,
-        config: InjectGlobalVariablesConfig,
-    ) -> Self {
+impl<'a> InjectGlobalVariables<'a> {
+    pub fn new(allocator: &'a Allocator, config: InjectGlobalVariablesConfig) -> Self {
         Self {
             ast: AstBuilder::new(allocator),
             config,
             dot_defines: vec![],
             replaced_dot_defines: vec![],
-            symbols,
-            scopes,
         }
     }
 
-    pub fn build(&mut self, program: &mut Program<'a>) {
+    pub fn build(
+        &mut self,
+        symbols: SymbolTable,
+        scopes: ScopeTree,
+        program: &mut Program<'a>,
+    ) -> InjectGlobalVariablesReturn {
+        let mut symbols = symbols;
+        let mut scopes = scopes;
         // Step 1: slow path where visiting the AST is required to replace dot defines.
         let dot_defines = self
             .config
@@ -157,7 +162,7 @@ impl<'a, 'b> InjectGlobalVariables<'a, 'b> {
 
         if !dot_defines.is_empty() {
             self.dot_defines = dot_defines;
-            self.visit_program(program);
+            (symbols, scopes) = traverse_mut(self, self.ast.allocator, program, symbols, scopes);
         }
 
         // Step 2: find all the injects that are referenced.
@@ -172,17 +177,19 @@ impl<'a, 'b> InjectGlobalVariables<'a, 'b> {
                 } else if self.replaced_dot_defines.iter().any(|d| d.0 == i.specifier.local()) {
                     false
                 } else {
-                    self.scopes.root_unresolved_references().contains_key(i.specifier.local())
+                    scopes.root_unresolved_references().contains_key(i.specifier.local())
                 }
             })
             .cloned()
             .collect::<Vec<_>>();
 
         if injects.is_empty() {
-            return;
+            return InjectGlobalVariablesReturn { symbols, scopes };
         }
 
         self.inject_imports(&injects, program);
+
+        InjectGlobalVariablesReturn { symbols, scopes }
     }
 
     fn inject_imports(&self, injects: &[InjectImport], program: &mut Program<'a>) {
@@ -227,10 +234,10 @@ impl<'a, 'b> InjectGlobalVariables<'a, 'b> {
         }
     }
 
-    fn replace_dot_defines(&mut self, expr: &mut Expression<'a>) {
+    fn replace_dot_defines(&mut self, expr: &mut Expression<'a>, ctx: &mut TraverseCtx<'a>) {
         if let Expression::StaticMemberExpression(member) = expr {
             for DotDefineState { dot_define, value_atom } in &mut self.dot_defines {
-                if ReplaceGlobalDefines::is_dot_define(self.symbols, dot_define, member) {
+                if ReplaceGlobalDefines::is_dot_define(ctx.symbols(), dot_define, member) {
                     // If this is first replacement made for this dot define,
                     // create `Atom` for replacement, and record in `replaced_dot_defines`
                     if value_atom.is_none() {

crates/oxc_minifier/src/plugins/replace_global_defines.rs

@@ -1,12 +1,13 @@
 use std::{cmp::Ordering, sync::Arc};
 
 use oxc_allocator::Allocator;
-use oxc_ast::{ast::*, visit::walk_mut, AstBuilder, VisitMut};
+use oxc_ast::ast::*;
 use oxc_diagnostics::OxcDiagnostic;
 use oxc_parser::Parser;
-use oxc_semantic::{IsGlobalReference, SymbolTable};
+use oxc_semantic::{IsGlobalReference, ScopeTree, SymbolTable};
 use oxc_span::{CompactStr, SourceType};
 use oxc_syntax::identifier::is_identifier_name;
+use oxc_traverse::{traverse_mut, Traverse, TraverseCtx};
 
 /// Configuration for [ReplaceGlobalDefines].
 ///
@@ -162,53 +163,57 @@ impl ReplaceGlobalDefinesConfig {
     }
 }
 
+#[must_use]
+pub struct ReplaceGlobalDefinesReturn {
+    pub symbols: SymbolTable,
+    pub scopes: ScopeTree,
+}
+
 /// Replace Global Defines.
 ///
 /// References:
 ///
 /// * <https://esbuild.github.io/api/#define>
 /// * <https://github.com/terser/terser?tab=readme-ov-file#conditional-compilation>
 /// * <https://github.com/evanw/esbuild/blob/9c13ae1f06dfa909eb4a53882e3b7e4216a503fe/internal/config/globals.go#L852-L1014>
-pub struct ReplaceGlobalDefines<'a, 'b> {
-    ast: AstBuilder<'a>,
-    symbols: &'b mut SymbolTable,
+pub struct ReplaceGlobalDefines<'a> {
+    allocator: &'a Allocator,
     config: ReplaceGlobalDefinesConfig,
 }
 
-impl<'a, 'b> VisitMut<'a> for ReplaceGlobalDefines<'a, 'b> {
-    fn visit_expression(&mut self, expr: &mut Expression<'a>) {
-        self.replace_identifier_defines(expr);
-        self.replace_dot_defines(expr);
-        walk_mut::walk_expression(self, expr);
+impl<'a> Traverse<'a> for ReplaceGlobalDefines<'a> {
+    fn enter_expression(&mut self, expr: &mut Expression<'a>, ctx: &mut TraverseCtx<'a>) {
+        self.replace_identifier_defines(expr, ctx);
+        self.replace_dot_defines(expr, ctx);
     }
 }
 
-impl<'a, 'b> ReplaceGlobalDefines<'a, 'b> {
-    pub fn new(
-        allocator: &'a Allocator,
-        symbols: &'b mut SymbolTable,
-        config: ReplaceGlobalDefinesConfig,
-    ) -> Self {
-        Self { ast: AstBuilder::new(allocator), symbols, config }
+impl<'a> ReplaceGlobalDefines<'a> {
+    pub fn new(allocator: &'a Allocator, config: ReplaceGlobalDefinesConfig) -> Self {
+        Self { allocator, config }
     }
 
-    pub fn build(&mut self, program: &mut Program<'a>) {
-        self.visit_program(program);
+    pub fn build(
+        &mut self,
+        symbols: SymbolTable,
+        scopes: ScopeTree,
+        program: &mut Program<'a>,
+    ) -> ReplaceGlobalDefinesReturn {
+        let (symbols, scopes) = traverse_mut(self, self.allocator, program, symbols, scopes);
+        ReplaceGlobalDefinesReturn { symbols, scopes }
     }
 
     // Construct a new expression because we don't have ast clone right now.
     fn parse_value(&self, source_text: &str) -> Expression<'a> {
         // Allocate the string lazily because replacement happens rarely.
-        let source_text = self.ast.allocator.alloc(source_text.to_string());
+        let source_text = self.allocator.alloc(source_text.to_string());
         // Unwrapping here, it should already be checked by [ReplaceGlobalDefinesConfig::new].
-        Parser::new(self.ast.allocator, source_text, SourceType::default())
-            .parse_expression()
-            .unwrap()
+        Parser::new(self.allocator, source_text, SourceType::default()).parse_expression().unwrap()
     }
 
-    fn replace_identifier_defines(&self, expr: &mut Expression<'a>) {
+    fn replace_identifier_defines(&self, expr: &mut Expression<'a>, ctx: &mut TraverseCtx<'a>) {
         let Expression::Identifier(ident) = expr else { return };
-        if !ident.is_global_reference(self.symbols) {
+        if !ident.is_global_reference(ctx.symbols()) {
             return;
         }
         for (key, value) in &self.config.0.identifier {
@@ -220,12 +225,12 @@ impl<'a, 'b> ReplaceGlobalDefines<'a, 'b> {
         }
     }
 
-    fn replace_dot_defines(&mut self, expr: &mut Expression<'a>) {
+    fn replace_dot_defines(&mut self, expr: &mut Expression<'a>, ctx: &mut TraverseCtx<'a>) {
         let Expression::StaticMemberExpression(member) = expr else {
             return;
         };
         for dot_define in &self.config.0.dot {
-            if Self::is_dot_define(self.symbols, dot_define, member) {
+            if Self::is_dot_define(ctx.symbols(), dot_define, member) {
                 let value = self.parse_value(&dot_define.value);
                 *expr = value;
                 return;

crates/oxc_minifier/tests/plugins/inject_global_variables.rs

@@ -16,11 +16,11 @@ pub(crate) fn test(source_text: &str, expected: &str, config: InjectGlobalVariab
     let allocator = Allocator::default();
     let ret = Parser::new(&allocator, source_text, source_type).parse();
     let program = allocator.alloc(ret.program);
-    let (mut symbols, mut scopes) = SemanticBuilder::new(source_text)
+    let (symbols, scopes) = SemanticBuilder::new(source_text)
         .build(program)
         .semantic
         .into_symbol_table_and_scope_tree();
-    InjectGlobalVariables::new(&allocator, &mut symbols, &mut scopes, config).build(program);
+    let _ = InjectGlobalVariables::new(&allocator, config).build(symbols, scopes, program);
     let result = CodeGenerator::new()
         .with_options(CodegenOptions { single_quote: true, ..CodegenOptions::default() })
         .build(program)

crates/oxc_minifier/tests/plugins/replace_global_defines.rs

@@ -12,11 +12,11 @@ pub(crate) fn test(source_text: &str, expected: &str, config: ReplaceGlobalDefin
     let allocator = Allocator::default();
     let ret = Parser::new(&allocator, source_text, source_type).parse();
     let program = allocator.alloc(ret.program);
-    let (mut symbols, _scopes) = SemanticBuilder::new(source_text)
+    let (symbols, scopes) = SemanticBuilder::new(source_text)
         .build(program)
         .semantic
         .into_symbol_table_and_scope_tree();
-    ReplaceGlobalDefines::new(&allocator, &mut symbols, config).build(program);
+    let _ = ReplaceGlobalDefines::new(&allocator, config).build(symbols, scopes, program);
     let result = CodeGenerator::new()
         .with_options(CodegenOptions { single_quote: true, ..CodegenOptions::default() })
         .build(program)