diff --git a/crates/oxc_semantic/src/builder.rs b/crates/oxc_semantic/src/builder.rs index 140fb66d148ff..e53d8447148cd 100644 --- a/crates/oxc_semantic/src/builder.rs +++ b/crates/oxc_semantic/src/builder.rs @@ -197,9 +197,17 @@ impl<'a> SemanticBuilder<'a> { } let ast_node = AstNode::new(kind, self.current_scope_id, self.cfg.current_node_ix, flags); - let parent_node_id = - if matches!(kind, AstKind::Program(_)) { None } else { Some(self.current_node_id) }; - self.current_node_id = self.nodes.add_node(ast_node, parent_node_id); + self.current_node_id = if matches!(kind, AstKind::Program(_)) { + let id = self.nodes.add_node(ast_node, None); + #[allow(unsafe_code)] + // SAFETY: `ast_node` is a `Program` and hence the root of the tree. + unsafe { + self.nodes.set_root(&ast_node); + } + id + } else { + self.nodes.add_node(ast_node, Some(self.current_node_id)) + }; } fn pop_ast_node(&mut self) { diff --git a/crates/oxc_semantic/src/node.rs b/crates/oxc_semantic/src/node.rs index 77987d9fcc363..3c9cac77a23fa 100644 --- a/crates/oxc_semantic/src/node.rs +++ b/crates/oxc_semantic/src/node.rs @@ -54,12 +54,23 @@ impl<'a> AstNode<'a> { } /// Untyped AST nodes flattened into an vec -#[derive(Debug, Default)] +#[derive(Debug)] pub struct AstNodes<'a> { + root: AstNodeId, nodes: IndexVec>, parent_ids: IndexVec>, } +impl<'a> Default for AstNodes<'a> { + fn default() -> Self { + Self { + root: AstNodeId::new(0), + nodes: IndexVec::default(), + parent_ids: IndexVec::default(), + } + } +} + impl<'a> AstNodes<'a> { pub fn iter(&self) -> impl Iterator> + '_ { self.nodes.iter() @@ -98,6 +109,36 @@ impl<'a> AstNodes<'a> { &mut self.nodes[ast_node_id] } + /// Get the root `AstNodeId`, It is always pointing to a `Program`. + pub fn root(&self) -> AstNodeId { + self.root + } + + /// Set the root node, + /// SAFETY: + /// The root `AstNode` should always point to a `Program` and this should be the real root of + /// the tree, It isn't possible to statically check for this so user should think about it before + /// using. + #[allow(unsafe_code)] + pub(super) unsafe fn set_root(&mut self, root: &AstNode<'a>) { + match root.kind() { + AstKind::Program(_) => { + self.root = root.id(); + } + _ => unreachable!("Expected a `Program` node as the root of the tree."), + } + } + + /// Get the root node as immutable reference, It is always guaranteed to be a `Program`. + pub fn root_node(&self) -> &AstNode<'a> { + self.get_node(self.root()) + } + + /// Get the root node as mutable reference, It is always guaranteed to be a `Program`. + pub fn root_node_mut(&mut self) -> &mut AstNode<'a> { + self.get_node_mut(self.root()) + } + /// Walk up the AST, iterating over each parent node. /// /// The first node produced by this iterator is the first parent of the node