1use slotmap::{SlotMap, new_key_type};
8
9new_key_type! {
10 pub struct GreenNodeId;
12}
13
14#[derive(Debug, Clone, PartialEq, Eq)]
17pub enum GreenNode {
18 Token {
20 token_index: usize,
21 width: usize, },
23
24 Internal {
26 kind: SyntaxKind,
27 children: Vec<GreenNodeId>,
28 width: usize, },
30}
31
32impl GreenNode {
33 pub fn width(&self) -> usize {
35 match self {
36 GreenNode::Token { width, .. } => *width,
37 GreenNode::Internal { width, .. } => *width,
38 }
39 }
40
41 pub fn kind(&self) -> Option<SyntaxKind> {
43 match self {
44 GreenNode::Token { .. } => None,
45 GreenNode::Internal { kind, .. } => Some(*kind),
46 }
47 }
48
49 pub fn children(&self) -> Option<&[GreenNodeId]> {
51 match self {
52 GreenNode::Token { .. } => None,
53 GreenNode::Internal { children, .. } => Some(children),
54 }
55 }
56}
57
58#[derive(Debug, Default, Clone)]
60pub struct GreenNodeArena {
61 nodes: SlotMap<GreenNodeId, GreenNode>,
62}
63
64impl GreenNodeArena {
65 pub fn new() -> Self {
67 Self {
68 nodes: SlotMap::with_key(),
69 }
70 }
71
72 pub fn alloc_token(&mut self, token_index: usize, width: usize) -> GreenNodeId {
74 self.nodes.insert(GreenNode::Token { token_index, width })
75 }
76
77 pub fn alloc_internal(&mut self, kind: SyntaxKind, children: Vec<GreenNodeId>) -> GreenNodeId {
79 let width = children.iter().map(|&id| self.nodes[id].width()).sum();
80 self.nodes.insert(GreenNode::Internal {
81 kind,
82 children,
83 width,
84 })
85 }
86
87 pub fn get(&self, id: GreenNodeId) -> &GreenNode {
89 &self.nodes[id]
90 }
91
92 pub fn width(&self, id: GreenNodeId) -> usize {
94 self.nodes[id].width()
95 }
96
97 pub fn kind(&self, id: GreenNodeId) -> Option<SyntaxKind> {
99 self.nodes[id].kind()
100 }
101
102 pub fn children(&self, id: GreenNodeId) -> Option<&[GreenNodeId]> {
104 self.nodes[id].children()
105 }
106
107 pub fn print_tree(
118 &self,
119 node_id: GreenNodeId,
120 tokens: &[crate::compiler::parser::Token],
121 source: &str,
122 indent: usize,
123 ) -> String {
124 use std::fmt::Write;
125 let mut result = String::new();
126 let indent_str = " ".repeat(indent);
127
128 match self.get(node_id) {
129 GreenNode::Token { token_index, .. } => {
130 let token = &tokens[*token_index];
131 let text = token.text(source);
132 let escaped_text = text
134 .replace('\\', "\\\\")
135 .replace('\n', "\\n")
136 .replace('\r', "\\r")
137 .replace('\t', "\\t");
138 let _ = writeln!(result, "{indent_str}{:?} \"{}\"", token.kind, escaped_text);
139 }
140 GreenNode::Internal { kind, children, .. } => {
141 let _ = writeln!(result, "{indent_str}{kind}");
142 for &child in children {
143 result.push_str(&self.print_tree(child, tokens, source, indent + 1));
144 }
145 }
146 }
147 result
148 }
149}
150
151#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
153pub enum SyntaxKind {
154 Program,
156 Statement,
157
158 FunctionDecl,
160 LetDecl,
161 LetRecDecl,
162
163 BinaryExpr,
165 UnaryExpr,
166 ParenExpr,
167 CallExpr,
168 FieldAccess,
169 IndexExpr,
170 AssignExpr,
171 ArrayExpr,
172 MacroExpansion,
173 BracketExpr,
174 EscapeExpr,
175 LambdaExpr,
176 IfExpr,
177 MatchExpr,
178 MatchArm,
179 MatchArmList,
180 MatchPattern,
181 ConstructorPattern,
182 BlockExpr,
183 TupleExpr,
184 RecordExpr,
185
186 IntLiteral,
188 FloatLiteral,
189 StringLiteral,
190 SelfLiteral,
191 NowLiteral,
192 SampleRateLiteral,
193 PlaceHolderLiteral,
194
195 Identifier,
197
198 TypeAnnotation,
200 PrimitiveType,
201 UnitType,
202 FunctionType,
203 TupleType,
204 RecordType,
205 ArrayType,
206 CodeType,
207 UnionType,
208 TypeIdent,
209
210 Pattern,
212 SinglePattern,
213 TuplePattern,
214 RecordPattern,
215
216 IncludeStmt,
218 StageDecl,
219 ModuleDecl,
220 UseStmt,
221 TypeDecl,
222 VariantDef,
223
224 QualifiedPath,
226 VisibilityPub,
227 UseTargetMultiple, UseTargetWildcard, ParamList,
232 ArgList,
233 ExprList,
234 ParamDefault,
235
236 Error, }
239
240impl std::fmt::Display for SyntaxKind {
241 fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
242 match self {
243 SyntaxKind::Program => write!(f, "Program"),
244 SyntaxKind::Statement => write!(f, "Statement"),
245 SyntaxKind::FunctionDecl => write!(f, "FunctionDecl"),
246 SyntaxKind::LetDecl => write!(f, "LetDecl"),
247 SyntaxKind::LetRecDecl => write!(f, "LetRecDecl"),
248 SyntaxKind::BinaryExpr => write!(f, "BinaryExpr"),
249 SyntaxKind::UnaryExpr => write!(f, "UnaryExpr"),
250 SyntaxKind::ParenExpr => write!(f, "ParenExpr"),
251 SyntaxKind::CallExpr => write!(f, "CallExpr"),
252 SyntaxKind::FieldAccess => write!(f, "FieldAccess"),
253 SyntaxKind::IndexExpr => write!(f, "IndexExpr"),
254 SyntaxKind::AssignExpr => write!(f, "AssignExpr"),
255 SyntaxKind::ArrayExpr => write!(f, "ArrayExpr"),
256 SyntaxKind::MacroExpansion => write!(f, "MacroExpansion"),
257 SyntaxKind::BracketExpr => write!(f, "BracketExpr"),
258 SyntaxKind::EscapeExpr => write!(f, "EscapeExpr"),
259 SyntaxKind::LambdaExpr => write!(f, "LambdaExpr"),
260 SyntaxKind::IfExpr => write!(f, "IfExpr"),
261 SyntaxKind::MatchExpr => write!(f, "MatchExpr"),
262 SyntaxKind::MatchArm => write!(f, "MatchArm"),
263 SyntaxKind::MatchArmList => write!(f, "MatchArmList"),
264 SyntaxKind::MatchPattern => write!(f, "MatchPattern"),
265 SyntaxKind::ConstructorPattern => write!(f, "ConstructorPattern"),
266 SyntaxKind::BlockExpr => write!(f, "BlockExpr"),
267 SyntaxKind::TupleExpr => write!(f, "TupleExpr"),
268 SyntaxKind::RecordExpr => write!(f, "RecordExpr"),
269 SyntaxKind::IntLiteral => write!(f, "IntLiteral"),
270 SyntaxKind::FloatLiteral => write!(f, "FloatLiteral"),
271 SyntaxKind::StringLiteral => write!(f, "StringLiteral"),
272 SyntaxKind::SelfLiteral => write!(f, "SelfLiteral"),
273 SyntaxKind::NowLiteral => write!(f, "NowLiteral"),
274 SyntaxKind::SampleRateLiteral => write!(f, "SampleRateLiteral"),
275 SyntaxKind::PlaceHolderLiteral => write!(f, "PlaceHolderLiteral"),
276 SyntaxKind::Identifier => write!(f, "Identifier"),
277 SyntaxKind::TypeAnnotation => write!(f, "TypeAnnotation"),
278 SyntaxKind::PrimitiveType => write!(f, "PrimitiveType"),
279 SyntaxKind::UnitType => write!(f, "UnitType"),
280 SyntaxKind::FunctionType => write!(f, "FunctionType"),
281 SyntaxKind::TupleType => write!(f, "TupleType"),
282 SyntaxKind::RecordType => write!(f, "RecordType"),
283 SyntaxKind::ArrayType => write!(f, "ArrayType"),
284 SyntaxKind::CodeType => write!(f, "CodeType"),
285 SyntaxKind::UnionType => write!(f, "UnionType"),
286 SyntaxKind::TypeIdent => write!(f, "TypeIdent"),
287 SyntaxKind::Pattern => write!(f, "Pattern"),
288 SyntaxKind::SinglePattern => write!(f, "SinglePattern"),
289 SyntaxKind::TuplePattern => write!(f, "TuplePattern"),
290 SyntaxKind::RecordPattern => write!(f, "RecordPattern"),
291 SyntaxKind::IncludeStmt => write!(f, "IncludeStmt"),
292 SyntaxKind::StageDecl => write!(f, "StageDecl"),
293 SyntaxKind::ModuleDecl => write!(f, "ModuleDecl"),
294 SyntaxKind::UseStmt => write!(f, "UseStmt"),
295 SyntaxKind::TypeDecl => write!(f, "TypeDecl"),
296 SyntaxKind::VariantDef => write!(f, "VariantDef"),
297 SyntaxKind::QualifiedPath => write!(f, "QualifiedPath"),
298 SyntaxKind::VisibilityPub => write!(f, "VisibilityPub"),
299 SyntaxKind::UseTargetMultiple => write!(f, "UseTargetMultiple"),
300 SyntaxKind::UseTargetWildcard => write!(f, "UseTargetWildcard"),
301 SyntaxKind::ParamList => write!(f, "ParamList"),
302 SyntaxKind::ArgList => write!(f, "ArgList"),
303 SyntaxKind::ExprList => write!(f, "ExprList"),
304 SyntaxKind::ParamDefault => write!(f, "ParamDefault"),
305 SyntaxKind::Error => write!(f, "Error"),
306 }
307 }
308}
309
310#[derive(Clone, Copy, Debug)]
313pub struct Marker {
314 pub pos: usize,
316}
317
318pub struct GreenTreeBuilder {
320 arena: GreenNodeArena,
321 stack: Vec<(SyntaxKind, Vec<GreenNodeId>)>,
322}
323
324impl GreenTreeBuilder {
325 pub fn new() -> Self {
326 Self {
327 arena: GreenNodeArena::new(),
328 stack: Vec::new(),
329 }
330 }
331
332 pub fn arena(&self) -> &GreenNodeArena {
334 &self.arena
335 }
336
337 pub fn into_arena(self) -> GreenNodeArena {
339 self.arena
340 }
341
342 pub fn marker(&self) -> Marker {
346 let pos = self
347 .stack
348 .last()
349 .map(|(_, children)| children.len())
350 .unwrap_or(0);
351 Marker { pos }
352 }
353
354 pub fn start_node(&mut self, kind: SyntaxKind) {
356 self.stack.push((kind, Vec::new()));
357 }
358
359 pub fn start_node_at(&mut self, marker: Marker, kind: SyntaxKind) {
363 if let Some((_, children)) = self.stack.last_mut() {
364 let wrapped_children: Vec<_> = children.drain(marker.pos..).collect();
366 self.stack.push((kind, wrapped_children));
368 } else {
369 self.stack.push((kind, Vec::new()));
371 }
372 }
373
374 pub fn add_token(&mut self, token_index: usize, width: usize) {
376 let token_id = self.arena.alloc_token(token_index, width);
377 if let Some((_, children)) = self.stack.last_mut() {
378 children.push(token_id);
379 }
380 }
381
382 pub fn finish_node(&mut self) -> Option<GreenNodeId> {
384 if let Some((kind, children)) = self.stack.pop() {
385 let node_id = self.arena.alloc_internal(kind, children);
386
387 if let Some((_, parent_children)) = self.stack.last_mut() {
389 parent_children.push(node_id);
390 }
391
392 Some(node_id)
393 } else {
394 None
395 }
396 }
397
398 pub fn is_root(&self) -> bool {
400 self.stack.len() <= 1
401 }
402}
403
404impl Default for GreenTreeBuilder {
405 fn default() -> Self {
406 Self::new()
407 }
408}
409
410#[cfg(test)]
411mod tests {
412 use super::*;
413
414 #[test]
415 fn test_green_node_token() {
416 let mut arena = GreenNodeArena::new();
417 let token_id = arena.alloc_token(0, 5);
418 assert_eq!(arena.width(token_id), 5);
419 assert_eq!(arena.kind(token_id), None);
420 }
421
422 #[test]
423 fn test_green_node_internal() {
424 let mut arena = GreenNodeArena::new();
425 let token1 = arena.alloc_token(0, 5);
426 let token2 = arena.alloc_token(1, 3);
427 let internal = arena.alloc_internal(SyntaxKind::BinaryExpr, vec![token1, token2]);
428
429 assert_eq!(arena.width(internal), 8);
430 assert_eq!(arena.kind(internal), Some(SyntaxKind::BinaryExpr));
431 assert_eq!(arena.children(internal).unwrap().len(), 2);
432 }
433
434 #[test]
435 fn test_builder() {
436 let mut builder = GreenTreeBuilder::new();
437
438 builder.start_node(SyntaxKind::Program);
439 builder.start_node(SyntaxKind::IntLiteral);
440 builder.add_token(0, 2); builder.finish_node();
442 let root_id = builder.finish_node().unwrap();
443
444 let arena = builder.arena();
445 assert_eq!(arena.kind(root_id), Some(SyntaxKind::Program));
446 assert_eq!(arena.width(root_id), 2);
447 }
448
449 #[test]
450 fn test_nested_nodes() {
451 let mut builder = GreenTreeBuilder::new();
452
453 builder.start_node(SyntaxKind::BinaryExpr);
454 builder.add_token(0, 1); builder.add_token(1, 1); builder.add_token(2, 1); let node_id = builder.finish_node().unwrap();
458
459 let arena = builder.arena();
460 assert_eq!(arena.width(node_id), 3);
461 assert_eq!(arena.children(node_id).unwrap().len(), 3);
462 }
463}