refactor(transformer): convert ModuleImports into common transform (#…

…6186)

An alternative version of #6177.

Convert `ModuleImports` into a common transform. Works much as before, but it inserts `import` / `require` statements by passing them to `TopLevelStatements` common transform, so they get inserted in one go with any other inserted top-level statements. This avoids shuffling up the `Vec<Statement>` multiple times, which can be slow with large files.

`VarDeclarations` also inserts any declarations via `TopLevelStatements` but runs after `ModuleImports`, so can control whether a `var` statement is inserted before or after `import` statements by inserting it via `VarDeclarations` (to appear after `import` statements) or directly into `TopLevelStatements` (to appear before `import` statements). Insertion order is not actually important, but allows us to match Babel's output and pass its tests.

crates/oxc_transformer/src/common/mod.rs

@@ -6,20 +6,24 @@ use oxc_traverse::{Traverse, TraverseCtx};
 
 use crate::TransformCtx;
 
+pub mod module_imports;
 pub mod top_level_statements;
 pub mod var_declarations;
 
+use module_imports::ModuleImports;
 use top_level_statements::TopLevelStatements;
 use var_declarations::VarDeclarations;
 
 pub struct Common<'a, 'ctx> {
+    module_imports: ModuleImports<'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 {
+            module_imports: ModuleImports::new(ctx),
             var_declarations: VarDeclarations::new(ctx),
             top_level_statements: TopLevelStatements::new(ctx),
         }
