oak_core/ast/
mod.rs

1//! Unified AST processing interface for Oaks core framework.
2//!
3//! This module provides a unified interface for processing ASTs in Oaks,
4//! including traversal, modification, and querying methods.
5
6use crate::{Language, RedLeaf, RedNode, RedTree};
7
8/// A unified interface for processing ASTs.
9///
10/// This trait provides a common set of methods for working with ASTs,
11/// regardless of the specific language implementation.
12pub trait AstProcessor<'a, L: Language> {
13    /// Traverses the AST in pre-order, applying the given visitor to each node.
14    fn traverse_pre_order(&self, root: RedTree<'a, L>, visitor: &mut impl Visitor<'a, L>);
15
16    /// Traverses the AST in post-order, applying the given visitor to each node.
17    fn traverse_post_order(&self, root: RedTree<'a, L>, visitor: &mut impl Visitor<'a, L>);
18
19    /// Finds all nodes of the specified type.
20    fn find_nodes_by_type(&self, root: RedTree<'a, L>, node_type: L::ElementType) -> Vec<RedNode<'a, L>>;
21
22    /// Finds all tokens of the specified type.
23    fn find_tokens_by_type(&self, root: RedTree<'a, L>, token_type: L::TokenType) -> Vec<RedLeaf<L>>;
24
25    /// Finds the node at the specified offset.
26    fn find_node_at_offset(&self, root: RedTree<'a, L>, offset: usize) -> Option<RedTree<'a, L>>;
27
28    /// Replaces a node with a new one.
29    ///
30    /// Note: This operation may require rebuilding parts of the tree.
31    fn replace_node(&self, old_node: RedNode<'a, L>, new_node: RedNode<'a, L>) -> Result<RedTree<'a, L>, AstError>;
32
33    /// Removes a node from the tree.
34    fn remove_node(&self, node: RedNode<'a, L>) -> Result<RedTree<'a, L>, AstError>;
35
36    /// Adds a child node to the specified parent.
37    fn add_child(&self, parent: RedNode<'a, L>, child: RedTree<'a, L>) -> Result<RedNode<'a, L>, AstError>;
38}
39
40/// A visitor for AST nodes.
41///
42/// This trait extends the existing Visitor trait with additional methods
43/// for more flexible traversal and modification.
44pub trait Visitor<'a, L: Language> {
45    /// Visits a node before its children.
46    fn visit_node_pre(&mut self, node: RedNode<'a, L>) -> VisitResult;
47
48    /// Visits a node after its children.
49    fn visit_node_post(&mut self, node: RedNode<'a, L>) -> VisitResult;
50
51    /// Visits a token.
52    fn visit_token(&mut self, token: RedLeaf<L>) -> VisitResult;
53
54    /// Default implementation that just continues traversal.
55    fn visit_node_pre_default(&mut self, _node: RedNode<'a, L>) -> VisitResult {
56        VisitResult::Continue
57    }
58
59    /// Default implementation that just continues traversal.
60    fn visit_node_post_default(&mut self, _node: RedNode<'a, L>) -> VisitResult {
61        VisitResult::Continue
62    }
63
64    /// Default implementation that just continues traversal.
65    fn visit_token_default(&mut self, _token: RedLeaf<L>) -> VisitResult {
66        VisitResult::Continue
67    }
68}
69
70/// Result of a visit operation.
71pub enum VisitResult {
72    /// Continue traversal.
73    Continue,
74    /// Skip the children of the current node.
75    SkipChildren,
76    /// Stop traversal entirely.
77    Stop,
78}
79
80/// Error type for AST operations.
81pub enum AstError {
82    /// The operation failed because the node was not found.
83    NodeNotFound,
84    /// The operation failed because the node is not modifiable.
85    NodeNotModifiable,
86    /// The operation failed because of an invalid operation.
87    InvalidOperation,
88    /// The operation failed because of a language-specific error.
89    LanguageSpecific(String),
90}
91
92impl std::fmt::Debug for AstError {
93    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
94        match self {
95            AstError::NodeNotFound => write!(f, "Node not found"),
96            AstError::NodeNotModifiable => write!(f, "Node not modifiable"),
97            AstError::InvalidOperation => write!(f, "Invalid operation"),
98            AstError::LanguageSpecific(msg) => write!(f, "Language-specific error: {}", msg),
99        }
100    }
101}
102
103impl std::fmt::Display for AstError {
104    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
105        std::fmt::Debug::fmt(self, f)
106    }
107}
108
109impl std::error::Error for AstError {}
110
111/// Default implementation of AstProcessor.
112pub struct DefaultAstProcessor;
113
114impl<'a, L: Language> AstProcessor<'a, L> for DefaultAstProcessor {
115    fn traverse_pre_order(&self, root: RedTree<'a, L>, visitor: &mut impl Visitor<'a, L>) {
116        self.traverse_pre_order_impl(root, visitor)
117    }
118
119    fn traverse_post_order(&self, root: RedTree<'a, L>, visitor: &mut impl Visitor<'a, L>) {
120        self.traverse_post_order_impl(root, visitor)
121    }
122
123    fn find_nodes_by_type(&self, root: RedTree<'a, L>, node_type: L::ElementType) -> Vec<RedNode<'a, L>> {
124        let mut result = Vec::new();
125        let mut visitor = TypeFinderVisitor { node_type, result: &mut result };
126        self.traverse_pre_order(root, &mut visitor);
127        result
128    }
129
130    fn find_tokens_by_type(&self, root: RedTree<'a, L>, token_type: L::TokenType) -> Vec<RedLeaf<L>> {
131        let mut result = Vec::new();
132        let mut visitor = TokenTypeFinderVisitor { token_type, result: &mut result };
133        self.traverse_pre_order(root, &mut visitor);
134        result
135    }
136
137    fn find_node_at_offset(&self, root: RedTree<'a, L>, offset: usize) -> Option<RedTree<'a, L>> {
138        match root {
139            RedTree::Node(node) => {
140                if offset >= node.span().start && offset < node.span().end {
141                    node.child_at_offset(offset)
142                }
143                else {
144                    None
145                }
146            }
147            RedTree::Leaf(leaf) => {
148                if offset >= leaf.span.start && offset < leaf.span.end {
149                    Some(RedTree::Leaf(leaf))
150                }
151                else {
152                    None
153                }
154            }
155        }
156    }
157
158    fn replace_node(&self, _old_node: RedNode<'a, L>, _new_node: RedNode<'a, L>) -> Result<RedTree<'a, L>, AstError> {
159        // TODO: Implement node replacement
160        Err(AstError::NodeNotModifiable)
161    }
162
163    fn remove_node(&self, _node: RedNode<'a, L>) -> Result<RedTree<'a, L>, AstError> {
164        // TODO: Implement node removal
165        Err(AstError::NodeNotModifiable)
166    }
167
168    fn add_child(&self, _parent: RedNode<'a, L>, _child: RedTree<'a, L>) -> Result<RedNode<'a, L>, AstError> {
169        // TODO: Implement child addition
170        Err(AstError::NodeNotModifiable)
171    }
172}
173
174impl DefaultAstProcessor {
175    /// Creates a new DefaultAstProcessor.
176    pub fn new() -> Self {
177        Self
178    }
179
180    /// Implementation of pre-order traversal.
181    fn traverse_pre_order_impl<'a, L: Language>(&self, node: RedTree<'a, L>, visitor: &mut impl Visitor<'a, L>) {
182        match node {
183            RedTree::Node(n) => match visitor.visit_node_pre(n) {
184                VisitResult::Continue => {
185                    for child in n.children() {
186                        self.traverse_pre_order_impl(child, visitor);
187                    }
188                    visitor.visit_node_post(n);
189                }
190                VisitResult::SkipChildren => {
191                    visitor.visit_node_post(n);
192                }
193                VisitResult::Stop => return,
194            },
195            RedTree::Leaf(t) => {
196                if let VisitResult::Stop = visitor.visit_token(t) {
197                    return;
198                }
199            }
200        }
201    }
202
203    /// Implementation of post-order traversal.
204    fn traverse_post_order_impl<'a, L: Language>(&self, node: RedTree<'a, L>, visitor: &mut impl Visitor<'a, L>) {
205        match node {
206            RedTree::Node(n) => {
207                if let VisitResult::Continue = visitor.visit_node_pre(n) {
208                    for child in n.children() {
209                        self.traverse_post_order_impl(child, visitor);
210                    }
211                }
212                if let VisitResult::Stop = visitor.visit_node_post(n) {
213                    return;
214                }
215            }
216            RedTree::Leaf(t) => {
217                if let VisitResult::Stop = visitor.visit_token(t) {
218                    return;
219                }
220            }
221        }
222    }
223}
224
225/// A visitor that finds nodes of a specific type.
226struct TypeFinderVisitor<'b, 'a, L: Language> {
227    node_type: L::ElementType,
228    result: &'b mut Vec<RedNode<'a, L>>,
229}
230
231impl<'b, 'a, L: Language> Visitor<'a, L> for TypeFinderVisitor<'b, 'a, L> {
232    fn visit_node_pre(&mut self, node: RedNode<'a, L>) -> VisitResult {
233        if node.element_type() == self.node_type {
234            self.result.push(node);
235        }
236        VisitResult::Continue
237    }
238
239    fn visit_node_post(&mut self, _node: RedNode<'a, L>) -> VisitResult {
240        VisitResult::Continue
241    }
242
243    fn visit_token(&mut self, _token: RedLeaf<L>) -> VisitResult {
244        VisitResult::Continue
245    }
246}
247
248/// A visitor that finds tokens of a specific type.
249struct TokenTypeFinderVisitor<'b, L: Language> {
250    token_type: L::TokenType,
251    result: &'b mut Vec<RedLeaf<L>>,
252}
253
254impl<'b, 'a, L: Language> Visitor<'a, L> for TokenTypeFinderVisitor<'b, L> {
255    fn visit_node_pre(&mut self, _node: RedNode<'a, L>) -> VisitResult {
256        VisitResult::Continue
257    }
258
259    fn visit_node_post(&mut self, _node: RedNode<'a, L>) -> VisitResult {
260        VisitResult::Continue
261    }
262
263    fn visit_token(&mut self, token: RedLeaf<L>) -> VisitResult {
264        if token.kind == self.token_type {
265            self.result.push(token);
266        }
267        VisitResult::Continue
268    }
269}
270
271/// Extension methods for RedTree to make it easier to work with the AST processor.
272pub trait RedTreeExt<'a, L: Language> {
273    /// Traverses the tree in pre-order.
274    fn traverse_pre_order(&self, visitor: &mut impl Visitor<'a, L>);
275
276    /// Traverses the tree in post-order.
277    fn traverse_post_order(&self, visitor: &mut impl Visitor<'a, L>);
278
279    /// Finds all nodes of the specified type.
280    fn find_nodes_by_type(&self, node_type: L::ElementType) -> Vec<RedNode<'a, L>>;
281
282    /// Finds all tokens of the specified type.
283    fn find_tokens_by_type(&self, token_type: L::TokenType) -> Vec<RedLeaf<L>>;
284
285    /// Finds the node at the specified offset.
286    fn find_node_at_offset(&self, offset: usize) -> Option<RedTree<'a, L>>;
287}
288
289impl<'a, L: Language> RedTreeExt<'a, L> for RedTree<'a, L> {
290    fn traverse_pre_order(&self, visitor: &mut impl Visitor<'a, L>) {
291        let processor = DefaultAstProcessor::new();
292        processor.traverse_pre_order(*self, visitor);
293    }
294
295    fn traverse_post_order(&self, visitor: &mut impl Visitor<'a, L>) {
296        let processor = DefaultAstProcessor::new();
297        processor.traverse_post_order(*self, visitor);
298    }
299
300    fn find_nodes_by_type(&self, node_type: L::ElementType) -> Vec<RedNode<'a, L>> {
301        let processor = DefaultAstProcessor::new();
302        processor.find_nodes_by_type(*self, node_type)
303    }
304
305    fn find_tokens_by_type(&self, token_type: L::TokenType) -> Vec<RedLeaf<L>> {
306        let processor = DefaultAstProcessor::new();
307        processor.find_tokens_by_type(*self, token_type)
308    }
309
310    fn find_node_at_offset(&self, offset: usize) -> Option<RedTree<'a, L>> {
311        let processor = DefaultAstProcessor::new();
312        processor.find_node_at_offset(*self, offset)
313    }
314}
oak_core/ast/mod.rs

oak_core/ast/
mod.rs
