From 00e28029aef3a9a37e860d82fe2f7cfa46155f16 Mon Sep 17 00:00:00 2001 From: overlookmotel <557937+overlookmotel@users.noreply.github.com> Date: Tue, 1 Oct 2024 07:40:17 +0000 Subject: [PATCH] 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` multiple times, which can be slow with large files. --- crates/oxc_transformer/src/common/mod.rs | 10 ++- .../src/common/top_level_statements.rs | 70 +++++++++++++++++++ .../src/common/var_declarations.rs | 44 ++++++++---- crates/oxc_transformer/src/context.rs | 8 ++- 4 files changed, 116 insertions(+), 16 deletions(-) create mode 100644 crates/oxc_transformer/src/common/top_level_statements.rs diff --git a/crates/oxc_transformer/src/common/mod.rs b/crates/oxc_transformer/src/common/mod.rs index d8785e7d2..3e008854a 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 000000000..bf28330df --- /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 fb3d21306..65566a40f 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 2570e70a8..2dbfce889 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(), } }