refactor(transformer): introduce TopLevelStatements common transform

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(),
         }
     }