refactor(transformer): introduce TopLevelStatements common transform (#6185)

Introduce `TopLevelStatements` common transform. It holds a `Vec` of statements to be inserted at the top of the program, and other transforms can push statements to it.

All statements will be inserted in one go at the end of traversal, to avoid shuffling up the `Vec<Statement>` multiple times, which can be slow with large files.

crates/oxc_transformer/src/common/mod.rs

@@ -6,24 +6,30 @@ use oxc_traverse::{Traverse, TraverseCtx};
 
 use crate::TransformCtx;
 
+pub mod top_level_statements;
 pub mod var_declarations;
 
+use top_level_statements::TopLevelStatements;
 use var_declarations::VarDeclarations;
 
 pub struct Common<'a, 'ctx> {
     var_declarations: VarDeclarations<'a, 'ctx>,
+    top_level_statements: TopLevelStatements<'a, 'ctx>,
 }
 
 impl<'a, 'ctx> Common<'a, 'ctx> {
     pub fn new(ctx: &'ctx TransformCtx<'a>) -> Self {
-        Self { var_declarations: VarDeclarations::new(ctx) }
+        Self {
+            var_declarations: VarDeclarations::new(ctx),
+            top_level_statements: TopLevelStatements::new(ctx),
+        }
     }
 }
 
 impl<'a, 'ctx> Traverse<'a> for Common<'a, 'ctx> {
-    #[inline] // Inline because it's no-op in release mode
     fn exit_program(&mut self, program: &mut Program<'a>, ctx: &mut TraverseCtx<'a>) {
         self.var_declarations.exit_program(program, ctx);
+        self.top_level_statements.exit_program(program, ctx);
     }
 
     fn enter_statements(&mut self, stmts: &mut Vec<'a, Statement<'a>>, ctx: &mut TraverseCtx<'a>) {

crates/oxc_transformer/src/common/top_level_statements.rs

@@ -0,0 +1,70 @@
+//! Utility transform to add statements to top of program.
+//!
+//! `TopLevelStatementsStore` contains a `Vec<Statement>`. It is stored on `TransformCtx`.
+//!
+//! `TopLevelStatements` transform inserts those statements at top of program.
+//!
+//! Other transforms can add statements to the store with `TopLevelStatementsStore::insert_statement`:
+//!
+//! ```rs
+//! self.ctx.top_level_statements.insert_statement(stmt);
+//! ```
+
+use std::cell::RefCell;
+
+use oxc_ast::ast::*;
+use oxc_traverse::{Traverse, TraverseCtx};
+
+use crate::TransformCtx;
+
+/// Transform that inserts any statements which have been requested insertion via `TopLevelStatementsStore`
+/// to top of the program.
+///
+/// Insertions are made after any existing `import` statements.
+///
+/// Must run after all other transforms.
+pub struct TopLevelStatements<'a, 'ctx> {
+    ctx: &'ctx TransformCtx<'a>,
+}
+
+impl<'a, 'ctx> TopLevelStatements<'a, 'ctx> {
+    pub fn new(ctx: &'ctx TransformCtx<'a>) -> Self {
+        Self { 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();
+        if stmts.is_empty() {
+            return;
+        }
+
+        // Insert statements after any existing `import` statements
+        let index = program
+            .body
+            .iter()
+            .rposition(|stmt| matches!(stmt, Statement::ImportDeclaration(_)))
+            .map_or(0, |i| i + 1);
+
+        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);
+    }
+}

crates/oxc_transformer/src/common/var_declarations.rs

@@ -25,7 +25,7 @@ use crate::{helpers::stack::SparseStack, TransformCtx};
 /// Transform that maintains the stack of `Vec<VariableDeclarator>`s, and adds a `var` statement
 /// to top of a statement block if another transform has requested that.
 ///
-/// Must run after all other transforms.
+/// Must run after all other transforms except `TopLevelStatements`.
 pub struct VarDeclarations<'a, 'ctx> {
     ctx: &'ctx TransformCtx<'a>,
 }
@@ -37,8 +37,12 @@ impl<'a, 'ctx> VarDeclarations<'a, 'ctx> {
 }
 
 impl<'a, 'ctx> Traverse<'a> for VarDeclarations<'a, 'ctx> {
-    #[inline] // Inline because it's no-op in release mode
-    fn exit_program(&mut self, _program: &mut Program<'a>, _ctx: &mut TraverseCtx<'a>) {
+    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 declarators = self.ctx.var_declarations.declarators.borrow();
         debug_assert!(declarators.len() == 1);
         debug_assert!(declarators.last().is_none());
@@ -54,17 +58,31 @@ impl<'a, 'ctx> Traverse<'a> for VarDeclarations<'a, 'ctx> {
     }
 
     fn exit_statements(&mut self, stmts: &mut Vec<'a, Statement<'a>>, ctx: &mut TraverseCtx<'a>) {
-        let mut declarators = self.ctx.var_declarations.declarators.borrow_mut();
-        if let Some(declarators) = declarators.pop() {
-            debug_assert!(!declarators.is_empty());
-            let variable = ctx.ast.alloc_variable_declaration(
-                SPAN,
-                VariableDeclarationKind::Var,
-                declarators,
-                false,
-            );
-            stmts.insert(0, Statement::VariableDeclaration(variable));
+        if ctx.ancestors_depth() == 2 {
+            // Top level. Handle in `exit_program` instead.
+            // (depth 1 = None, depth 2 = Program)
+            return;
         }
+
+        if let Some(stmt) = self.get_var_statement(ctx) {
+            stmts.insert(0, stmt);
+        }
+    }
+}
+
+impl<'a, 'ctx> VarDeclarations<'a, 'ctx> {
+    fn get_var_statement(&mut self, ctx: &mut TraverseCtx<'a>) -> Option<Statement<'a>> {
+        let mut declarators = self.ctx.var_declarations.declarators.borrow_mut();
+        let declarators = declarators.pop()?;
+        debug_assert!(!declarators.is_empty());
+
+        let stmt = Statement::VariableDeclaration(ctx.ast.alloc_variable_declaration(
+            SPAN,
+            VariableDeclarationKind::Var,
+            declarators,
+            false,
+        ));
+        Some(stmt)
     }
 }
 

crates/oxc_transformer/src/context.rs

@@ -10,7 +10,10 @@ use oxc_diagnostics::OxcDiagnostic;
 use oxc_span::SourceType;
 
 use crate::{
-    common::var_declarations::VarDeclarationsStore, helpers::module_imports::ModuleImports,
+    common::{
+        top_level_statements::TopLevelStatementsStore, var_declarations::VarDeclarationsStore,
+    },
+    helpers::module_imports::ModuleImports,
     TransformOptions,
 };
 
@@ -36,6 +39,8 @@ pub struct TransformCtx<'a> {
     pub module_imports: ModuleImports<'a>,
     /// Manage inserting `var` statements globally
     pub var_declarations: VarDeclarationsStore<'a>,
+    /// Manage inserting statements at top of program globally
+    pub top_level_statements: TopLevelStatementsStore<'a>,
 }
 
 impl<'a> TransformCtx<'a> {
@@ -65,6 +70,7 @@ impl<'a> TransformCtx<'a> {
             trivias,
             module_imports: ModuleImports::new(),
             var_declarations: VarDeclarationsStore::new(),
+            top_level_statements: TopLevelStatementsStore::new(),
         }
     }