Skip to main content

oak_core/tree/
red_tree.rs

1//! Red-green tree implementation with position information for incremental parsing.
2//!
3//! This module provides the "red" side of the red-green tree architecture.
4
5use crate::{
6    Language,
7    tree::green_tree::{GreenNode, GreenTree},
8};
9use core::range::Range;
10use std::fmt;
11
12/// A red tree element with absolute position information.
13pub enum RedTree<'a, L: Language> {
14    /// A red node with child elements
15    Node(RedNode<'a, L>),
16    /// A red leaf kind
17    Leaf(RedLeaf<L>),
18}
19
20// Manually implement Clone/Copy to avoid L: Copy bound
21impl<'a, L: Language> Clone for RedTree<'a, L> {
22    fn clone(&self) -> Self {
23        *self
24    }
25}
26
27impl<'a, L: Language> Copy for RedTree<'a, L> {}
28
29impl<'a, L: Language> fmt::Debug for RedTree<'a, L> {
30    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
31        match self {
32            Self::Node(node) => fmt::Debug::fmt(node, f),
33            Self::Leaf(leaf) => fmt::Debug::fmt(leaf, f),
34        }
35    }
36}
37
38impl<'a, L: Language> PartialEq for RedTree<'a, L> {
39    fn eq(&self, other: &Self) -> bool {
40        match (self, other) {
41            (Self::Node(l0), Self::Node(r0)) => l0 == r0,
42            (Self::Leaf(l0), Self::Leaf(r0)) => l0 == r0,
43            _ => false,
44        }
45    }
46}
47
48impl<'a, L: Language> Eq for RedTree<'a, L> {}
49
50impl<'a, L: Language> RedTree<'a, L> {
51    /// Returns the absolute byte span of this red tree element.
52    #[inline]
53    pub fn span(&self) -> Range<usize> {
54        match self {
55            RedTree::Node(n) => n.span(),
56            RedTree::Leaf(t) => t.span,
57        }
58    }
59
60    /// Returns the kind of this red tree element.
61    pub fn kind<T>(&self) -> T
62    where
63        T: From<L::ElementType> + From<L::TokenType>,
64    {
65        match self {
66            RedTree::Node(n) => T::from(n.green.kind),
67            RedTree::Leaf(l) => T::from(l.kind),
68        }
69    }
70
71    /// Returns the text of this red tree element from the source.
72    pub fn text<'s, S: crate::source::Source + ?Sized>(&self, source: &'s S) -> std::borrow::Cow<'s, str> {
73        source.get_text_in(self.span())
74    }
75
76    /// Returns an iterator over the child elements if this is a node.
77    pub fn children(&self) -> RedChildren<'a, L> {
78        match self {
79            RedTree::Node(n) => n.children(),
80            RedTree::Leaf(_) => RedChildren::empty(),
81        }
82    }
83
84    /// Returns this element as a node if it is one.
85    pub fn as_node(&self) -> Option<RedNode<'a, L>> {
86        match self {
87            RedTree::Node(n) => Some(*n),
88            RedTree::Leaf(_) => None,
89        }
90    }
91
92    /// Returns this element as a leaf if it is one.
93    pub fn as_leaf(&self) -> Option<RedLeaf<L>> {
94        match self {
95            RedTree::Node(_) => None,
96            RedTree::Leaf(l) => Some(*l),
97        }
98    }
99}
100
101/// A red node that wraps a green node with absolute offset information.
102pub struct RedNode<'a, L: Language> {
103    /// The underlying green node that contains the structural information
104    pub green: &'a GreenNode<'a, L>,
105    /// The absolute byte offset of this node in the source text
106    pub offset: usize,
107}
108
109// Manually implement Clone/Copy to avoid L: Copy bound
110impl<'a, L: Language> Clone for RedNode<'a, L> {
111    fn clone(&self) -> Self {
112        *self
113    }
114}
115
116impl<'a, L: Language> Copy for RedNode<'a, L> {}
117
118impl<'a, L: Language> PartialEq for RedNode<'a, L> {
119    fn eq(&self, other: &Self) -> bool {
120        self.green == other.green && self.offset == other.offset
121    }
122}
123
124impl<'a, L: Language> Eq for RedNode<'a, L> {}
125
126impl<'a, L: Language> fmt::Debug for RedNode<'a, L> {
127    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
128        f.debug_struct("RedNode").field("green", &self.green).field("offset", &self.offset).finish()
129    }
130}
131
132/// A red leaf kind with absolute position information.
133pub struct RedLeaf<L: Language> {
134    /// The kind kind/category
135    pub kind: L::TokenType,
136    /// The absolute byte span of this kind in the source text
137    pub span: Range<usize>,
138}
139
140// Manually implement Clone/Copy to avoid L: Copy bound
141impl<L: Language> Clone for RedLeaf<L> {
142    fn clone(&self) -> Self {
143        *self
144    }
145}
146
147impl<L: Language> Copy for RedLeaf<L> {}
148
149impl<L: Language> PartialEq for RedLeaf<L> {
150    fn eq(&self, other: &Self) -> bool {
151        self.kind == other.kind && self.span == other.span
152    }
153}
154
155impl<L: Language> Eq for RedLeaf<L> {}
156
157impl<L: Language> fmt::Debug for RedLeaf<L> {
158    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
159        f.debug_struct("RedLeaf").field("kind", &self.kind).field("span", &self.span).finish()
160    }
161}
162
163/// An iterator over the child elements of a red node.
164pub struct RedChildren<'a, L: Language> {
165    node: Option<RedNode<'a, L>>,
166    index: usize,
167    offset: usize,
168}
169
170impl<'a, L: Language> RedChildren<'a, L> {
171    /// Creates an empty iterator.
172    pub fn empty() -> Self {
173        Self { node: None, index: 0, offset: 0 }
174    }
175}
176
177impl<'a, L: Language> RedNode<'a, L> {
178    /// Creates a new red node from a green node and offset.
179    #[inline]
180    pub fn new(green: &'a GreenNode<'a, L>, offset: usize) -> Self {
181        Self { green, offset }
182    }
183
184    /// Returns the absolute byte span of this red node.
185    #[inline]
186    pub fn span(&self) -> Range<usize> {
187        Range { start: self.offset, end: self.offset + self.green.text_len() as usize }
188    }
189
190    /// Returns the kind of this red node.
191    pub fn kind<T>(&self) -> T
192    where
193        T: From<L::ElementType>,
194    {
195        T::from(self.green.kind)
196    }
197
198    /// Gets the child element at the specified index.
199    pub fn child_at(&self, idx: usize) -> RedTree<'a, L> {
200        let children = self.green.children();
201        let green_child = &children[idx];
202
203        // Calculate offset by summing lengths of previous siblings
204        let mut offset = self.offset;
205        for i in 0..idx {
206            offset += children[i].len() as usize;
207        }
208
209        match green_child {
210            GreenTree::Node(n) => RedTree::Node(RedNode::new(n, offset)),
211            GreenTree::Leaf(t) => RedTree::Leaf(RedLeaf { kind: t.kind, span: Range { start: offset, end: offset + t.length as usize } }),
212        }
213    }
214
215    /// Returns an iterator over the child elements of this red node.
216    pub fn children(&self) -> RedChildren<'a, L> {
217        RedChildren { node: Some(*self), index: 0, offset: self.offset }
218    }
219
220    /// Finds the child element at the specified absolute byte offset.
221    pub fn child_index_at_offset(&self, offset: usize) -> Option<usize> {
222        if offset < self.offset || offset >= self.offset + self.green.text_len() as usize {
223            return None;
224        }
225
226        let relative_offset = (offset - self.offset) as u32;
227        let mut current_pos = 0;
228
229        for (idx, child) in self.green.children().iter().enumerate() {
230            let len = child.len();
231            if relative_offset < current_pos + len {
232                return Some(idx);
233            }
234            current_pos += len;
235        }
236
237        None
238    }
239
240    /// Finds the child element at the specified absolute byte offset.
241    pub fn child_at_offset(&self, offset: usize) -> Option<RedTree<'a, L>> {
242        self.child_index_at_offset(offset).map(|idx| self.child_at(idx))
243    }
244
245    /// Finds the leaf element at the specified absolute byte offset by traversing down the tree.
246    pub fn leaf_at_offset(&self, offset: usize) -> Option<RedLeaf<L>> {
247        let mut current = *self;
248        loop {
249            match current.child_at_offset(offset)? {
250                RedTree::Node(n) => current = n,
251                RedTree::Leaf(l) => return Some(l),
252            }
253        }
254    }
255}
256
257impl<'a, L: Language> Iterator for RedChildren<'a, L> {
258    type Item = RedTree<'a, L>;
259
260    fn next(&mut self) -> Option<Self::Item> {
261        let node = self.node.as_ref()?;
262        let children = node.green.children();
263        if self.index >= children.len() {
264            return None;
265        }
266
267        let ch = &children[self.index];
268        let offset = self.offset;
269        let elem = match ch {
270            GreenTree::Node(n) => RedTree::Node(RedNode::new(n, offset)),
271            GreenTree::Leaf(t) => RedTree::Leaf(RedLeaf { kind: t.kind, span: Range { start: offset, end: offset + t.length as usize } }),
272        };
273
274        self.offset += ch.len() as usize;
275        self.index += 1;
276        Some(elem)
277    }
278}