refactor(semantic)!: simplify node creation

crates/oxc_semantic/src/builder.rs

@@ -24,7 +24,7 @@ use crate::{
     jsdoc::JSDocBuilder,
     label::LabelBuilder,
     module_record::ModuleRecordBuilder,
-    node::{AstNode, AstNodeId, AstNodes, NodeFlags},
+    node::{AstNodeId, AstNodes, NodeFlags},
     reference::{Reference, ReferenceFlag, ReferenceId},
     scope::{ScopeFlags, ScopeId, ScopeTree, UnresolvedReferences},
     symbol::{SymbolFlags, SymbolId, SymbolTable},
@@ -229,28 +229,17 @@ impl<'a> SemanticBuilder<'a> {
 
     fn create_ast_node(&mut self, kind: AstKind<'a>) {
         let mut flags = self.current_node_flags;
-
         if self.build_jsdoc && self.jsdoc.retrieve_attached_jsdoc(&kind) {
             flags |= NodeFlags::JSDoc;
         }
 
-        let ast_node = AstNode::new(
+        self.current_node_id = self.nodes.add_node(
             kind,
             self.current_scope_id,
+            self.current_node_id,
             control_flow!(|self, cfg| cfg.current_node_ix),
             flags,
         );
-        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))
-        };
         self.record_ast_node();
     }
 

crates/oxc_semantic/src/node.rs

@@ -22,13 +22,14 @@ pub struct AstNode<'a> {
 }
 
 impl<'a> AstNode<'a> {
-    pub fn new(
+    pub(crate) fn new(
         kind: AstKind<'a>,
         scope_id: ScopeId,
         cfg_id: BasicBlockId,
         flags: NodeFlags,
+        id: AstNodeId,
     ) -> Self {
-        Self { id: AstNodeId::new(0), kind, cfg_id, scope_id, flags }
+        Self { id, kind, scope_id, cfg_id, flags }
     }
 
     pub fn id(&self) -> AstNodeId {
@@ -59,6 +60,9 @@ impl<'a> AstNode<'a> {
 /// 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<AstNodeId>,
     nodes: IndexVec<AstNodeId, AstNode<'a>>,
     parent_ids: IndexVec<AstNodeId, Option<AstNodeId>>,
@@ -116,21 +120,6 @@ impl<'a> AstNodes<'a> {
         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 = Some(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`.
     /// Returns `None` if root node isn't set.
     pub fn root_node(&self) -> Option<&AstNode<'a>> {
@@ -152,11 +141,24 @@ impl<'a> AstNodes<'a> {
         std::iter::successors(Some(ast_node_id), |node_id| parent_ids[*node_id])
     }
 
-    /// Adds an `AstNode` to the `AstNodes` tree and returns its `AstNodeId`.
-    pub fn add_node(&mut self, node: AstNode<'a>, parent_id: Option<AstNodeId>) -> AstNodeId {
-        let mut node = node;
-        let ast_node_id = self.parent_ids.push(parent_id);
-        node.id = ast_node_id;
+    /// Create and add an `AstNode` to the `AstNodes` tree and returns its `AstNodeId`.
+    pub fn add_node(
+        &mut self,
+        kind: AstKind<'a>,
+        scope_id: ScopeId,
+        parent_node_id: AstNodeId,
+        cfg_id: BasicBlockId,
+        flags: NodeFlags,
+    ) -> AstNodeId {
+        let ast_node_id = match kind {
+            AstKind::Program(_) => {
+                let id = self.parent_ids.push(None);
+                self.root = Some(id);
+                id
+            }
+            _ => self.parent_ids.push(Some(parent_node_id)),
+        };
+        let node = AstNode::new(kind, scope_id, cfg_id, flags, ast_node_id);
         self.nodes.push(node);
         ast_node_id
     }