diff --git a/crates/oxc_transformer/src/common/mod.rs b/crates/oxc_transformer/src/common/mod.rs index d8785e7d29119..3e008854ab13e 100644 --- a/crates/oxc_transformer/src/common/mod.rs +++ b/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>) { diff --git a/crates/oxc_transformer/src/common/top_level_statements.rs b/crates/oxc_transformer/src/common/top_level_statements.rs new file mode 100644 index 0000000000000..bf28330dfe64f --- /dev/null +++ b/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`. 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>>, +} + +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); + } +} diff --git a/crates/oxc_transformer/src/common/var_declarations.rs b/crates/oxc_transformer/src/common/var_declarations.rs index fb3d213062d81..65566a40ff141 100644 --- a/crates/oxc_transformer/src/common/var_declarations.rs +++ b/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`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> { + 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) } } diff --git a/crates/oxc_transformer/src/context.rs b/crates/oxc_transformer/src/context.rs index 2570e70a814cc..2dbfce8894e35 100644 --- a/crates/oxc_transformer/src/context.rs +++ b/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(), } }