Skip to content

Commit 1c6bc31

Browse files
author
Marlon Costa
committed
feat(node): Add utility methods for AST traversal
- Add has_ancestor() for checking if any ancestor matches a predicate - Add all_occurrences() to Search trait for finding all matching nodes
1 parent 37e5d83 commit 1c6bc31

2 files changed

Lines changed: 45 additions & 0 deletions

File tree

src/node.rs

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -183,6 +183,21 @@ impl<'a> Node<'a> {
183183
}
184184
res
185185
}
186+
187+
/// Checks if this node has any ancestor that meets the given predicate.
188+
///
189+
/// Traverses up the tree from this node's parent to the root,
190+
/// returning true if any ancestor satisfies the predicate.
191+
pub fn has_ancestor<F: Fn(&Node) -> bool>(&self, pred: F) -> bool {
192+
let mut node = *self;
193+
while let Some(parent) = node.parent() {
194+
if pred(&parent) {
195+
return true;
196+
}
197+
node = parent;
198+
}
199+
false
200+
}
186201
}
187202

188203
/// An `AST` cursor.
@@ -236,6 +251,35 @@ impl<'a> Search<'a> for Node<'a> {
236251
None
237252
}
238253

254+
fn all_occurrences(&self, pred: fn(u16) -> bool) -> Vec<Node<'a>> {
255+
let mut cursor = self.cursor();
256+
let mut stack = Vec::new();
257+
let mut children = Vec::new();
258+
let mut results = Vec::new();
259+
260+
stack.push(*self);
261+
262+
while let Some(node) = stack.pop() {
263+
if pred(node.kind_id()) {
264+
results.push(node);
265+
}
266+
cursor.reset(&node);
267+
if cursor.goto_first_child() {
268+
loop {
269+
children.push(cursor.node());
270+
if !cursor.goto_next_sibling() {
271+
break;
272+
}
273+
}
274+
for child in children.drain(..).rev() {
275+
stack.push(child);
276+
}
277+
}
278+
}
279+
280+
results
281+
}
282+
239283
fn act_on_node(&self, action: &mut dyn FnMut(&Node<'a>)) {
240284
let mut cursor = self.cursor();
241285
let mut stack = Vec::new();

src/traits.rs

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -68,6 +68,7 @@ pub trait ParserTrait {
6868

6969
pub(crate) trait Search<'a> {
7070
fn first_occurrence(&self, pred: fn(u16) -> bool) -> Option<Node<'a>>;
71+
fn all_occurrences(&self, pred: fn(u16) -> bool) -> Vec<Node<'a>>;
7172
fn act_on_node(&self, pred: &mut dyn FnMut(&Node<'a>));
7273
fn first_child(&self, pred: fn(u16) -> bool) -> Option<Node<'a>>;
7374
fn act_on_child(&self, action: &mut dyn FnMut(&Node<'a>));

0 commit comments

Comments
 (0)