refactor(semantic,linter): assert that Program is always the first node (#12123)

ulrichstark · web-flow · commit 228cff54189d · 2025-07-08T11:35:47.000+01:00
Follow on to #12087 > We might be able to improve perf of AstNodes::ancestor_ids a little by using a custom iterator, instead of std::iter::successors. > Also, node_id == NodeId::DUMMY would be a cheaper check than parent_id == node_id. The NodeId of Program in practice is always NodeId::DUMMY (0), but we don't have a static guarantee of that at present. But we could get that guarantee by adding assert!(self.parent_ids.is_empty()) to the top of AstNodes::add_program_node (or maybe just debug_assert!). > @overlookmotel in #12087 (comment) My changes: - added constant NodeId::ROOT (same value as NodeId::DUMMY, but leads to more intuitive code) - added asserts in `add_program_node` that nodes are empty and node kind that's being added is Program - ancestor_kinds and ancestors now build upon ancestor_ids, making std::iter::successors implementation obsolete - make AstNodeParentIter simpler by iterating over node ids and by making use of previous assertions of `add_program_node` - remove root, root_node and root_node_mut functions and replace their use with either new program function or NodeId::ROOT Results in general much less code, especially less boilerplate code in the linter around getting the program node. Also might hopefully yield some small benchmark improvements. Edit: Only very minor if not zero improvements on the linter benchmarks :( I also noticed while working on this that the ancestor iterator methods where documented as "The first node produced by this iterator is the first parent of the node pointed to by node_id". That statement was incorrect because all the iterators didn't yield the parent of node_id, but node_id itself as first item. I updated the documentation of these methods to correctly reflect their implementation.
diff --git a/crates/oxc_linter/src/rules/eslint/no_eval.rs b/crates/oxc_linter/src/rules/eslint/no_eval.rs
@@ -214,11 +214,8 @@ impl Rule for NoEval {
                         return;
                     }
 
-                    let root = ctx.nodes().get_node(ctx.nodes().root().unwrap());
-                    let program = root.kind().as_program().unwrap();
-
                     let is_valid = if scope_flags.is_top() {
-                        program.source_type.is_script()
+                        ctx.nodes().program().unwrap().source_type.is_script()
                     } else {
                         let node = ctx.nodes().get_node(ctx.scoping().get_node_id(scope_id));
                         ast_util::is_default_this_binding(ctx, node, true)
diff --git a/crates/oxc_linter/src/rules/eslint/no_unreachable.rs b/crates/oxc_linter/src/rules/eslint/no_unreachable.rs
@@ -8,6 +8,7 @@ use oxc_cfg::{
 };
 use oxc_diagnostics::OxcDiagnostic;
 use oxc_macros::declare_oxc_lint;
+use oxc_semantic::NodeId;
 use oxc_span::{GetSpan, Span};
 
 use crate::{context::LintContext, rule::Rule};
@@ -54,7 +55,7 @@ declare_oxc_lint!(
 impl Rule for NoUnreachable {
     fn run_once(&self, ctx: &LintContext) {
         let nodes = ctx.nodes();
-        let Some(root) = nodes.root_node() else { return };
+        let root = nodes.get_node(NodeId::ROOT);
         let cfg = ctx.cfg();
         let graph = cfg.graph();
 
diff --git a/crates/oxc_linter/src/rules/eslint/sort_imports.rs b/crates/oxc_linter/src/rules/eslint/sort_imports.rs
@@ -6,10 +6,7 @@ use std::{
 
 use cow_utils::CowUtils;
 use itertools::Itertools;
-use oxc_ast::{
-    AstKind,
-    ast::{ImportDeclaration, ImportDeclarationSpecifier, Statement},
-};
+use oxc_ast::ast::{ImportDeclaration, ImportDeclarationSpecifier, Statement};
 use oxc_diagnostics::OxcDiagnostic;
 use oxc_macros::declare_oxc_lint;
 use oxc_span::Span;
@@ -137,10 +134,7 @@ impl Rule for SortImports {
     }
 
     fn run_once(&self, ctx: &LintContext) {
-        let Some(root) = ctx.nodes().root_node() else {
-            return;
-        };
-        let AstKind::Program(program) = root.kind() else { unreachable!() };
+        let program = ctx.nodes().program().unwrap();
 
         let mut import_declarations = vec![];
 
diff --git a/crates/oxc_linter/src/rules/import/exports_last.rs b/crates/oxc_linter/src/rules/import/exports_last.rs
@@ -1,8 +1,5 @@
 use itertools::Itertools;
-use oxc_ast::{
-    AstKind,
-    ast::{ModuleDeclaration, Statement},
-};
+use oxc_ast::ast::{ModuleDeclaration, Statement};
 use oxc_diagnostics::OxcDiagnostic;
 use oxc_macros::declare_oxc_lint;
 use oxc_span::{GetSpan, Span};
@@ -54,19 +51,15 @@ declare_oxc_lint!(
 impl Rule for ExportsLast {
     fn run_once(&self, ctx: &LintContext<'_>) {
         // find last non export declaration index
-        let Some(root) = ctx.nodes().root_node() else {
-            return;
-        };
-        if let AstKind::Program(program) = root.kind() {
-            let body = &program.body;
-            let find_res =
-                body.iter().rev().find_position(|statement| !is_exports_declaration(statement));
-            if let Some((index, _)) = find_res {
-                let end = body.len() - index;
-                for statement in &body[0..end] {
-                    if is_exports_declaration(statement) {
-                        ctx.diagnostic(exports_last_diagnostic(statement.span()));
-                    }
+        let program = ctx.nodes().program().unwrap();
+        let body = &program.body;
+        let find_res =
+            body.iter().rev().find_position(|statement| !is_exports_declaration(statement));
+        if let Some((index, _)) = find_res {
+            let end = body.len() - index;
+            for statement in &body[0..end] {
+                if is_exports_declaration(statement) {
+                    ctx.diagnostic(exports_last_diagnostic(statement.span()));
                 }
             }
         }
diff --git a/crates/oxc_linter/src/rules/import/first.rs b/crates/oxc_linter/src/rules/import/first.rs
@@ -1,9 +1,6 @@
 use std::convert::From;
 
-use oxc_ast::{
-    AstKind,
-    ast::{Statement, TSModuleReference},
-};
+use oxc_ast::ast::{Statement, TSModuleReference};
 use oxc_diagnostics::OxcDiagnostic;
 use oxc_macros::declare_oxc_lint;
 use oxc_span::Span;
@@ -113,10 +110,7 @@ impl Rule for First {
         let mut non_import_count = 0;
         let mut any_relative = false;
 
-        let Some(root) = ctx.nodes().root_node() else {
-            return;
-        };
-        let AstKind::Program(program) = root.kind() else { unreachable!() };
+        let program = ctx.nodes().program().unwrap();
 
         for statement in &program.body {
             match statement {
diff --git a/crates/oxc_linter/src/rules/jest/no_duplicate_hooks.rs b/crates/oxc_linter/src/rules/jest/no_duplicate_hooks.rs
@@ -104,18 +104,15 @@ declare_oxc_lint!(
 
 impl Rule for NoDuplicateHooks {
     fn run_once(&self, ctx: &LintContext) {
-        let Some(root_node) = ctx.nodes().root_node() else {
-            return;
-        };
         let mut hook_contexts: FxHashMap<NodeId, Vec<FxHashMap<String, i32>>> =
             FxHashMap::default();
-        hook_contexts.insert(root_node.id(), Vec::new());
+        hook_contexts.insert(NodeId::ROOT, Vec::new());
 
         let mut possibles_jest_nodes = collect_possible_jest_call_node(ctx);
         possibles_jest_nodes.sort_by_key(|n| n.node.id());
 
         for possible_jest_node in possibles_jest_nodes {
-            Self::run(&possible_jest_node, root_node.id(), &mut hook_contexts, ctx);
+            Self::run(&possible_jest_node, NodeId::ROOT, &mut hook_contexts, ctx);
         }
     }
 }
diff --git a/crates/oxc_linter/src/rules/nextjs/no_async_client_component.rs b/crates/oxc_linter/src/rules/nextjs/no_async_client_component.rs
@@ -90,10 +90,7 @@ declare_oxc_lint!(
 
 impl Rule for NoAsyncClientComponent {
     fn run_once(&self, ctx: &LintContext) {
-        let Some(root) = ctx.nodes().root_node() else {
-            return;
-        };
-        let AstKind::Program(program) = root.kind() else { unreachable!() };
+        let program = ctx.nodes().program().unwrap();
 
         if program.directives.iter().any(|directive| directive.directive.as_str() == "use client") {
             for node in &program.body {
diff --git a/crates/oxc_linter/src/rules/typescript/consistent_type_imports.rs b/crates/oxc_linter/src/rules/typescript/consistent_type_imports.rs
@@ -576,8 +576,7 @@ fn get_type_only_named_import<'a>(
     ctx: &LintContext<'a>,
     source: &str,
 ) -> Option<&'a ImportDeclaration<'a>> {
-    let root = ctx.nodes().root_node()?;
-    let program = root.kind().as_program()?;
+    let program = ctx.nodes().program().unwrap();
 
     for stmt in &program.body {
         let Statement::ImportDeclaration(import_decl) = stmt else {
diff --git a/crates/oxc_linter/src/rules/typescript/triple_slash_reference.rs b/crates/oxc_linter/src/rules/typescript/triple_slash_reference.rs
@@ -1,7 +1,4 @@
-use oxc_ast::{
-    AstKind,
-    ast::{Statement, TSModuleReference},
-};
+use oxc_ast::ast::{Statement, TSModuleReference};
 use oxc_diagnostics::OxcDiagnostic;
 use oxc_macros::declare_oxc_lint;
 use oxc_span::{GetSpan, Span};
@@ -109,10 +106,7 @@ impl Rule for TripleSlashReference {
     }
 
     fn run_once(&self, ctx: &LintContext) {
-        let Some(root) = ctx.nodes().root_node() else {
-            return;
-        };
-        let AstKind::Program(program) = root.kind() else { unreachable!() };
+        let program = ctx.nodes().program().unwrap();
 
         // We don't need to iterate over all comments since Triple-slash directives are only valid at the top of their containing file.
         // We are trying to get the first statement start potioin, falling back to the program end if statement does not exist
diff --git a/crates/oxc_linter/src/rules/unicorn/no_empty_file.rs b/crates/oxc_linter/src/rules/unicorn/no_empty_file.rs
@@ -1,4 +1,3 @@
-use oxc_ast::AstKind;
 use oxc_diagnostics::OxcDiagnostic;
 use oxc_macros::declare_oxc_lint;
 use oxc_span::Span;
@@ -44,11 +43,7 @@ declare_oxc_lint!(
 
 impl Rule for NoEmptyFile {
     fn run_once(&self, ctx: &LintContext) {
-        let Some(root) = ctx.nodes().root_node() else {
-            return;
-        };
-
-        let AstKind::Program(program) = root.kind() else { unreachable!() };
+        let program = ctx.nodes().program().unwrap();
         if program.body.iter().any(|node| !is_empty_stmt(node)) {
             return;
         }
diff --git a/crates/oxc_linter/src/rules/unicorn/no_process_exit.rs b/crates/oxc_linter/src/rules/unicorn/no_process_exit.rs
@@ -63,11 +63,7 @@ impl Rule for NoProcessExit {
 }
 
 fn has_hashbang(ctx: &LintContext) -> bool {
-    let Some(root) = ctx.nodes().root_node() else {
-        return false;
-    };
-    let AstKind::Program(program) = root.kind() else { unreachable!() };
-    program.hashbang.is_some()
+    ctx.nodes().program().unwrap().hashbang.is_some()
 }
 
 fn is_inside_process_event_handler(ctx: &LintContext, node: &AstNode) -> bool {
diff --git a/crates/oxc_semantic/src/node.rs b/crates/oxc_semantic/src/node.rs
@@ -1,7 +1,7 @@
 use std::iter::FusedIterator;
 
 use oxc_allocator::{Address, GetAddress};
-use oxc_ast::AstKind;
+use oxc_ast::{AstKind, ast::Program};
 use oxc_cfg::BlockNodeId;
 use oxc_index::IndexVec;
 use oxc_span::{GetSpan, Span};
@@ -98,10 +98,7 @@ impl GetAddress for AstNode<'_> {
 /// Untyped AST nodes flattened into an vec
 #[derive(Debug, Default)]
 pub struct AstNodes<'a> {
-    /// The root node should always point to a `Program`, which is the real
-    /// root of the tree. It isn't possible to statically check for this, so
-    /// users should beware.
-    root: Option<NodeId>,
+    program: Option<&'a Program<'a>>,
     nodes: IndexVec<NodeId, AstNode<'a>>,
     /// `node` -> `parent`
     parent_ids: IndexVec<NodeId, NodeId>,
@@ -125,25 +122,34 @@ impl<'a> AstNodes<'a> {
         self.nodes.is_empty()
     }
 
-    /// Walk up the AST, iterating over each parent [`AstNode`].
+    /// Walk up the AST, iterating over each parent [`NodeId`].
     ///
-    /// The first node produced by this iterator is the first parent of the node
-    /// pointed to by `node_id`. The last node will usually be a `Program`.
+    /// The first node produced by this iterator is the node pointed to by `node_id`.
+    /// The last node will always be [`AstKind::Program`].
     #[inline]
-    pub fn ancestors(&self, node_id: NodeId) -> impl Iterator<Item = &AstNode<'a>> + Clone + '_ {
-        AstNodeParentIter { current_node_id: Some(node_id), nodes: self }
+    pub fn ancestor_ids(&self, node_id: NodeId) -> impl Iterator<Item = NodeId> + Clone + '_ {
+        AstNodeIdParentIter { next_node_id: Some(node_id), nodes: self }
     }
 
     /// Walk up the AST, iterating over each parent [`AstKind`].
     ///
-    /// The first node produced by this iterator is the first parent of the node
-    /// pointed to by `node_id`. The last node will is a [`AstKind::Program`].
+    /// The first node produced by this iterator is the node pointed to by `node_id`.
+    /// The last node will always be [`AstKind::Program`].
     #[inline]
     pub fn ancestor_kinds(
         &self,
         node_id: NodeId,
     ) -> impl Iterator<Item = AstKind<'a>> + Clone + '_ {
-        self.ancestors(node_id).map(AstNode::kind)
+        self.ancestor_ids(node_id).map(|id| self.kind(id))
+    }
+
+    /// Walk up the AST, iterating over each parent [`AstNode`].
+    ///
+    /// The first node produced by this iterator is the node pointed to by `node_id`.
+    /// The last node will always be [`AstKind::Program`].
+    #[inline]
+    pub fn ancestors(&self, node_id: NodeId) -> impl Iterator<Item = &AstNode<'a>> + Clone + '_ {
+        self.ancestor_ids(node_id).map(|id| self.get_node(id))
     }
 
     /// Access the underlying struct from [`oxc_ast`].
@@ -178,50 +184,13 @@ impl<'a> AstNodes<'a> {
         &mut self.nodes[node_id]
     }
 
-    /// Get the root [`NodeId`]. This always points to a [`Program`] node.
+    /// Get the [`Program`] that's also the root of the AST.
     ///
     /// Returns [`None`] if root node isn't set. This will never happen if you
     /// are obtaining an [`AstNodes`] that has already been constructed.
-    ///
-    /// [`Program`]: oxc_ast::ast::Program
     #[inline]
-    pub fn root(&self) -> Option<NodeId> {
-        self.root
-    }
-
-    /// Get the root node as immutable reference, It is always guaranteed to be a [`Program`].
-    ///
-    /// Returns [`None`] if root node isn't set. This will never happen if you
-    /// are obtaining an [`AstNodes`] that has already been constructed.
-    ///
-    /// [`Program`]: oxc_ast::ast::Program
-    #[inline]
-    pub fn root_node(&self) -> Option<&AstNode<'a>> {
-        self.root().map(|id| self.get_node(id))
-    }
-
-    /// Get the root node as mutable reference, It is always guaranteed to be a [`Program`].
-    ///
-    /// Returns [`None`] if root node isn't set. This will never happen if you
-    /// are obtaining an [`AstNodes`] that has already been constructed.
-    ///
-    /// [`Program`]: oxc_ast::ast::Program
-    #[inline]
-    pub fn root_node_mut(&mut self) -> Option<&mut AstNode<'a>> {
-        self.root().map(|id| self.get_node_mut(id))
-    }
-
-    /// Walk up the AST, iterating over each parent [`NodeId`].
-    ///
-    /// The first node produced by this iterator is the first parent of the node
-    /// pointed to by `node_id`. The last node will always be a [`Program`].
-    ///
-    /// [`Program`]: oxc_ast::ast::Program
-    pub fn ancestor_ids(&self, node_id: NodeId) -> impl Iterator<Item = NodeId> + '_ {
-        std::iter::successors(Some(node_id), |&node_id| {
-            let parent_id = self.parent_ids[node_id];
-            if parent_id == node_id { None } else { Some(parent_id) }
-        })
+    pub fn program(&self) -> Option<&'a Program<'a>> {
+        self.program
     }
 
     /// Create and add an [`AstNode`] to the [`AstNodes`] tree and get its [`NodeId`].
@@ -245,16 +214,23 @@ impl<'a> AstNodes<'a> {
     }
 
     /// Create and add an [`AstNode`] to the [`AstNodes`] tree and get its [`NodeId`].
+    ///
+    /// # Panics
+    ///
+    /// Panics if this is not the first node being added to the AST
     pub fn add_program_node(
         &mut self,
         kind: AstKind<'a>,
         scope_id: ScopeId,
         cfg_id: BlockNodeId,
         flags: NodeFlags,
     ) -> NodeId {
-        let node_id = self.parent_ids.push(NodeId::DUMMY);
-        self.parent_ids[node_id] = node_id;
-        self.root = Some(node_id);
+        assert!(self.parent_ids.is_empty(), "Program node must be the first node in the AST.");
+        let AstKind::Program(program) = kind else {
+            panic!("Program node must be of kind `AstKind::Program`.");
+        };
+        self.program = Some(program);
+        let node_id = self.parent_ids.push(NodeId::ROOT);
         let node = AstNode::new(kind, scope_id, cfg_id, flags, node_id);
         self.nodes.push(node);
         node_id
@@ -277,23 +253,23 @@ impl<'a, 'n> IntoIterator for &'n AstNodes<'a> {
 }
 
 #[derive(Debug, Clone)]
-pub struct AstNodeParentIter<'s, 'a> {
-    current_node_id: Option<NodeId>,
+pub struct AstNodeIdParentIter<'s, 'a> {
+    next_node_id: Option<NodeId>,
     nodes: &'s AstNodes<'a>,
 }
 
-impl<'s, 'a> Iterator for AstNodeParentIter<'s, 'a> {
-    type Item = &'s AstNode<'a>;
+impl Iterator for AstNodeIdParentIter<'_, '_> {
+    type Item = NodeId;
 
     fn next(&mut self) -> Option<Self::Item> {
-        if let Some(node_id) = self.current_node_id {
-            let parent_id = self.nodes.parent_ids[node_id];
-            self.current_node_id = if parent_id == node_id { None } else { Some(parent_id) };
-            Some(self.nodes.get_node(node_id))
+        if let Some(node_id) = self.next_node_id {
+            self.next_node_id =
+                if node_id == NodeId::ROOT { None } else { Some(self.nodes.parent_ids[node_id]) };
+            Some(node_id)
         } else {
             None
         }
     }
 }
 
-impl FusedIterator for AstNodeParentIter<'_, '_> {}
+impl FusedIterator for AstNodeIdParentIter<'_, '_> {}
diff --git a/crates/oxc_syntax/src/node.rs b/crates/oxc_syntax/src/node.rs
@@ -15,6 +15,9 @@ impl NodeId {
     /// This is used for synthetically-created AST nodes, among other things.
     pub const DUMMY: Self = NodeId::new(0);
 
+    /// Node id of the Program node.
+    pub const ROOT: Self = NodeId::new(0);
+
     /// Create `NodeId` from `u32`.
     ///
     /// # Panics
diff --git a/tasks/coverage/src/driver.rs b/tasks/coverage/src/driver.rs

Original file line number	Diff line number	Diff line change
`@@ -104,18 +104,15 @@ declare_oxc_lint!(`
`104`	`104`
`105`	`105`	`impl Rule for NoDuplicateHooks {`
`106`	`106`	`fn run_once(&self, ctx: &LintContext) {`
`107`		`- let Some(root_node) = ctx.nodes().root_node() else {`
`108`		`- return;`
`109`		`- };`
`110`	`107`	`let mut hook_contexts: FxHashMap<NodeId, Vec<FxHashMap<String, i32>>> =`
`111`	`108`	`FxHashMap::default();`
`112`		`- hook_contexts.insert(root_node.id(), Vec::new());`
	`109`	`+ hook_contexts.insert(NodeId::ROOT, Vec::new());`
`113`	`110`
`114`	`111`	`let mut possibles_jest_nodes = collect_possible_jest_call_node(ctx);`
`115`	`112`	`possibles_jest_nodes.sort_by_key(\|n\| n.node.id());`
`116`	`113`
`117`	`114`	`for possible_jest_node in possibles_jest_nodes {`
`118`		`- Self::run(&possible_jest_node, root_node.id(), &mut hook_contexts, ctx);`
	`115`	`+ Self::run(&possible_jest_node, NodeId::ROOT, &mut hook_contexts, ctx);`
`119`	`116`	`}`
`120`	`117`	`}`
`121`	`118`	`}`