refactor(transformer): move functionality of common transforms into stores (#6243)

In common transforms, move all functionality into the `*Store` structs. The `Traverse` impls become just thin wrappers which call into methods on the `*Store` structs.

I think this is clearer because now reading these files from top-to-bottom, you have the code that inserts into the stores before the code that reads from the stores and acts on it.

crates/oxc_transformer/src/common/module_imports.rs

@@ -48,12 +48,7 @@ impl<'a, 'ctx> ModuleImports<'a, 'ctx> {
 
 impl<'a, 'ctx> Traverse<'a> for ModuleImports<'a, 'ctx> {
     fn exit_program(&mut self, _program: &mut Program<'a>, ctx: &mut TraverseCtx<'a>) {
-        let mut imports = self.ctx.module_imports.imports.borrow_mut();
+        self.ctx.module_imports.insert_into_program(self.ctx, ctx);
-        if self.ctx.source_type.is_script() {
-            self.insert_require_statements(&mut imports, ctx);
-        } else {
-            self.insert_import_statements(&mut imports, ctx);
-        }
     }
 }
 
@@ -69,22 +64,85 @@ impl<'a> NamedImport<'a> {
     }
 }
 
-impl<'a, 'ctx> ModuleImports<'a, 'ctx> {
+/// Store for `import` / `require` statements to be added at top of program.
+///
+/// TODO(improve-on-babel): Insertion order does not matter. We only have to use `IndexMap`
+/// to produce output that's the same as Babel's.
+/// Substitute `FxHashMap` once we don't need to match Babel's output exactly.
+pub struct ModuleImportsStore<'a> {
+    imports: RefCell<IndexMap<Atom<'a>, Vec<NamedImport<'a>>>>,
+}
+
+// Public methods
+impl<'a> ModuleImportsStore<'a> {
+    /// Create new `ModuleImportsStore`.
+    pub fn new() -> Self {
+        Self { imports: RefCell::new(IndexMap::default()) }
+    }
+
+    /// Add `import` or `require` to top of program.
+    ///
+    /// Which it will be depends on the source type.
+    ///
+    /// * `import { named_import } from 'source';` or
+    /// * `var named_import = require('source');`
+    ///
+    /// If `front` is `true`, `import`/`require` is added to front of the `import`s/`require`s.
+    /// TODO(improve-on-babel): `front` option is only required to pass one of Babel's tests. Output
+    /// without it is still valid. Remove this once our output doesn't need to match Babel exactly.
+    pub fn add_import(&self, source: Atom<'a>, import: NamedImport<'a>, front: bool) {
+        match self.imports.borrow_mut().entry(source) {
+            IndexMapEntry::Occupied(mut entry) => {
+                entry.get_mut().push(import);
+                if front && entry.index() != 0 {
+                    entry.move_index(0);
+                }
+            }
+            IndexMapEntry::Vacant(entry) => {
+                let named_imports = vec![import];
+                if front {
+                    entry.shift_insert(0, named_imports);
+                } else {
+                    entry.insert(named_imports);
+                }
+            }
+        }
+    }
+
+    /// Returns `true` if no imports have been scheduled for insertion.
+    pub fn is_empty(&self) -> bool {
+        self.imports.borrow().is_empty()
+    }
+}
+
+// Internal methods
+impl<'a> ModuleImportsStore<'a> {
+    /// Insert `import` / `require` statements at top of program.
+    fn insert_into_program(&self, transform_ctx: &TransformCtx<'a>, ctx: &mut TraverseCtx<'a>) {
+        if transform_ctx.source_type.is_script() {
+            self.insert_require_statements(transform_ctx, ctx);
+        } else {
+            self.insert_import_statements(transform_ctx, ctx);
+        }
+    }
+
     fn insert_import_statements(
-        &mut self,
+        &self,
-        imports: &mut IndexMap<Atom<'a>, Vec<NamedImport<'a>>>,
+        transform_ctx: &TransformCtx<'a>,
         ctx: &mut TraverseCtx<'a>,
     ) {
+        let mut imports = self.imports.borrow_mut();
         let stmts =
             imports.drain(..).map(|(source, names)| Self::get_named_import(source, names, ctx));
-        self.ctx.top_level_statements.insert_statements(stmts);
+        transform_ctx.top_level_statements.insert_statements(stmts);
     }
 
     fn insert_require_statements(
-        &mut self,
+        &self,
-        imports: &mut IndexMap<Atom<'a>, Vec<NamedImport<'a>>>,
+        transform_ctx: &TransformCtx<'a>,
         ctx: &mut TraverseCtx<'a>,
     ) {
+        let mut imports = self.imports.borrow_mut();
         if imports.is_empty() {
             return;
         }
@@ -93,7 +151,7 @@ impl<'a, 'ctx> ModuleImports<'a, 'ctx> {
         let stmts = imports
             .drain(..)
             .map(|(source, names)| Self::get_require(source, names, require_symbol_id, ctx));
-        self.ctx.top_level_statements.insert_statements(stmts);
+        transform_ctx.top_level_statements.insert_statements(stmts);
     }
 
     fn get_named_import(
@@ -157,51 +215,3 @@ impl<'a, 'ctx> ModuleImports<'a, 'ctx> {
         ctx.ast.statement_declaration(var_decl)
     }
 }
-
-/// Store for `import` / `require` statements to be added at top of program.
-///
-/// TODO(improve-on-babel): Insertion order does not matter. We only have to use `IndexMap`
-/// to produce output that's the same as Babel's.
-/// Substitute `FxHashMap` once we don't need to match Babel's output exactly.
-pub struct ModuleImportsStore<'a> {
-    imports: RefCell<IndexMap<Atom<'a>, Vec<NamedImport<'a>>>>,
-}
-
-impl<'a> ModuleImportsStore<'a> {
-    pub fn new() -> ModuleImportsStore<'a> {
-        Self { imports: RefCell::new(IndexMap::default()) }
-    }
-
-    /// Add `import` or `require` to top of program.
-    ///
-    /// Which it will be depends on the source type.
-    ///
-    /// * `import { named_import } from 'source';` or
-    /// * `var named_import = require('source');`
-    ///
-    /// If `front` is `true`, `import`/`require` is added to front of the `import`s/`require`s.
-    /// TODO(improve-on-babel): `front` option is only required to pass one of Babel's tests. Output
-    /// without it is still valid. Remove this once our output doesn't need to match Babel exactly.
-    pub fn add_import(&self, source: Atom<'a>, import: NamedImport<'a>, front: bool) {
-        match self.imports.borrow_mut().entry(source) {
-            IndexMapEntry::Occupied(mut entry) => {
-                entry.get_mut().push(import);
-                if front && entry.index() != 0 {
-                    entry.move_index(0);
-                }
-            }
-            IndexMapEntry::Vacant(entry) => {
-                let named_imports = vec![import];
-                if front {
-                    entry.shift_insert(0, named_imports);
-                } else {
-                    entry.insert(named_imports);
-                }
-            }
-        }
-    }
-
-    pub fn is_empty(&self) -> bool {
-        self.imports.borrow().is_empty()
-    }
-}

