Skip to content

Commit

Permalink
refactor(transformer): move functionality of common transforms into s…
Browse files Browse the repository at this point in the history 
…tores (#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.
  • Loading branch information
@overlookmotel @7086cmd
overlookmotel authored and 7086cmd committed Oct 2, 2024
1 parent 2a89dfd commit c113f2c
Show file tree
Hide file tree
Showing 3 changed files with 149 additions and 113 deletions.
132 changes: 71 additions & 61 deletions crates/oxc_transformer/src/common/module_imports.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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();
if self.ctx.source_type.is_script() {
self.insert_require_statements(&mut imports, ctx);
} else {
self.insert_import_statements(&mut imports, ctx);
}
self.ctx.module_imports.insert_into_program(self.ctx, ctx);
}
}

Expand All @@ -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,
imports: &mut IndexMap<Atom<'a>, Vec<NamedImport<'a>>>,
&self,
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,
imports: &mut IndexMap<Atom<'a>, Vec<NamedImport<'a>>>,
&self,
transform_ctx: &TransformCtx<'a>,
ctx: &mut TraverseCtx<'a>,
) {
let mut imports = self.imports.borrow_mut();
if imports.is_empty() {
return;
}
Expand All @@ -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(
Expand Down Expand Up @@ -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()
}
}
38 changes: 23 additions & 15 deletions crates/oxc_transformer/src/common/top_level_statements.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -35,19 +35,7 @@ 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();
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(..));
self.ctx.top_level_statements.insert_into_program(program);
}
}

Expand All @@ -56,13 +44,13 @@ 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![]) }
}
}

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);
Expand All @@ -73,3 +61,23 @@ impl<'a> TopLevelStatementsStore<'a> {
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;
}

// 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(..));
}
}
92 changes: 55 additions & 37 deletions crates/oxc_transformer/src/common/var_declarations.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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();
stack.push(None);
self.ctx.var_declarations.record_entering_statements();
}

fn exit_statements(&mut self, stmts: &mut Vec<'a, Statement<'a>>, ctx: &mut TraverseCtx<'a>) {
if matches!(ctx.parent(), Ancestor::ProgramBody(_)) {
// Handle in `exit_program` instead
return;
}

if let Some(stmt) = self.get_var_statement(ctx) {
stmts.insert(0, stmt);
}
self.ctx.var_declarations.insert_into_statements(stmts, ctx);
}
}

impl<'a, 'ctx> VarDeclarations<'a, 'ctx> {
fn get_var_statement(&mut self, ctx: &mut TraverseCtx<'a>) -> Option<Statement<'a>> {
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)
fn exit_program(&mut self, _program: &mut Program<'a>, ctx: &mut TraverseCtx<'a>) {
self.ctx.var_declarations.insert_into_program(self.ctx, ctx);
}
}

Expand All @@ -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(
Expand Down Expand Up @@ -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)
}
}

0 comments on commit c113f2c

Please sign in to comment.