perf(semantic): calculate number of nodes, scopes, symbols, references before visiting AST

crates/oxc_semantic/src/builder.rs

@@ -6,6 +6,7 @@ use std::{
     sync::Arc,
 };
 
+use oxc_ast::visit::walk::{walk_ts_enum_member_name, walk_ts_module_declaration_name};
 #[allow(clippy::wildcard_imports)]
 use oxc_ast::{ast::*, AstKind, Trivias, Visit};
 use oxc_cfg::{
@@ -100,6 +101,94 @@ pub struct SemanticBuilderReturn<'a> {
     pub errors: Vec<OxcDiagnostic>,
 }
 
+#[derive(Default, Debug)]
+pub struct Collector {
+    node: usize,
+    scope: usize,
+    symbol: usize,
+    reference: usize,
+}
+
+impl<'a> Visit<'a> for Collector {
+    #[inline]
+    fn enter_node(&mut self, _: AstKind<'a>) {
+        self.node += 1;
+    }
+    #[inline]
+    fn enter_scope(&mut self, _: ScopeFlags, _: &Cell<Option<ScopeId>>) {
+        self.scope += 1;
+    }
+
+    #[inline]
+    fn visit_binding_identifier(&mut self, _: &BindingIdentifier<'a>) {
+        self.symbol += 1;
+        self.node += 1;
+    }
+
+    #[inline]
+    fn visit_identifier_reference(&mut self, _: &IdentifierReference<'a>) {
+        self.reference += 1;
+        self.node += 1;
+    }
+
+    #[inline]
+    fn visit_jsx_identifier(&mut self, _: &JSXIdentifier<'a>) {
+        self.reference += 1;
+        self.node += 1;
+    }
+
+    #[inline]
+    fn visit_jsx_member_expression_object(&mut self, it: &JSXMemberExpressionObject<'a>) {
+        self.node += 1;
+        if let JSXMemberExpressionObject::MemberExpression(expr) = &it {
+            self.visit_jsx_member_expression(expr);
+        } else {
+            self.node += 1;
+        }
+    }
+
+    #[inline]
+    fn visit_jsx_element_name(&mut self, it: &JSXElementName<'a>) {
+        self.node += 1;
+        match it {
+            JSXElementName::Identifier(ident) => {
+                self.node += 1;
+                if ident.name.chars().next().is_some_and(char::is_uppercase) {
+                    self.reference += 1;
+                }
+            }
+            JSXElementName::NamespacedName(name) => self.visit_jsx_namespaced_name(name),
+            JSXElementName::MemberExpression(expr) => self.visit_jsx_member_expression(expr),
+        }
+    }
+
+    #[inline]
+    fn visit_jsx_attribute_name(&mut self, it: &JSXAttributeName<'a>) {
+        // NOTE: AstKind doesn't exists!
+        // self.node += 1;
+
+        if let JSXAttributeName::NamespacedName(name) = it {
+            self.visit_jsx_namespaced_name(name);
+        } else {
+            self.node += 1;
+        }
+    }
+
+    #[inline]
+    fn visit_ts_enum_member_name(&mut self, it: &TSEnumMemberName<'a>) {
+        if !it.is_expression() {
+            self.symbol += 1;
+        }
+        walk_ts_enum_member_name(self, it);
+    }
+
+    #[inline]
+    fn visit_ts_module_declaration_name(&mut self, it: &TSModuleDeclarationName<'a>) {
+        self.symbol += 1;
+        walk_ts_module_declaration_name(self, it);
+    }
+}
+
 impl<'a> SemanticBuilder<'a> {
     pub fn new(source_text: &'a str, source_type: SourceType) -> Self {
         let scope = ScopeTree::default();
@@ -192,8 +281,19 @@ impl<'a> SemanticBuilder<'a> {
             let scope_id = self.scope.add_scope(None, AstNodeId::DUMMY, ScopeFlags::Top);
             program.scope_id.set(Some(scope_id));
         } else {
+            let mut collector = Collector::default();
+            collector.visit_program(program);
+            self.nodes.reserve(collector.node);
+            self.scope.reserve(collector.scope);
+            self.symbols.reserve(collector.symbol, collector.reference);
+
             self.visit_program(program);
 
+            debug_assert_eq!(self.nodes.len(), collector.node);
+            debug_assert_eq!(self.scope.len(), collector.scope);
+            debug_assert_eq!(self.symbols.references.len(), collector.reference);
+            debug_assert!(collector.symbol >= self.symbols.len());
+
             // Checking syntax error on module record requires scope information from the previous AST pass
             if self.check_syntax_error {
                 checker::check_module_record(&self);

crates/oxc_semantic/src/node.rs

@@ -171,6 +171,11 @@ impl<'a> AstNodes<'a> {
         self.nodes.push(node);
         ast_node_id
     }
+
+    pub fn reserve(&mut self, additional: usize) {
+        self.nodes.reserve(additional);
+        self.parent_ids.reserve(additional);
+    }
 }
 
 #[derive(Debug)]

crates/oxc_semantic/src/scope.rs

@@ -218,4 +218,12 @@ impl ScopeTree {
     pub fn remove_binding(&mut self, scope_id: ScopeId, name: &CompactStr) {
         self.bindings[scope_id].shift_remove(name);
     }
+
+    pub fn reserve(&mut self, additional: usize) {
+        self.parent_ids.reserve(additional);
+        self.child_ids.reserve(additional);
+        self.flags.reserve(additional);
+        self.bindings.reserve(additional);
+        self.node_ids.reserve(additional);
+    }
 }

crates/oxc_semantic/src/symbol.rs

@@ -193,4 +193,15 @@ impl SymbolTable {
             _ => false,
         }
     }
+
+    pub fn reserve(&mut self, additional: usize, reference: usize) {
+        self.spans.reserve(additional);
+        self.names.reserve(additional);
+        self.flags.reserve(additional);
+        self.scope_ids.reserve(additional);
+        self.declarations.reserve(additional);
+        self.resolved_references.reserve(additional);
+        self.redeclare_variables.reserve(additional);
+        self.references.reserve(reference);
+    }
 }