crates/oxc_transformer/src/common/top_level_statements.rs

@@ -35,7 +35,38 @@ impl<'a, 'ctx> TopLevelStatements<'a, 'ctx> {
 
 impl<'a, 'ctx> Traverse<'a> for TopLevelStatements<'a, 'ctx> {
     fn exit_program(&mut self, program: &mut Program<'a>, _ctx: &mut TraverseCtx<'a>) {
-        let mut stmts = self.ctx.top_level_statements.stmts.borrow_mut();
+        self.ctx.top_level_statements.insert_into_program(program);
+    }
+}
+
+/// Store for statements to be added at top of program
+pub struct TopLevelStatementsStore<'a> {
+    stmts: RefCell<Vec<Statement<'a>>>,
+}
+
+// Public methods
+impl<'a> TopLevelStatementsStore<'a> {
+    /// Create new `TopLevelStatementsStore`.
+    pub fn new() -> Self {
+        Self { stmts: RefCell::new(vec![]) }
+    }
+
+    /// Add a statement to be inserted at top of program.
+    pub fn insert_statement(&self, stmt: Statement<'a>) {
+        self.stmts.borrow_mut().push(stmt);
+    }
+
+    /// Add statements to be inserted at top of program.
+    pub fn insert_statements<I: IntoIterator<Item = Statement<'a>>>(&self, stmts: I) {
+        self.stmts.borrow_mut().extend(stmts);
+    }
+}
+
+// Internal methods
+impl<'a> TopLevelStatementsStore<'a> {
+    /// Insert statements at top of program.
+    fn insert_into_program(&self, program: &mut Program<'a>) {
+        let mut stmts = self.stmts.borrow_mut();
         if stmts.is_empty() {
             return;
         }
@@ -50,26 +81,3 @@ impl<'a, 'ctx> Traverse<'a> for TopLevelStatements<'a, 'ctx> {
         program.body.splice(index..index, stmts.drain(..));
     }
 }
-
-/// Store for statements to be added at top of program
-pub struct TopLevelStatementsStore<'a> {
-    stmts: RefCell<Vec<Statement<'a>>>,
-}
-
-impl<'a> TopLevelStatementsStore<'a> {
-    pub fn new() -> Self {
-        Self { stmts: RefCell::new(vec![]) }
-    }
-}
-
-impl<'a> TopLevelStatementsStore<'a> {
-    /// Add a statement to be inserted at top of program.
-    pub fn insert_statement(&self, stmt: Statement<'a>) {
-        self.stmts.borrow_mut().push(stmt);
-    }
-
-    /// Add statements to be inserted at top of program.
-    pub fn insert_statements<I: IntoIterator<Item = Statement<'a>>>(&self, stmts: I) {
-        self.stmts.borrow_mut().extend(stmts);
-    }
-}

crates/oxc_transformer/src/common/var_declarations.rs

