diff --git a/src/node.rs b/src/node.rs index 49b8df104..f6a014706 100644 --- a/src/node.rs +++ b/src/node.rs @@ -183,6 +183,27 @@ impl<'a> Node<'a> { } res } + + /// Checks if this node has any ancestor that meets the given predicate. + /// + /// This method walks up the tree from this node's parent to the root, + /// returning `true` as soon as it finds any ancestor for which `pred` + /// returns `true`. + /// + /// Note: This differs from [`Node::has_ancestors`], which checks for a + /// specific pattern of ancestors (immediate parent and grandparent + /// satisfying two separate predicates). `has_ancestor` instead uses a + /// single predicate and considers all ancestors in the hierarchy. + pub fn has_ancestor bool>(&self, pred: F) -> bool { + let mut node = *self; + while let Some(parent) = node.parent() { + if pred(&parent) { + return true; + } + node = parent; + } + false + } } /// An `AST` cursor. @@ -236,6 +257,37 @@ impl<'a> Search<'a> for Node<'a> { None } + /// Performs a depth-first search starting from this node (including `self`) + /// and returns all descendant nodes whose `kind_id` satisfies the given predicate. + fn all_occurrences(&self, pred: fn(u16) -> bool) -> Vec> { + let mut cursor = self.cursor(); + let mut stack = Vec::new(); + let mut children = Vec::new(); + let mut results = Vec::new(); + + stack.push(*self); + + while let Some(node) = stack.pop() { + if pred(node.kind_id()) { + results.push(node); + } + cursor.reset(&node); + if cursor.goto_first_child() { + loop { + children.push(cursor.node()); + if !cursor.goto_next_sibling() { + break; + } + } + for child in children.drain(..).rev() { + stack.push(child); + } + } + } + + results + } + fn act_on_node(&self, action: &mut dyn FnMut(&Node<'a>)) { let mut cursor = self.cursor(); let mut stack = Vec::new(); diff --git a/src/traits.rs b/src/traits.rs index 16d4ed9cb..d04692263 100644 --- a/src/traits.rs +++ b/src/traits.rs @@ -68,6 +68,7 @@ pub trait ParserTrait { pub(crate) trait Search<'a> { fn first_occurrence(&self, pred: fn(u16) -> bool) -> Option>; + fn all_occurrences(&self, pred: fn(u16) -> bool) -> Vec>; fn act_on_node(&self, pred: &mut dyn FnMut(&Node<'a>)); fn first_child(&self, pred: fn(u16) -> bool) -> Option>; fn act_on_child(&self, action: &mut dyn FnMut(&Node<'a>));