wat_syntax/red/
node.rs

1use super::{element::SyntaxElement, traversal::Descendants};
2use crate::{
3    AmberNode, GreenNode, GreenToken, NodeOrToken, SyntaxKind, SyntaxKindMatch, SyntaxNodeChildren, SyntaxToken,
4    TokenAtOffset, green::GreenChild,
5};
6use std::{fmt, ptr::NonNull, rc::Rc};
7use text_size::{TextRange, TextSize};
8
9#[derive(Clone, PartialEq, Eq, Hash)]
10/// Node in the red syntax tree.
11pub struct SyntaxNode {
12    pub(crate) data: Rc<NodeData>,
13}
14
15impl SyntaxNode {
16    #[inline]
17    /// Build a new syntax tree on top of a green tree.
18    pub fn new_root(green: GreenNode) -> Self {
19        SyntaxNode {
20            data: Rc::new(NodeData {
21                range: TextRange::new(0.into(), green.text_len()),
22                level: NodeLevel::Root { green },
23                index: 0,
24            }),
25        }
26    }
27
28    #[inline]
29    pub(crate) fn new(index: u32, green: &GreenNode, offset: TextSize, parent: Rc<NodeData>) -> Self {
30        SyntaxNode {
31            data: Rc::new(NodeData {
32                range: TextRange::new(offset, offset + green.text_len()),
33                level: NodeLevel::Child {
34                    green: NonNull::from(green),
35                    parent,
36                },
37                index,
38            }),
39        }
40    }
41    #[inline]
42    pub(crate) fn new_child(&self, index: u32, green: &GreenNode, offset: TextSize) -> Self {
43        SyntaxNode::new(index, green, self.data.range.start() + offset, Rc::clone(&self.data))
44    }
45    #[inline]
46    pub(crate) fn new_token(&self, index: u32, green: &GreenToken, offset: TextSize) -> SyntaxToken {
47        SyntaxToken::new(index, green, self.data.range.start() + offset, Rc::clone(&self.data))
48    }
49
50    #[inline]
51    /// Kind of this node.
52    pub fn kind(&self) -> SyntaxKind {
53        self.green().kind()
54    }
55
56    #[inline]
57    /// The range that this node covers in the original text.
58    pub fn text_range(&self) -> TextRange {
59        self.data.range
60    }
61
62    #[inline]
63    /// The underlying green node of this red node.
64    pub fn green(&self) -> &GreenNode {
65        self.data.green()
66    }
67
68    #[inline]
69    /// The corresponding amber node of this red node.
70    pub fn amber(&self) -> AmberNode<'_> {
71        self.into()
72    }
73
74    #[inline]
75    /// Parent of this node. It returns `None` if this node is the root.
76    pub fn parent(&self) -> Option<SyntaxNode> {
77        match &self.data.level {
78            NodeLevel::Root { .. } => None,
79            NodeLevel::Child { parent, .. } => Some(SyntaxNode {
80                data: Rc::clone(parent),
81            }),
82        }
83    }
84
85    #[inline]
86    /// Iterator along the chain of parents of this node.
87    pub fn ancestors(&self) -> impl Iterator<Item = SyntaxNode> {
88        std::iter::successors(Some(self.clone()), SyntaxNode::parent)
89    }
90
91    #[inline]
92    /// Iterator over the child nodes of this node.
93    ///
94    /// If you want to iterate over both nodes and tokens, use [`children_with_tokens`](Self::children_with_tokens) instead.
95    ///
96    /// Though you can filter specific kinds of children on this iterator manually,
97    /// it is more efficient to use [`children_by_kind`](Self::children_by_kind) instead.
98    pub fn children(&self) -> SyntaxNodeChildren {
99        SyntaxNodeChildren {
100            parent: self.clone(),
101            green: match &self.data.level {
102                NodeLevel::Root { green } => NonNull::from(green),
103                NodeLevel::Child { green, .. } => *green,
104            },
105            index: 0,
106        }
107    }
108
109    #[inline]
110    /// Iterator over specific kinds of child nodes of this node.
111    /// This is more efficient than filtering with [`children`](Self::children) manually.
112    pub fn children_by_kind<M>(&self, matcher: M) -> impl DoubleEndedIterator<Item = SyntaxNode> + use<'_, M>
113    where
114        M: SyntaxKindMatch,
115    {
116        self.green()
117            .slice()
118            .iter()
119            .enumerate()
120            .filter_map(move |(i, child)| match child {
121                GreenChild::Node { offset, node } if matcher.matches(node.kind()) => {
122                    Some(self.new_child(i as u32, node, *offset))
123                }
124                _ => None,
125            })
126    }
127
128    #[inline]
129    /// Iterator over specific kinds of child tokens of this node.
130    pub fn tokens_by_kind<M>(&self, matcher: M) -> impl DoubleEndedIterator<Item = SyntaxToken> + use<'_, M>
131    where
132        M: SyntaxKindMatch,
133    {
134        self.green()
135            .slice()
136            .iter()
137            .enumerate()
138            .filter_map(move |(i, child)| match child {
139                GreenChild::Token { offset, token } if matcher.matches(token.kind()) => {
140                    Some(self.new_token(i as u32, token, *offset))
141                }
142                _ => None,
143            })
144    }
145
146    #[inline]
147    /// Iterator over the child nodes and tokens of this node.
148    pub fn children_with_tokens(&self) -> impl DoubleEndedIterator<Item = SyntaxElement> {
149        self.green().slice().iter().enumerate().map(|(i, child)| match child {
150            GreenChild::Node { offset, node } => self.new_child(i as u32, node, *offset).into(),
151            GreenChild::Token { offset, token } => self.new_token(i as u32, token, *offset).into(),
152        })
153    }
154
155    #[inline]
156    /// Check if this node has specific kinds of child nodes or tokens.
157    ///
158    /// This is an efficient alternative to `node.children_with_tokens().any(...)`
159    /// since it won't create any nodes or tokens.
160    pub fn has_child_or_token_by_kind<M>(&self, matcher: M) -> bool
161    where
162        M: SyntaxKindMatch,
163    {
164        self.green().slice().iter().any(|child| match child {
165            GreenChild::Node { node, .. } => matcher.matches(node.kind()),
166            GreenChild::Token { token, .. } => matcher.matches(token.kind()),
167        })
168    }
169
170    #[inline]
171    /// Nodes that come immediately after this node.
172    ///
173    /// If you want to iterate over both nodes and tokens, use [`next_siblings_with_tokens`](Self::next_siblings_with_tokens) instead.
174    ///
175    /// Unlike rowan, the iterator doesn't contain the current node itself.
176    pub fn next_siblings(&self) -> impl Iterator<Item = SyntaxNode> {
177        match &self.data.level {
178            NodeLevel::Root { .. } => None,
179            NodeLevel::Child { parent, .. } => Some(parent),
180        }
181        .into_iter()
182        .flat_map(|parent| parent.next_children(self.data.index))
183    }
184
185    #[inline]
186    /// Node or token that comes immediately after this node.
187    pub fn next_sibling_or_token(&self) -> Option<SyntaxElement> {
188        match &self.data.level {
189            NodeLevel::Root { .. } => None,
190            NodeLevel::Child { parent, .. } => parent.next_child_or_token(self.data.index),
191        }
192    }
193
194    #[inline]
195    /// Nodes and tokens that come immediately after this node.
196    ///
197    /// Unlike rowan, the iterator doesn't contain the current node itself.
198    pub fn next_siblings_with_tokens(&self) -> impl Iterator<Item = SyntaxElement> {
199        match &self.data.level {
200            NodeLevel::Root { .. } => None,
201            NodeLevel::Child { parent, .. } => Some(parent),
202        }
203        .into_iter()
204        .flat_map(|parent| parent.next_children_with_tokens(self.data.index))
205    }
206
207    #[inline]
208    /// Consecutive tokens sequence that come immediately after this node without any nodes in between.
209    pub fn next_consecutive_tokens(&self) -> impl Iterator<Item = SyntaxToken> {
210        match &self.data.level {
211            NodeLevel::Root { .. } => None,
212            NodeLevel::Child { parent, .. } => Some(parent),
213        }
214        .into_iter()
215        .flat_map(|parent| parent.next_consecutive_tokens(self.data.index))
216    }
217
218    #[inline]
219    /// Nodes that come immediately before this node.
220    ///
221    /// If you want to iterate over both nodes and tokens, use [`prev_siblings_with_tokens`](Self::prev_siblings_with_tokens) instead.
222    ///
223    /// Unlike rowan, the iterator doesn't contain the current node itself.
224    pub fn prev_siblings(&self) -> impl Iterator<Item = SyntaxNode> {
225        match &self.data.level {
226            NodeLevel::Root { .. } => None,
227            NodeLevel::Child { parent, .. } => Some(parent),
228        }
229        .into_iter()
230        .flat_map(|parent| parent.prev_children(self.data.index))
231    }
232
233    #[inline]
234    /// Node or token that comes immediately before this node.
235    pub fn prev_sibling_or_token(&self) -> Option<SyntaxElement> {
236        match &self.data.level {
237            NodeLevel::Root { .. } => None,
238            NodeLevel::Child { parent, .. } => parent.prev_child_or_token(self.data.index),
239        }
240    }
241
242    #[inline]
243    /// Nodes and tokens that come immediately before this node.
244    ///
245    /// Unlike rowan, the iterator doesn't contain the current node itself.
246    pub fn prev_siblings_with_tokens(&self) -> impl Iterator<Item = SyntaxElement> {
247        match &self.data.level {
248            NodeLevel::Root { .. } => None,
249            NodeLevel::Child { parent, .. } => Some(parent),
250        }
251        .into_iter()
252        .flat_map(|parent| parent.prev_children_with_tokens(self.data.index))
253    }
254
255    #[inline]
256    /// Consecutive tokens sequence that come immediately before this node without any nodes in between.
257    pub fn prev_consecutive_tokens(&self) -> impl Iterator<Item = SyntaxToken> {
258        match &self.data.level {
259            NodeLevel::Root { .. } => None,
260            NodeLevel::Child { parent, .. } => Some(parent),
261        }
262        .into_iter()
263        .flat_map(|parent| parent.prev_consecutive_tokens(self.data.index))
264    }
265
266    #[inline]
267    /// Iterator over all nodes in the subtree, including this node itself.
268    pub fn descendants(&self) -> impl Iterator<Item = SyntaxNode> {
269        Descendants::new(self.clone())
270    }
271
272    /// Find a token in the subtree corresponding to this node, which covers the offset.
273    // Copied from rowan with modification.
274    pub fn token_at_offset(&self, offset: TextSize) -> TokenAtOffset {
275        let range = self.data.range;
276        let relative_offset = offset - range.start();
277        let mut children = self
278            .green()
279            .slice()
280            .iter()
281            .enumerate()
282            .filter(|(_, child)| {
283                let (start, end) = match child {
284                    GreenChild::Node { offset, node } => (*offset, offset + node.text_len()),
285                    GreenChild::Token { offset, token } => (*offset, offset + token.text_len()),
286                };
287                start <= relative_offset && relative_offset <= end
288            })
289            .map(|(i, child)| match child {
290                GreenChild::Node { offset, node } => NodeOrToken::Node(self.new_child(i as u32, node, *offset)),
291                GreenChild::Token { offset, token } => NodeOrToken::Token(self.new_token(i as u32, token, *offset)),
292            });
293
294        let Some(left) = children.next() else {
295            return TokenAtOffset::None;
296        };
297        let right = children.next();
298        if let Some(right) = right {
299            match (left.token_at_offset(offset), right.token_at_offset(offset)) {
300                (TokenAtOffset::Single(left), TokenAtOffset::Single(right)) => TokenAtOffset::Between(left, right),
301                _ => TokenAtOffset::None,
302            }
303        } else {
304            left.token_at_offset(offset)
305        }
306    }
307
308    #[inline]
309    /// Find a child node that intersects with the given range.
310    pub fn child_at_range(&self, range: TextRange) -> Option<SyntaxNode> {
311        if !self.data.range.contains_range(range) {
312            return None;
313        }
314        let relative_range = range - self.data.range.start();
315        let slice = self.green().slice();
316        let i = slice
317            .binary_search_by(|child| match child {
318                GreenChild::Node { offset, node } => {
319                    TextRange::new(*offset, offset + node.text_len()).ordering(relative_range)
320                }
321                GreenChild::Token { offset, token } => {
322                    TextRange::new(*offset, offset + token.text_len()).ordering(relative_range)
323                }
324            })
325            .unwrap_or_else(|i| i.saturating_sub(1)); // not sure why but rowan does it
326        slice.get(i).and_then(|child| match child {
327            GreenChild::Node { offset, node } => {
328                if TextRange::new(*offset, offset + node.text_len()).contains_range(relative_range) {
329                    Some(self.new_child(i as u32, node, *offset))
330                } else {
331                    None
332                }
333            }
334            GreenChild::Token { .. } => None,
335        })
336    }
337
338    #[inline]
339    /// Replace this node with new green node. It returns new *root* green node with replaced node.
340    pub fn replace_with(&self, replacement: GreenNode) -> GreenNode {
341        match &self.data.level {
342            NodeLevel::Root { .. } => replacement,
343            NodeLevel::Child { parent, .. } => {
344                let parent = SyntaxNode {
345                    data: Rc::clone(parent),
346                };
347                let new_parent = parent.green().replace_child(self.data.index as usize, replacement);
348                parent.replace_with(new_parent)
349            }
350        }
351    }
352}
353
354impl fmt::Debug for SyntaxNode {
355    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
356        debug_print_node(self, f, 0)
357    }
358}
359fn debug_print_node(node: &SyntaxNode, f: &mut fmt::Formatter<'_>, level: usize) -> fmt::Result {
360    writeln!(f, "{}{:?}@{:?}", "  ".repeat(level), node.kind(), node.text_range())?;
361    node.children_with_tokens()
362        .try_for_each(|node_or_token| match node_or_token {
363            NodeOrToken::Node(node) => debug_print_node(&node, f, level + 1),
364            NodeOrToken::Token(token) => writeln!(f, "{}{token:?}", "  ".repeat(level + 1)),
365        })
366}
367
368impl fmt::Display for SyntaxNode {
369    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
370        self.green().fmt(f)
371    }
372}
373
374#[derive(PartialEq, Eq, Hash)]
375pub(crate) struct NodeData {
376    range: TextRange,
377    level: NodeLevel,
378    index: u32,
379}
380
381impl NodeData {
382    #[inline]
383    fn green(&self) -> &GreenNode {
384        match &self.level {
385            NodeLevel::Root { green } => green,
386            NodeLevel::Child { green, .. } => unsafe { green.as_ref() },
387        }
388    }
389
390    #[inline]
391    pub fn next_children(self: &Rc<Self>, index: u32) -> impl Iterator<Item = SyntaxNode> {
392        self.green()
393            .slice()
394            .iter()
395            .enumerate()
396            .skip(index as usize + 1)
397            .filter_map(|(i, child)| match child {
398                GreenChild::Node { offset, node } => Some(SyntaxNode::new(
399                    i as u32,
400                    node,
401                    self.range.start() + offset,
402                    Rc::clone(self),
403                )),
404                _ => None,
405            })
406    }
407
408    #[inline]
409    pub fn next_child_or_token(self: &Rc<Self>, index: u32) -> Option<SyntaxElement> {
410        let i = index + 1;
411        self.green().slice().get(i as usize).map(|child| match child {
412            GreenChild::Node { offset, node } => {
413                SyntaxNode::new(i, node, self.range.start() + offset, Rc::clone(self)).into()
414            }
415            GreenChild::Token { offset, token } => {
416                SyntaxToken::new(i, token, self.range.start() + offset, Rc::clone(self)).into()
417            }
418        })
419    }
420
421    #[inline]
422    pub fn next_children_with_tokens(self: &Rc<Self>, index: u32) -> impl Iterator<Item = SyntaxElement> {
423        self.green()
424            .slice()
425            .iter()
426            .enumerate()
427            .skip(index as usize + 1)
428            .map(|(i, child)| match child {
429                GreenChild::Node { offset, node } => {
430                    SyntaxNode::new(i as u32, node, self.range.start() + offset, Rc::clone(self)).into()
431                }
432                GreenChild::Token { offset, token } => {
433                    SyntaxToken::new(i as u32, token, self.range.start() + offset, Rc::clone(self)).into()
434                }
435            })
436    }
437
438    #[inline]
439    pub fn next_consecutive_tokens(self: &Rc<Self>, index: u32) -> impl Iterator<Item = SyntaxToken> {
440        self.green()
441            .slice()
442            .iter()
443            .enumerate()
444            .skip(index as usize + 1)
445            .map_while(|(i, child)| match child {
446                GreenChild::Node { .. } => None,
447                GreenChild::Token { offset, token } => Some(SyntaxToken::new(
448                    i as u32,
449                    token,
450                    self.range.start() + offset,
451                    Rc::clone(self),
452                )),
453            })
454    }
455
456    #[inline]
457    pub fn prev_children(self: &Rc<Self>, index: u32) -> impl Iterator<Item = SyntaxNode> {
458        let slice = self.green().slice();
459        slice
460            .iter()
461            .enumerate()
462            .rev()
463            .skip(slice.len() - index as usize)
464            .filter_map(|(i, child)| match child {
465                GreenChild::Node { offset, node } => Some(SyntaxNode::new(
466                    i as u32,
467                    node,
468                    self.range.start() + offset,
469                    Rc::clone(self),
470                )),
471                _ => None,
472            })
473    }
474
475    #[inline]
476    pub fn prev_child_or_token(self: &Rc<Self>, index: u32) -> Option<SyntaxElement> {
477        let i = index.checked_sub(1)?;
478        self.green().slice().get(i as usize).map(|child| match child {
479            GreenChild::Node { offset, node } => {
480                SyntaxNode::new(i, node, self.range.start() + offset, Rc::clone(self)).into()
481            }
482            GreenChild::Token { offset, token } => {
483                SyntaxToken::new(i, token, self.range.start() + offset, Rc::clone(self)).into()
484            }
485        })
486    }
487
488    #[inline]
489    pub fn prev_children_with_tokens(self: &Rc<Self>, index: u32) -> impl Iterator<Item = SyntaxElement> {
490        let slice = self.green().slice();
491        slice
492            .iter()
493            .enumerate()
494            .rev()
495            .skip(slice.len() - index as usize)
496            .map(|(i, child)| match child {
497                GreenChild::Node { offset, node } => {
498                    SyntaxNode::new(i as u32, node, self.range.start() + offset, Rc::clone(self)).into()
499                }
500                GreenChild::Token { offset, token } => {
501                    SyntaxToken::new(i as u32, token, self.range.start() + offset, Rc::clone(self)).into()
502                }
503            })
504    }
505
506    #[inline]
507    pub fn prev_consecutive_tokens(self: &Rc<Self>, index: u32) -> impl Iterator<Item = SyntaxToken> {
508        let slice = self.green().slice();
509        slice
510            .iter()
511            .enumerate()
512            .rev()
513            .skip(slice.len() - index as usize)
514            .map_while(|(i, child)| match child {
515                GreenChild::Node { .. } => None,
516                GreenChild::Token { offset, token } => Some(SyntaxToken::new(
517                    i as u32,
518                    token,
519                    self.range.start() + offset,
520                    Rc::clone(self),
521                )),
522            })
523    }
524}
525
526#[derive(PartialEq, Eq, Hash)]
527enum NodeLevel {
528    Root {
529        green: GreenNode,
530    },
531    Child {
532        green: NonNull<GreenNode>,
533        parent: Rc<NodeData>,
534    },
535}
wat_syntax/red/node.rs

wat_syntax/red/
node.rs