@@ -38,51 +38,20 @@ impl<'a, 'ctx> VarDeclarations<'a, 'ctx> {
 }
 
 impl<'a, 'ctx> Traverse<'a> for VarDeclarations<'a, 'ctx> {
-    fn exit_program(&mut self, _program: &mut Program<'a>, ctx: &mut TraverseCtx<'a>) {
-        if let Some(stmt) = self.get_var_statement(ctx) {
-            // Delegate to `TopLevelStatements`
-            self.ctx.top_level_statements.insert_statement(stmt);
-        }
-
-        let stack = self.ctx.var_declarations.stack.borrow();
-        debug_assert!(stack.len() == 1);
-        debug_assert!(stack.last().is_none());
-    }
-
     fn enter_statements(
         &mut self,
         _stmts: &mut Vec<'a, Statement<'a>>,
         _ctx: &mut TraverseCtx<'a>,
     ) {
-        let mut stack = self.ctx.var_declarations.stack.borrow_mut();
+        self.ctx.var_declarations.record_entering_statements();
-        stack.push(None);
     }
 
     fn exit_statements(&mut self, stmts: &mut Vec<'a, Statement<'a>>, ctx: &mut TraverseCtx<'a>) {
-        if matches!(ctx.parent(), Ancestor::ProgramBody(_)) {
+        self.ctx.var_declarations.insert_into_statements(stmts, ctx);
-            // Handle in `exit_program` instead
-            return;
-        }
-
-        if let Some(stmt) = self.get_var_statement(ctx) {
-            stmts.insert(0, stmt);
-        }
     }
-}
 
-impl<'a, 'ctx> VarDeclarations<'a, 'ctx> {
+    fn exit_program(&mut self, _program: &mut Program<'a>, ctx: &mut TraverseCtx<'a>) {
-    fn get_var_statement(&mut self, ctx: &mut TraverseCtx<'a>) -> Option<Statement<'a>> {
+        self.ctx.var_declarations.insert_into_program(self.ctx, ctx);
-        let mut stack = self.ctx.var_declarations.stack.borrow_mut();
-        let declarators = stack.pop()?;
-        debug_assert!(!declarators.is_empty());
-
-        let stmt = Statement::VariableDeclaration(ctx.ast.alloc_variable_declaration(
-            SPAN,
-            VariableDeclarationKind::Var,
-            declarators,
-            false,
-        ));
-        Some(stmt)
     }
 }
 
@@ -91,13 +60,13 @@ pub struct VarDeclarationsStore<'a> {
     stack: RefCell<SparseStack<Vec<'a, VariableDeclarator<'a>>>>,
 }
 
+// Public methods
 impl<'a> VarDeclarationsStore<'a> {
+    /// Create new `VarDeclarationsStore`.
     pub fn new() -> Self {
         Self { stack: RefCell::new(SparseStack::new()) }
     }
-}
 
-impl<'a> VarDeclarationsStore<'a> {
     /// Add a `VariableDeclarator` to be inserted at top of current enclosing statement block,
     /// given `name` and `symbol_id`.
     pub fn insert(
@@ -132,3 +101,52 @@ impl<'a> VarDeclarationsStore<'a> {
         stack.last_mut_or_init(|| ctx.ast.vec()).push(declarator);
     }
 }
+
+// Internal methods
+impl<'a> VarDeclarationsStore<'a> {
+    fn record_entering_statements(&self) {
+        let mut stack = self.stack.borrow_mut();
+        stack.push(None);
+    }
+
+    fn insert_into_statements(
+        &self,
+        stmts: &mut Vec<'a, Statement<'a>>,
+        ctx: &mut TraverseCtx<'a>,
+    ) {
+        if matches!(ctx.parent(), Ancestor::ProgramBody(_)) {
+            // Handle in `insert_into_program` instead
+            return;
+        }
+
+        if let Some(stmt) = self.get_var_statement(ctx) {
+            stmts.insert(0, stmt);
+        }
+    }
+
+    fn insert_into_program(&self, transform_ctx: &TransformCtx<'a>, ctx: &mut TraverseCtx<'a>) {
+        if let Some(stmt) = self.get_var_statement(ctx) {
+            // Delegate to `TopLevelStatements`
+            transform_ctx.top_level_statements.insert_statement(stmt);
+        }
+
+        // Check stack is emptied
+        let stack = self.stack.borrow();
+        debug_assert!(stack.len() == 1);
+        debug_assert!(stack.last().is_none());
+    }
+
+    fn get_var_statement(&self, ctx: &mut TraverseCtx<'a>) -> Option<Statement<'a>> {
+        let mut stack = self.stack.borrow_mut();
+        let declarators = stack.pop()?;
+        debug_assert!(!declarators.is_empty());
+
+        let stmt = Statement::VariableDeclaration(ctx.ast.alloc_variable_declaration(
+            SPAN,
+            VariableDeclarationKind::Var,
+            declarators,
+            false,
+        ));
+        Some(stmt)
+    }
+}