use tree_sitter::Node;
pub const MAX_TREE_DEPTH: usize = 512;
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct WalkStats {
pub visited: usize,
pub skipped_subtrees: usize,
}
impl WalkStats {
pub fn truncated(&self) -> bool {
self.skipped_subtrees > 0
}
}
pub fn walk<'t, F>(root: Node<'t>, visit: &mut F) -> WalkStats
where
F: FnMut(Node<'t>),
{
let mut stats = WalkStats::default();
let mut cursor = root.walk();
let mut depth = 0usize;
loop {
visit(cursor.node());
stats.visited += 1;
if depth < MAX_TREE_DEPTH {
if cursor.goto_first_child() {
depth += 1;
continue;
}
} else if cursor.node().child_count() > 0 {
stats.skipped_subtrees += 1;
}
loop {
if depth == 0 {
return stats;
}
if cursor.goto_next_sibling() {
break;
}
cursor.goto_parent();
depth -= 1;
}
}
}
pub fn node_text<'a>(node: Node<'_>, source: &'a [u8]) -> Option<&'a str> {
node.utf8_text(source).ok()
}
pub fn field_text<'a>(node: Node<'_>, field: &str, source: &'a [u8]) -> Option<&'a str> {
node_text(node.child_by_field_name(field)?, source)
}
pub fn start_line(node: Node<'_>) -> u32 {
node.start_position().row as u32 + 1
}
pub fn end_line(node: Node<'_>) -> u32 {
node.end_position().row as u32 + 1
}
pub fn first_line_of(node: Node<'_>, source: &[u8]) -> Option<String> {
let text = node_text(node, source)?;
Some(text.lines().next().unwrap_or("").trim_end().to_string())
}
pub fn has_parse_error(root: Node<'_>) -> bool {
root.has_error()
}
pub fn first_error_line(root: Node<'_>) -> Option<u32> {
if !root.has_error() {
return None;
}
let mut line = None;
walk(root, &mut |node| {
if line.is_none() && (node.is_error() || node.is_missing()) {
line = Some(start_line(node));
}
});
line
}
#[cfg(test)]
mod tests {
use super::*;
fn parse_rust(source: &str) -> tree_sitter::Tree {
let mut parser = tree_sitter::Parser::new();
parser
.set_language(&tree_sitter_rust::language())
.expect("rust grammar loads");
parser.parse(source, None).expect("parser returns a tree")
}
#[test]
fn walk_visits_in_source_order() {
let tree = parse_rust("struct A; struct B; struct C;");
let src = "struct A; struct B; struct C;".as_bytes();
let mut names = Vec::new();
walk(tree.root_node(), &mut |node| {
if node.kind() == "type_identifier" {
if let Some(t) = node_text(node, src) {
names.push(t.to_string());
}
}
});
assert_eq!(names, vec!["A", "B", "C"], "sibling order must be preserved");
}
#[test]
fn walk_covers_the_whole_subtree() {
let tree = parse_rust("fn f() { let x = Foo { a: 1 }; }");
let mut count = 0usize;
let stats = walk(tree.root_node(), &mut |_| count += 1);
assert_eq!(stats.visited, count);
assert!(count > 10, "a non-trivial function has many nodes");
assert!(!stats.truncated());
}
#[test]
fn walk_bounds_pathological_nesting_instead_of_overflowing() {
let depth = MAX_TREE_DEPTH * 4;
let source = format!("fn f() {{ let x = {}1{}; }}", "(".repeat(depth), ")".repeat(depth));
let tree = parse_rust(&source);
let stats = walk(tree.root_node(), &mut |_| {});
assert!(
stats.truncated(),
"input nests deeper than the limit, so truncation must be reported"
);
assert!(stats.visited >= MAX_TREE_DEPTH);
}
#[test]
fn walk_stays_inside_the_requested_subtree() {
let source = "struct A; struct B;";
let tree = parse_rust(source);
let first = tree.root_node().child(0).expect("first item exists");
let mut seen = Vec::new();
walk(first, &mut |node| {
if node.kind() == "type_identifier" {
if let Some(t) = node_text(node, source.as_bytes()) {
seen.push(t.to_string());
}
}
});
assert_eq!(seen, vec!["A"], "walking one item must not reach its sibling");
}
#[test]
fn parse_errors_are_detected_and_located() {
let tree = parse_rust("fn broken( {");
assert!(has_parse_error(tree.root_node()));
assert!(first_error_line(tree.root_node()).is_some());
let clean = parse_rust("fn ok() {}");
assert!(!has_parse_error(clean.root_node()));
assert_eq!(first_error_line(clean.root_node()), None);
}
}