perf(parser): register import / export statements in module record directly (#12807)

Currently when parsing a list of statements in `parse_directives_and_statements`, we:

1. Check on every turn of the loop whether the statement is a `ModuleDeclaration`, and whether at top level.
2. If so, we call `ModuleRecordBuilder::visit_module_declaration`.
3. `visit_module_declaration` branches again on the type of the `ModuleDeclaration`.

Instead, remove `ModuleRecordBuilder::visit_module_declaration` method, and instead call individual methods of `ModuleRecordBuilder` for each type of declaration directly, when parsing that particular declaration.

This removes the need to check *every* statement in the loop in `parse_directives_and_statements` just in case it's a module declaration.

Gives a small perf boost (+0.2%) on parser benchmarks. Might be a little bit more in real world, as this removes an unpredictable branch, which CodSpeed doesn't measure.

Author: overlookmotel

crates/oxc_parser/src/context.rs

@@ -164,4 +164,8 @@ impl StatementContext {
     pub(crate) fn is_single_statement(self) -> bool {
         !matches!(self, Self::StatementList | Self::TopLevelStatementList)
     }
+
+    pub(crate) fn is_top_level(self) -> bool {
+        self == Self::TopLevelStatementList
+    }
 }

crates/oxc_parser/src/js/module.rs

@@ -5,7 +5,7 @@ use rustc_hash::FxHashMap;
 
 use super::FunctionKind;
 use crate::{
-    Context, ParserImpl, diagnostics,
+    Context, ParserImpl, StatementContext, diagnostics,
     lexer::Kind,
     modifiers::{Modifier, ModifierFlags, ModifierKind, Modifiers},
 };
@@ -61,7 +61,11 @@ impl<'a> ParserImpl<'a> {
     }
 
     /// Section 16.2.2 Import Declaration
-    pub(crate) fn parse_import_declaration(&mut self, span: u32) -> Statement<'a> {
+    pub(crate) fn parse_import_declaration(
+        &mut self,
+        span: u32,
+        should_record_module_record: bool,
+    ) -> Statement<'a> {
         let token_after_import = self.cur_token();
         let mut identifier_after_import: Option<BindingIdentifier<'_>> =
             if self.cur_kind().is_binding_identifier() {
@@ -186,16 +190,20 @@ impl<'a> ParserImpl<'a> {
         self.asi();
         let span = self.end_span(span);
 
-        self.ast
-            .module_declaration_import_declaration(
-                span,
-                specifiers,
-                source,
-                phase,
-                with_clause,
-                import_kind,
-            )
-            .into()
+        let import_decl = self.ast.alloc_import_declaration(
+            span,
+            specifiers,
+            source,
+            phase,
+            with_clause,
+            import_kind,
+        );
+
+        if should_record_module_record {
+            self.module_record_builder.visit_import_declaration(&import_decl);
+        }
+
+        Statement::ImportDeclaration(import_decl)
     }
 
     // Full Syntax: <https://developer.mozilla.org/en-US/docs/Web/JavaScript/Reference/Statements/import#syntax>
@@ -322,21 +330,29 @@ impl<'a> ParserImpl<'a> {
     pub(crate) fn parse_ts_export_assignment_declaration(
         &mut self,
         start_span: u32,
+        stmt_ctx: StatementContext,
     ) -> Box<'a, TSExportAssignment<'a>> {
         self.expect(Kind::Eq);
         let expression = self.parse_assignment_expression_or_higher();
         self.asi();
+        if stmt_ctx.is_top_level() {
+            self.module_record_builder.found_ts_export();
+        }
         self.ast.alloc_ts_export_assignment(self.end_span(start_span), expression)
     }
 
     pub(crate) fn parse_ts_export_namespace(
         &mut self,
         start_span: u32,
+        stmt_ctx: StatementContext,
     ) -> Box<'a, TSNamespaceExportDeclaration<'a>> {
         self.expect(Kind::As);
         self.expect(Kind::Namespace);
         let id = self.parse_identifier_name();
         self.asi();
+        if stmt_ctx.is_top_level() {
+            self.module_record_builder.found_ts_export();
+        }
         self.ast.alloc_ts_namespace_export_declaration(self.end_span(start_span), id)
     }
 
@@ -345,23 +361,31 @@ impl<'a> ParserImpl<'a> {
         &mut self,
         span: u32,
         mut decorators: Vec<'a, Decorator<'a>>,
+        stmt_ctx: StatementContext,
     ) -> Statement<'a> {
         self.bump_any(); // bump `export`
         let decl = match self.cur_kind() {
             // `export import A = B`
             Kind::Import => {
                 let import_span = self.start_span();
                 self.bump_any();
-                let stmt = self.parse_import_declaration(import_span);
+                // Pass `should_record_module_record: false` to prevent an `import` module record
+                // being created. It's an export not an import.
+                let stmt = self.parse_import_declaration(import_span, false);
                 if stmt.is_declaration() {
-                    self.ast.module_declaration_export_named_declaration(
+                    let export_named_decl = self.ast.alloc_export_named_declaration(
                         self.end_span(span),
                         Some(stmt.into_declaration()),
                         self.ast.vec(),
                         None,
                         ImportOrExportKind::Value,
                         NONE,
-                    )
+                    );
+                    if stmt_ctx.is_top_level() {
+                        self.module_record_builder
+                            .visit_export_named_declaration(&export_named_decl);
+                    }
+                    ModuleDeclaration::ExportNamedDeclaration(export_named_decl)
                 } else {
                     return self.fatal_error(diagnostics::unexpected_export(stmt.span()));
                 }
@@ -378,52 +402,56 @@ impl<'a> ParserImpl<'a> {
                 let modifiers = self.parse_modifiers(false, false);
                 let class_decl = self.parse_class_declaration(class_span, &modifiers, decorators);
                 let decl = Declaration::ClassDeclaration(class_decl);
-                self.ast.module_declaration_export_named_declaration(
+                let export_named_decl = self.ast.alloc_export_named_declaration(
                     self.end_span(span),
                     Some(decl),
                     self.ast.vec(),
                     None,
                     ImportOrExportKind::Value,
                     NONE,
-                )
+                );
+                if stmt_ctx.is_top_level() {
+                    self.module_record_builder.visit_export_named_declaration(&export_named_decl);
+                }
+                ModuleDeclaration::ExportNamedDeclaration(export_named_decl)
             }
             Kind::Eq if self.is_ts => ModuleDeclaration::TSExportAssignment(
-                self.parse_ts_export_assignment_declaration(span),
+                self.parse_ts_export_assignment_declaration(span, stmt_ctx),
             ),
             Kind::As if self.is_ts && self.lexer.peek_token().kind() == Kind::Namespace => {
                 // `export as namespace ...`
                 ModuleDeclaration::TSNamespaceExportDeclaration(
-                    self.parse_ts_export_namespace(span),
+                    self.parse_ts_export_namespace(span, stmt_ctx),
                 )
             }
             Kind::Default => ModuleDeclaration::ExportDefaultDeclaration(
-                self.parse_export_default_declaration(span, decorators),
+                self.parse_export_default_declaration(span, decorators, stmt_ctx),
+            ),
+            Kind::Star => ModuleDeclaration::ExportAllDeclaration(
+                self.parse_export_all_declaration(span, stmt_ctx),
+            ),
+            Kind::LCurly => ModuleDeclaration::ExportNamedDeclaration(
+                self.parse_export_named_specifiers(span, stmt_ctx),
             ),
-            Kind::Star => {
-                ModuleDeclaration::ExportAllDeclaration(self.parse_export_all_declaration(span))
-            }
-            Kind::LCurly => {
-                ModuleDeclaration::ExportNamedDeclaration(self.parse_export_named_specifiers(span))
-            }
             Kind::Type if self.is_ts => {
                 let next_kind = self.lexer.peek_token().kind();
 
                 match next_kind {
                     // `export type { ...`
                     Kind::LCurly => ModuleDeclaration::ExportNamedDeclaration(
-                        self.parse_export_named_specifiers(span),
+                        self.parse_export_named_specifiers(span, stmt_ctx),
                     ),
                     // `export type * as ...`
                     Kind::Star => ModuleDeclaration::ExportAllDeclaration(
-                        self.parse_export_all_declaration(span),
+                        self.parse_export_all_declaration(span, stmt_ctx),
                     ),
                     _ => ModuleDeclaration::ExportNamedDeclaration(
-                        self.parse_export_named_declaration(span, decorators),
+                        self.parse_export_named_declaration(span, decorators, stmt_ctx),
                     ),
                 }
             }
             _ => ModuleDeclaration::ExportNamedDeclaration(
-                self.parse_export_named_declaration(span, decorators),
+                self.parse_export_named_declaration(span, decorators, stmt_ctx),
             ),
         };
         Statement::from(decl)
@@ -440,7 +468,11 @@ impl<'a> ParserImpl<'a> {
     // ExportSpecifier :
     //   ModuleExportName
     //   ModuleExportName as ModuleExportName
-    fn parse_export_named_specifiers(&mut self, span: u32) -> Box<'a, ExportNamedDeclaration<'a>> {
+    fn parse_export_named_specifiers(
+        &mut self,
+        span: u32,
+        stmt_ctx: StatementContext,
+    ) -> Box<'a, ExportNamedDeclaration<'a>> {
         let export_kind = self.parse_import_or_export_kind();
         self.expect(Kind::LCurly);
         let (mut specifiers, _) = self.context(Context::empty(), self.ctx, |p| {
@@ -496,21 +528,26 @@ impl<'a> ParserImpl<'a> {
 
         self.asi();
         let span = self.end_span(span);
-        self.ast.alloc_export_named_declaration(
+        let export_named_decl = self.ast.alloc_export_named_declaration(
             span,
             None,
             specifiers,
             source,
             export_kind,
             with_clause,
-        )
+        );
+        if stmt_ctx.is_top_level() {
+            self.module_record_builder.visit_export_named_declaration(&export_named_decl);
+        }
+        export_named_decl
     }
 
     // export Declaration
     fn parse_export_named_declaration(
         &mut self,
         span: u32,
         decorators: Vec<'a, Decorator<'a>>,
+        stmt_ctx: StatementContext,
     ) -> Box<'a, ExportNamedDeclaration<'a>> {
         let decl_span = self.start_span();
         let reserved_ctx = self.ctx;
@@ -525,14 +562,18 @@ impl<'a> ParserImpl<'a> {
             ImportOrExportKind::Value
         };
         self.ctx = reserved_ctx;
-        self.ast.alloc_export_named_declaration(
+        let export_named_decl = self.ast.alloc_export_named_declaration(
             self.end_span(span),
             Some(declaration),
             self.ast.vec(),
             None,
             export_kind,
             NONE,
-        )
+        );
+        if stmt_ctx.is_top_level() {
+            self.module_record_builder.visit_export_named_declaration(&export_named_decl);
+        }
+        export_named_decl
     }
 
     // export default HoistableDeclaration[~Yield, +Await, +Default]
@@ -542,12 +583,18 @@ impl<'a> ParserImpl<'a> {
         &mut self,
         span: u32,
         decorators: Vec<'a, Decorator<'a>>,
+        stmt_ctx: StatementContext,
     ) -> Box<'a, ExportDefaultDeclaration<'a>> {
         let exported = self.parse_keyword_identifier(Kind::Default);
         let declaration = self.parse_export_default_declaration_kind(decorators);
         let exported = ModuleExportName::IdentifierName(exported);
         let span = self.end_span(span);
-        self.ast.alloc_export_default_declaration(span, exported, declaration)
+        let export_default_decl =
+            self.ast.alloc_export_default_declaration(span, exported, declaration);
+        if stmt_ctx.is_top_level() {
+            self.module_record_builder.visit_export_default_declaration(&export_default_decl);
+        }
+        export_default_decl
     }
 
     fn parse_export_default_declaration_kind(
@@ -672,7 +719,11 @@ impl<'a> ParserImpl<'a> {
     //   *
     //   * as ModuleExportName
     //   NamedExports
-    fn parse_export_all_declaration(&mut self, span: u32) -> Box<'a, ExportAllDeclaration<'a>> {
+    fn parse_export_all_declaration(
+        &mut self,
+        span: u32,
+        stmt_ctx: StatementContext,
+    ) -> Box<'a, ExportAllDeclaration<'a>> {
         let export_kind = self.parse_import_or_export_kind();
         self.bump_any(); // bump `star`
         let exported = self.eat(Kind::As).then(|| self.parse_module_export_name());
@@ -681,7 +732,12 @@ impl<'a> ParserImpl<'a> {
         let with_clause = self.parse_import_attributes();
         self.asi();
         let span = self.end_span(span);
-        self.ast.alloc_export_all_declaration(span, exported, source, with_clause, export_kind)
+        let export_all_decl =
+            self.ast.alloc_export_all_declaration(span, exported, source, with_clause, export_kind);
+        if stmt_ctx.is_top_level() {
+            self.module_record_builder.visit_export_all_declaration(&export_all_decl);
+        }
+        export_all_decl
     }
 
     // ImportSpecifier :

crates/oxc_parser/src/js/statement.rs

@@ -49,12 +49,6 @@ impl<'a> ParserImpl<'a> {
             }
             let stmt = self.parse_statement_list_item(stmt_ctx);
 
-            if is_top_level {
-                if let Some(module_decl) = stmt.as_module_declaration() {
-                    self.module_record_builder.visit_module_declaration(module_decl);
-                }
-            }
-
             // Section 11.2.1 Directive Prologue
             // The only way to get a correct directive is to parse the statement first and check if it is a string literal.
             // All other method are flawed, see test cases in [babel](https://github.com/babel/babel/blob/v7.26.2/packages/babel-parser/test/fixtures/core/categorized/not-directive/input.js)
@@ -109,7 +103,9 @@ impl<'a> ParserImpl<'a> {
                 &Modifiers::empty(),
                 self.ast.vec(),
             ),
-            Kind::Export => self.parse_export_declaration(self.start_span(), self.ast.vec()),
+            Kind::Export => {
+                self.parse_export_declaration(self.start_span(), self.ast.vec(), stmt_ctx)
+            }
             // [+Return] ReturnStatement[?Yield, ?Await]
             Kind::Return => self.parse_return_statement(),
             Kind::Var => {
@@ -124,7 +120,7 @@ impl<'a> ParserImpl<'a> {
             Kind::At => self.parse_decorated_statement(stmt_ctx),
             Kind::Let if !self.cur_token().escaped() => self.parse_let(stmt_ctx),
             Kind::Async => self.parse_async_statement(self.start_span(), stmt_ctx),
-            Kind::Import => self.parse_import_statement(),
+            Kind::Import => self.parse_import_statement(stmt_ctx),
             Kind::Const => self.parse_const_statement(stmt_ctx),
             Kind::Using if self.is_using_declaration() => self.parse_using_statement(),
             Kind::Await if self.is_using_statement() => self.parse_using_statement(),
@@ -705,7 +701,7 @@ impl<'a> ParserImpl<'a> {
     }
 
     /// Parse import statement or import expression.
-    fn parse_import_statement(&mut self) -> Statement<'a> {
+    fn parse_import_statement(&mut self, stmt_ctx: StatementContext) -> Statement<'a> {
         let checkpoint = self.checkpoint();
         let span = self.start_span();
         self.bump_any();
@@ -714,7 +710,7 @@ impl<'a> ParserImpl<'a> {
             self.rewind(checkpoint);
             self.parse_expression_or_labeled_statement()
         } else {
-            self.parse_import_declaration(span)
+            self.parse_import_declaration(span, stmt_ctx.is_top_level())
         }
     }
 
@@ -740,7 +736,7 @@ impl<'a> ParserImpl<'a> {
         let kind = self.cur_kind();
         if kind == Kind::Export {
             // Export span.start starts after decorators.
-            return self.parse_export_declaration(self.start_span(), decorators);
+            return self.parse_export_declaration(self.start_span(), decorators, stmt_ctx);
         }
         let modifiers = self.parse_modifiers(false, false);
         if self.at(Kind::Class) {