1use core::{
2 fmt,
3 ops::{Index, IndexMut},
4};
5
6use alloc::{boxed::Box, vec::Vec};
7
8use crate::compile::{Range, Token};
9
10#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)]
11pub struct NodeId(usize);
12impl From<usize> for NodeId {
13 fn from(id: usize) -> Self {
14 NodeId(id)
15 }
16}
17impl From<NodeId> for usize {
18 fn from(id: NodeId) -> usize {
19 id.0
20 }
21}
22
23#[derive(Clone)]
24pub struct Node {
25 pub token: Token,
26 pub range: Range,
27 pub parent: Option<NodeId>,
28 pub next_sibling: Option<NodeId>,
29 pub child: Option<NodeId>,
30}
31impl Node {
32 pub fn children<'a>(&self, arena: &'a Tree) -> NodeIter<'a> {
33 NodeIter(arena, self.child)
34 }
35}
36impl fmt::Debug for Node {
37 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
38 write!(f, "{:?} {:?}", self.token, self.range)?;
39 Ok(())
40 }
41}
42
43pub struct NodeIter<'a>(&'a Tree, Option<NodeId>);
44impl Iterator for NodeIter<'_> {
45 type Item = NodeId;
46 fn next(&mut self) -> Option<Self::Item> {
47 if let Some(next) = self.1 {
48 self.1 = self.0[next].next_sibling;
49 Some(next)
50 } else {
51 None
52 }
53 }
54}
55
56pub struct Checkpoint {
57 cursor: Option<NodeId>,
58}
59
60#[derive(Default)]
61pub struct TreeBuilder {
62 arena: Vec<Node>,
63 parent: Option<NodeId>,
64 cursor: Option<NodeId>,
65}
66impl TreeBuilder {
67 fn insert(&mut self, token: Token, range: Range) -> NodeId {
68 let id = NodeId::from(self.arena.len());
69 self.arena.push(Node {
70 token,
71 range,
72 parent: self.parent,
73 next_sibling: None,
74 child: None,
75 });
76 if let Some(prev) = self.cursor {
77 self.arena[usize::from(prev)].next_sibling = Some(id);
78 }
79 if let Some(parent) = self.parent {
80 self.arena[usize::from(parent)].child =
81 self.arena[usize::from(parent)].child.or(Some(id));
82 }
83 id
84 }
85 pub fn leaf(&mut self, token: Token, range: Range) {
86 self.cursor = Some(self.insert(token, range));
87 }
88 pub fn start_internal(&mut self, token: Token, range: Range) {
89 self.parent = Some(self.insert(token, range));
90 self.cursor = None;
91 }
92 pub fn finish_internal(&mut self) {
93 self.cursor = self.parent;
94 self.parent = self
95 .parent
96 .and_then(|parent| self.arena[usize::from(parent)].parent);
97 }
98 pub fn checkpoint(&self) -> Checkpoint {
99 Checkpoint {
100 cursor: self.cursor,
101 }
102 }
103 pub fn start_internal_at(&mut self, checkpoint: Checkpoint, token: Token, range: Range) {
104 let parent = if let Some(from) = checkpoint.cursor {
105 let id = NodeId::from(self.arena.len());
106 self.arena.push(Node {
107 token,
108 range,
109 parent: self.parent,
110 next_sibling: None,
111 child: self.arena[usize::from(from)].next_sibling,
112 });
113 self.arena[usize::from(from)].next_sibling = Some(id);
114 id
115 } else if let Some(parent) = self.parent {
116 let id = NodeId::from(self.arena.len());
117 self.arena.push(Node {
118 token,
119 range,
120 parent: self.parent,
121 next_sibling: None,
122 child: self.arena[usize::from(parent)].child,
123 });
124 self.arena[usize::from(parent)].child = Some(id);
125 id
126 } else {
127 let id = NodeId::from(self.arena.len());
128 self.arena.push(Node {
129 token,
130 range,
131 parent: None,
132 next_sibling: None,
133 child: self.cursor,
134 });
135 id
136 };
137 let mut next = self.arena[usize::from(parent)].child;
139 while let Some(node) = next {
140 let node = &mut self.arena[usize::from(node)];
141 next = node.next_sibling;
142 node.parent = Some(parent);
143 }
144 self.parent = Some(parent);
145 self.cursor = None;
146 }
147 pub fn finish(self) -> Tree {
148 assert!(self.cursor.is_some(), "no item");
149 let cursor = self.cursor.unwrap();
150
151 Tree {
152 arena: self.arena.into_boxed_slice(),
153 root: cursor,
154 }
155 }
156}
157
158#[derive(Clone)]
159pub struct Tree {
160 pub arena: Box<[Node]>,
161 pub root: NodeId,
162}
163impl Index<NodeId> for Tree {
164 type Output = Node;
165 fn index(&self, index: NodeId) -> &Node {
166 &self.arena[usize::from(index)]
167 }
168}
169impl IndexMut<NodeId> for Tree {
170 fn index_mut(&mut self, index: NodeId) -> &mut Node {
171 &mut self.arena[usize::from(index)]
172 }
173}
174impl fmt::Debug for Tree {
175 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
176 let mut next = Some(self.root);
177 let mut nested: usize = 0;
178 'outer: while let Some(id) = next {
179 let node = &self[id];
180 writeln!(f, "{:indent$}{:?}", "", node, indent = nested * 2)?;
181 if node.child.is_some() {
182 nested += 1;
183 next = node.child;
184 } else {
185 let mut me = Some(node);
186 while me.map(|me| me.next_sibling.is_none()).unwrap_or(false) {
187 match nested.checked_sub(1) {
188 Some(new) => nested = new,
189 None => break 'outer,
190 }
191 me = me.unwrap().parent.map(|id| &self[id]);
192 }
193 next = me.and_then(|me| me.next_sibling);
194 }
195 }
196 Ok(())
197 }
198}
199
200#[cfg(test)]
201mod tests {
202 use alloc::format;
203
204 use super::*;
205
206 fn sanity_check(tree: &Tree) {
207 let mut next = Some(tree.root);
208 let mut parent = None;
209 while let Some(id) = next {
210 let node = &tree[id];
211 assert_eq!(parent, node.parent);
212
213 if let Some(child) = node.child {
214 next = Some(child);
215 parent = Some(id);
216 } else {
217 let mut node = Some(id);
218 while node
219 .map(|node| tree[node].next_sibling.is_none())
220 .unwrap_or(false)
221 {
222 node = tree[node.unwrap()].parent;
223 }
224 next = node.and_then(|node| tree[node].next_sibling);
225 parent = node.and_then(|node| tree[node].parent);
226 }
227 }
228 }
229
230 #[test]
231 fn simple_builder() {
232 let mut builder = TreeBuilder::default();
233 builder.start_internal(Token::Root, Range(1, Some(1)));
234 builder.start_internal(Token::Alternative, Range(1, Some(1)));
235 builder.leaf(Token::Start, Range(1, Some(1)));
236 builder.start_internal(Token::Group(1), Range(1, Some(1)));
237 builder.start_internal(Token::Alternative, Range(1, Some(1)));
238 builder.leaf(Token::Any, Range(1, Some(1)));
239 builder.finish_internal();
240 builder.finish_internal();
241 builder.finish_internal();
242 builder.start_internal(Token::Alternative, Range(1, Some(1)));
243 builder.leaf(Token::End, Range(1, Some(1)));
244 builder.finish_internal();
245 builder.finish_internal();
246
247 let tree = builder.finish();
248 sanity_check(&tree);
249
250 assert_eq!(
251 format!("{:?}", tree),
252 "\
253Root 1..1
254 Alternative 1..1
255 ^ 1..1
256 Group(1) 1..1
257 Alternative 1..1
258 . 1..1
259 Alternative 1..1
260 $ 1..1
261"
262 );
263 }
264 #[test]
265 fn builder_checkpoint() {
266 let mut builder = TreeBuilder::default();
267 builder.start_internal(Token::Root, Range(1, Some(1)));
268 let mut alternation = builder.checkpoint();
269 builder.leaf(Token::Start, Range(1, Some(1)));
270 let group = builder.checkpoint();
271 builder.start_internal(Token::Alternative, Range(1, Some(1)));
272 builder.leaf(Token::Any, Range(1, Some(1)));
273 builder.finish_internal();
274 builder.start_internal_at(group, Token::Group(1), Range(1, Some(1)));
275 builder.finish_internal();
276 builder.start_internal_at(alternation, Token::Alternative, Range(1, Some(1)));
277 builder.finish_internal();
278 alternation = builder.checkpoint();
279 builder.leaf(Token::End, Range(1, Some(1)));
280 builder.start_internal_at(alternation, Token::Alternative, Range(1, Some(1)));
281 builder.finish_internal();
282 builder.finish_internal();
283
284 let tree = builder.finish();
285 sanity_check(&tree);
286
287 assert_eq!(
288 format!("{:?}", tree),
289 "\
290Root 1..1
291 Alternative 1..1
292 ^ 1..1
293 Group(1) 1..1
294 Alternative 1..1
295 . 1..1
296 Alternative 1..1
297 $ 1..1
298"
299 );
300 }
301}