feat(semantic): add root node to the AstNodes structure. (#3032)

Adds a way to fetch the root node without iterating over all ancestors
which has a nondeterministic time - best case O(1) worst case O(n!) - It
is only possible to set this field in the semantic builder.

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) {

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<AstNodeId, AstNode<'a>>,
     parent_ids: IndexVec<AstNodeId, Option<AstNodeId>>,
 }
 
+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<Item = &AstNode<'a>> + '_ {
         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