Skip to main content

qanglang_core/frontend/
ast_node_arena.rs

1use crate::{
2    arena::{Arena, Index},
3    frontend::{node_array_arena::NodeArrayArena, nodes::*},
4};
5
6#[derive(Debug, Default, Clone, Copy, Eq, PartialEq, Hash)]
7pub struct NodeId(Index);
8
9impl NodeId {
10    pub fn new(index: Index) -> Self {
11        Self(index)
12    }
13
14    pub fn get(&self) -> Index {
15        self.0
16    }
17}
18
19#[derive(Debug, Clone, Copy)]
20pub struct TypedNodeRef<T: Clone + Copy> {
21    pub id: NodeId,
22    pub node: T,
23}
24
25impl<T: Clone + Copy> TypedNodeRef<T> {
26    pub fn new(id: NodeId, node: T) -> Self {
27        Self { id, node }
28    }
29}
30
31#[derive(Default, Clone, Debug)]
32pub struct AstNodeArena {
33    nodes: Arena<AstNode>,
34    pub array: NodeArrayArena,
35}
36
37impl AstNodeArena {
38    pub fn new() -> Self {
39        Self {
40            nodes: Arena::new(),
41            array: NodeArrayArena::new(),
42        }
43    }
44
45    pub fn create_node(&mut self, node: AstNode) -> NodeId {
46        NodeId::new(self.nodes.insert(node))
47    }
48
49    pub fn get_node(&self, node_id: NodeId) -> &AstNode {
50        &self.nodes[node_id.get()]
51    }
52
53    pub fn get_decl_node(&self, node_id: NodeId) -> TypedNodeRef<DeclNode> {
54        let node = self.nodes[node_id.get()];
55        TypedNodeRef::new(node_id, node.try_into().unwrap())
56    }
57
58    pub fn get_stmt_node(&self, node_id: NodeId) -> TypedNodeRef<StmtNode> {
59        let node = self.nodes[node_id.get()];
60        TypedNodeRef::new(node_id, node.try_into().unwrap())
61    }
62
63    pub fn get_expr_node(&self, node_id: NodeId) -> TypedNodeRef<ExprNode> {
64        let node = self.nodes[node_id.get()];
65        TypedNodeRef::new(node_id, node.try_into().unwrap())
66    }
67
68    pub fn get_call_operation_node(&self, node_id: NodeId) -> TypedNodeRef<CallOperationNode> {
69        let node = self.nodes[node_id.get()];
70        TypedNodeRef::new(node_id, node.try_into().unwrap())
71    }
72
73    pub fn get_assignment_target_node(
74        &self,
75        node_id: NodeId,
76    ) -> TypedNodeRef<AssignmentTargetNode> {
77        let node = self.nodes[node_id.get()];
78        TypedNodeRef::new(node_id, node.try_into().unwrap())
79    }
80
81    pub fn get_primary_node(&self, node_id: NodeId) -> TypedNodeRef<PrimaryNode> {
82        let node = self.nodes[node_id.get()];
83        TypedNodeRef::new(node_id, node.try_into().unwrap())
84    }
85
86    pub fn get_var_decl_node(&self, node_id: NodeId) -> TypedNodeRef<VariableDeclNode> {
87        let node = self.nodes[node_id.get()];
88
89        match node {
90            AstNode::VariableDecl(decl) => TypedNodeRef::new(node_id, decl),
91            _ => panic!("Expected VariableDeclNode."),
92        }
93    }
94
95    pub fn get_program_node(&self, node_id: NodeId) -> TypedNodeRef<Module> {
96        let node = self.nodes[node_id.get()];
97
98        match node {
99            AstNode::Module(program) => TypedNodeRef::new(node_id, program),
100            _ => panic!("Expected ProgramNode"),
101        }
102    }
103
104    pub fn get_identifier_node(&self, node_id: NodeId) -> TypedNodeRef<IdentifierNode> {
105        let node = self.nodes[node_id.get()];
106        match node {
107            AstNode::Identifier(identifier) => TypedNodeRef::new(node_id, identifier),
108            _ => panic!("Expected IdentifierNode"),
109        }
110    }
111
112    pub fn get_func_expr_node(&self, node_id: NodeId) -> TypedNodeRef<FunctionExprNode> {
113        let node = self.nodes[node_id.get()];
114
115        match node {
116            AstNode::FunctionExpr(func) => TypedNodeRef::new(node_id, func),
117            _ => panic!("Expected FunctionExprNode"),
118        }
119    }
120
121    pub fn get_lambda_expr_node(&self, node_id: NodeId) -> TypedNodeRef<LambdaExprNode> {
122        let node = self.nodes[node_id.get()];
123
124        match node {
125            AstNode::LambdaExpr(lambda) => TypedNodeRef::new(node_id, lambda),
126            _ => panic!("Expected LambdaExprNode"),
127        }
128    }
129
130    pub fn get_block_stmt_node(&self, node_id: NodeId) -> TypedNodeRef<BlockStmtNode> {
131        let node = self.nodes[node_id.get()];
132
133        match node {
134            AstNode::BlockStmt(block) => TypedNodeRef::new(node_id, block),
135            _ => panic!("Expected BlockStmtNode."),
136        }
137    }
138
139    pub fn get_obj_entry_node(&self, node_id: NodeId) -> TypedNodeRef<ObjectEntryNode> {
140        let node = self.nodes[node_id.get()];
141
142        match node {
143            AstNode::ObjectEntry(entry) => TypedNodeRef::new(node_id, entry),
144            _ => panic!("Expected ObjectEntryNode."),
145        }
146    }
147
148    pub fn get_lambda_body_node(&self, node_id: NodeId) -> TypedNodeRef<LambdaBodyNode> {
149        let node = self.nodes[node_id.get()];
150        TypedNodeRef::new(node_id, node.try_into().unwrap())
151    }
152
153    pub fn get_for_initializer_node(&self, node_id: NodeId) -> TypedNodeRef<ForInitializerNode> {
154        let node = self.nodes[node_id.get()];
155        TypedNodeRef::new(node_id, node.try_into().unwrap())
156    }
157
158    pub fn get_class_member_node(&self, node_id: NodeId) -> TypedNodeRef<ClassMemberNode> {
159        let node = self.nodes[node_id.get()];
160        TypedNodeRef::new(node_id, node.try_into().unwrap())
161    }
162
163    pub fn get_when_expr_node(&self, node_id: NodeId) -> TypedNodeRef<WhenExprNode> {
164        let node = self.nodes[node_id.get()];
165
166        match node {
167            AstNode::WhenExpr(when_expr) => TypedNodeRef::new(node_id, when_expr),
168            _ => panic!("Expected WhenExprNode"),
169        }
170    }
171
172    pub fn get_when_branch_node(&self, node_id: NodeId) -> TypedNodeRef<WhenBranchNode> {
173        let node = self.nodes[node_id.get()];
174
175        match node {
176            AstNode::WhenBranch(when_branch) => TypedNodeRef::new(node_id, when_branch),
177            _ => panic!("Expected WhenBranchNode"),
178        }
179    }
180}