@@ -28,6 +32,7 @@ impl<'a, 'ctx> Common<'a, 'ctx> {
 
 impl<'a, 'ctx> Traverse<'a> for Common<'a, 'ctx> {
     fn exit_program(&mut self, program: &mut Program<'a>, ctx: &mut TraverseCtx<'a>) {
+        self.module_imports.exit_program(program, ctx);
         self.var_declarations.exit_program(program, ctx);
         self.top_level_statements.exit_program(program, ctx);
     }

crates/oxc_transformer/src/helpers/module_imports.rs → crates/oxc_transformer/src/common/module_imports.rs

@@ -1,12 +1,62 @@
+//! Utility transform to add `import` / `require` statements to top of program.
+//!
+//! `ModuleImportsStore` contains an `IndexMap<ImportType<'a>, Vec<NamedImport<'a>>>`.
+//! It is stored on `TransformCtx`.
+//!
+//! `ModuleImports` transform
+//!
+//! Other transforms can add `import`s / `require`s to the store by calling methods of `ModuleImportsStore`:
+//!
+//! ```rs
+//! // import { jsx as _jsx } from 'react';
+//! self.ctx.module_imports.add_import(
+//!     Atom::from("react"),
+//!     NamedImport::new(Atom::from("jsx"), Some(Atom::from("_jsx")), symbol_id)
+//! );
+//!
+//! // var _react = require('react');
+//! self.ctx.module_imports.add_require(
+//!     Atom::from("react"),
+//!     NamedImport::new(Atom::from("_react"), None, symbol_id)
+//! );
+//! ```
+//!
+//! Based on `@babel/helper-module-imports`
+//! <https://github.com/nicolo-ribaudo/babel/tree/main/packages/babel-helper-module-imports>
+
 use std::cell::RefCell;
 
 use indexmap::IndexMap;
-use oxc_allocator::Vec;
+
 use oxc_ast::{ast::*, NONE};
 use oxc_semantic::ReferenceFlags;
 use oxc_span::{Atom, SPAN};
 use oxc_syntax::symbol::SymbolId;
-use oxc_traverse::TraverseCtx;
+use oxc_traverse::{Traverse, TraverseCtx};
+
+use crate::TransformCtx;
+
+pub struct ModuleImports<'a, 'ctx> {
+    ctx: &'ctx TransformCtx<'a>,
+}
+
+impl<'a, 'ctx> ModuleImports<'a, 'ctx> {
+    pub fn new(ctx: &'ctx TransformCtx<'a>) -> Self {
+        Self { 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();
+        let Some((first_import_type, _)) = imports.first() else { return };
+        // Assume all imports are of the same kind
+        match first_import_type.kind {
+            ImportKind::Import => self.insert_import_statements(&mut imports, ctx),
+            ImportKind::Require => self.insert_require_statements(&mut imports, ctx),
+        }
+    }
+}
 
 pub struct NamedImport<'a> {
     imported: Atom<'a>,
@@ -20,7 +70,7 @@ impl<'a> NamedImport<'a> {
     }
 }
 
-#[derive(Hash, Eq, PartialEq)]
+#[derive(Clone, Copy, Hash, Eq, PartialEq)]
 pub enum ImportKind {
     Import,
     Require,
@@ -38,51 +88,34 @@ impl<'a> ImportType<'a> {
     }
 }
 
-/// Manage import statement globally
-/// <https://github.com/nicolo-ribaudo/babel/tree/main/packages/babel-helper-module-imports>
-pub struct ModuleImports<'a> {
-    imports: RefCell<IndexMap<ImportType<'a>, std::vec::Vec<NamedImport<'a>>>>,
-}
-
-impl<'a> ModuleImports<'a> {
-    pub fn new() -> ModuleImports<'a> {
-        Self { imports: RefCell::new(IndexMap::default()) }
-    }
-
-    /// Add `import { named_import } from 'source'`
-    pub fn add_import(&self, source: Atom<'a>, import: NamedImport<'a>) {
-        self.imports
-            .borrow_mut()
-            .entry(ImportType::new(ImportKind::Import, source))
-            .or_default()
-            .push(import);
-    }
-
-    /// Add `var named_import from 'source'`
-    pub fn add_require(&self, source: Atom<'a>, import: NamedImport<'a>, front: bool) {
-        let len = self.imports.borrow().len();
-        self.imports
-            .borrow_mut()
-            .entry(ImportType::new(ImportKind::Require, source))
-            .or_default()
-            .push(import);
-        if front {
-            self.imports.borrow_mut().move_index(len, 0);
-        }
+impl<'a, 'ctx> ModuleImports<'a, 'ctx> {
+    fn insert_import_statements(
+        &mut self,
+        imports: &mut IndexMap<ImportType<'a>, Vec<NamedImport<'a>>>,
+        ctx: &mut TraverseCtx<'a>,
+    ) {
+        let stmts = imports.drain(..).map(|(import_type, names)| {
+            debug_assert!(import_type.kind == ImportKind::Import);
+            Self::get_named_import(import_type.source, names, ctx)
+        });
+        self.ctx.top_level_statements.insert_statements(stmts);
     }
 
-    pub fn get_import_statements(&self, ctx: &mut TraverseCtx<'a>) -> Vec<'a, Statement<'a>> {
-        ctx.ast.vec_from_iter(self.imports.borrow_mut().drain(..).map(|(import_type, names)| {
-            match import_type.kind {
-                ImportKind::Import => Self::get_named_import(import_type.source, names, ctx),
-                ImportKind::Require => Self::get_require(import_type.source, names, ctx),
-            }
-        }))
+    fn insert_require_statements(
+        &mut self,
+        imports: &mut IndexMap<ImportType<'a>, Vec<NamedImport<'a>>>,
+        ctx: &mut TraverseCtx<'a>,
+    ) {
+        let stmts = imports.drain(..).map(|(import_type, names)| {
+            debug_assert!(import_type.kind == ImportKind::Require);
+            Self::get_require(import_type.source, names, ctx)
+        });
+        self.ctx.top_level_statements.insert_statements(stmts);
     }
 
     fn get_named_import(
         source: Atom<'a>,
-        names: std::vec::Vec<NamedImport<'a>>,
+        names: Vec<NamedImport<'a>>,
         ctx: &mut TraverseCtx<'a>,
     ) -> Statement<'a> {
         let specifiers = ctx.ast.vec_from_iter(names.into_iter().map(|name| {
@@ -137,3 +170,44 @@ impl<'a> ModuleImports<'a> {
         ctx.ast.statement_declaration(var_decl)
     }
 }
+
+/// Store for `import` / `require` statements to be added at top of program
+pub struct ModuleImportsStore<'a> {
+    imports: RefCell<IndexMap<ImportType<'a>, Vec<NamedImport<'a>>>>,
+}
+
+impl<'a> ModuleImportsStore<'a> {
+    pub fn new() -> ModuleImportsStore<'a> {
+        Self { imports: RefCell::new(IndexMap::default()) }
+    }
+
+    /// Add `import { named_import } from 'source'`
+    pub fn add_import(&self, source: Atom<'a>, import: NamedImport<'a>) {
+        self.imports
+            .borrow_mut()
+            .entry(ImportType::new(ImportKind::Import, source))
+            .or_default()
+            .push(import);
+    }
+
+    /// Add `var named_import from 'source'`.
+    ///
+    /// If `front` is true, `require` is added to top of the `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_require(&self, source: Atom<'a>, import: NamedImport<'a>, front: bool) {
+        let len = self.imports.borrow().len();
+        self.imports
+            .borrow_mut()
+            .entry(ImportType::new(ImportKind::Require, source))
+            .or_default()
+            .push(import);
+        if front {
+            self.imports.borrow_mut().move_index(len, 0);
+        }
+    }
+
+    pub fn is_empty(&self) -> bool {
+        self.imports.borrow().is_empty()
+    }
+}

crates/oxc_transformer/src/common/top_level_statements.rs

@@ -67,4 +67,9 @@ impl<'a> TopLevelStatementsStore<'a> {
     pub fn insert_statement(&self, stmt: Statement<'a>) {
         self.stmts.borrow_mut().push(stmt);
     }
+
+    /// Add statements to be inserted at top of program.
+    pub fn insert_statements<I: IntoIterator<Item = Statement<'a>>>(&self, stmts: I) {
+        self.stmts.borrow_mut().extend(stmts);
+    }
 }

crates/oxc_transformer/src/context.rs

@@ -11,9 +11,9 @@ use oxc_span::SourceType;
 
 use crate::{
     common::{
-        top_level_statements::TopLevelStatementsStore, var_declarations::VarDeclarationsStore,
+        module_imports::ModuleImportsStore, top_level_statements::TopLevelStatementsStore,
+        var_declarations::VarDeclarationsStore,
     },
-    helpers::module_imports::ModuleImports,
     TransformOptions,
 };
 
@@ -36,7 +36,7 @@ pub struct TransformCtx<'a> {
 
     // Helpers
     /// Manage import statement globally
-    pub module_imports: ModuleImports<'a>,
+    pub module_imports: ModuleImportsStore<'a>,
     /// Manage inserting `var` statements globally
     pub var_declarations: VarDeclarationsStore<'a>,
     /// Manage inserting statements at top of program globally
@@ -68,7 +68,7 @@ impl<'a> TransformCtx<'a> {
             source_type,
             source_text,
             trivias,
-            module_imports: ModuleImports::new(),
+            module_imports: ModuleImportsStore::new(),
             var_declarations: VarDeclarationsStore::new(),
             top_level_statements: TopLevelStatementsStore::new(),
         }

crates/oxc_transformer/src/lib.rs

@@ -29,7 +29,6 @@ mod plugins;
 
 mod helpers {
     pub mod bindings;
-    pub mod module_imports;
     pub mod stack;
 }
 

crates/oxc_transformer/src/react/jsx.rs

@@ -106,8 +106,7 @@ pub use super::{
     options::{JsxOptions, JsxRuntime},
 };
 use crate::{
-    helpers::{bindings::BoundIdentifier, module_imports::NamedImport},
-    TransformCtx,
+    common::module_imports::NamedImport, helpers::bindings::BoundIdentifier, TransformCtx,
 };
 
 pub struct ReactJsx<'a, 'ctx> {
@@ -170,6 +169,9 @@ impl<'a, 'ctx> AutomaticScriptBindings<'a, 'ctx> {
         if self.require_create_element.is_none() {
             let source =
                 get_import_source(self.jsx_runtime_importer.as_str(), self.react_importer_len);
+            // We have to insert this `require` above `require("react/jsx-runtime")`
+            // just to pass one of Babel's tests, but the order doesn't actually matter.
+            // TODO(improve-on-babel): Remove this once we don't need our output to match Babel exactly.
             let id = self.add_require_statement("react", source, true, ctx);
             self.require_create_element = Some(id);
         }
@@ -444,8 +446,8 @@ impl<'a, 'ctx> ReactJsx<'a, 'ctx> {
 }
 
 impl<'a, 'ctx> Traverse<'a> for ReactJsx<'a, 'ctx> {
-    fn exit_program(&mut self, program: &mut Program<'a>, ctx: &mut TraverseCtx<'a>) {
-        self.add_runtime_imports(program, ctx);
+    fn exit_program(&mut self, _program: &mut Program<'a>, ctx: &mut TraverseCtx<'a>) {
+        self.insert_var_file_name_statement(ctx);
     }
 
     fn exit_expression(&mut self, expr: &mut Expression<'a>, ctx: &mut TraverseCtx<'a>) {
@@ -468,31 +470,26 @@ impl<'a, 'ctx> ReactJsx<'a, 'ctx> {
         self.ctx.ast
     }
 
-    fn add_runtime_imports(&mut self, program: &mut Program<'a>, ctx: &mut TraverseCtx<'a>) {
-        if self.bindings.is_classic() {
-            if let Some(stmt) = self.jsx_source.get_var_file_name_statement() {
-                program.body.insert(0, stmt);
-            }
-            return;
-        }
-
-        let imports = self.ctx.module_imports.get_import_statements(ctx);
-        let mut index = program
-            .body
-            .iter()
-            .rposition(|stmt| matches!(stmt, Statement::ImportDeclaration(_)))
-            .map_or(0, |i| i + 1);
-
-        if let Some(stmt) = self.jsx_source.get_var_file_name_statement() {
-            program.body.insert(index, stmt);
-            // If source type is module then we need to add the import statement after the var file name statement
-            // Follow the same behavior as babel
-            if !self.is_script() {
-                index += 1;
-            }
+    fn insert_var_file_name_statement(&mut self, ctx: &mut TraverseCtx<'a>) {
+        let Some(declarator) = self.jsx_source.get_var_file_name_declarator() else { return };
+
+        // If is a module, add filename statements before `import`s. If script, then after `require`s.
+        // This is the same behavior as Babel.
+        // If in classic mode, then there are no import statements, so it doesn't matter either way.
+        // TODO(improve-on-babel): Simplify this once we don't need to follow Babel exactly.
+        if self.bindings.is_classic() || !self.is_script() {
+            // Insert before imports - add to `top_level_statements` immediately
+            let stmt = Statement::VariableDeclaration(ctx.ast.alloc_variable_declaration(
+                SPAN,
+                VariableDeclarationKind::Var,
+                self.ctx.ast.vec1(declarator),
+                false,
+            ));
+            self.ctx.top_level_statements.insert_statement(stmt);
+        } else {
+            // Insert after imports - add to `var_declarations`, which are inserted after `require` statements
+            self.ctx.var_declarations.insert_declarator(declarator, ctx);
         }
-
-        program.body.splice(index..index, imports);
     }
 
     fn transform_jsx<'b>(