use std::ops::Range;
use tree_sitter::{Node, Parser};
use crate::lang::{Lang, parsable};
#[derive(Debug, Clone, Default)]
pub struct Classified {
pub comments: Vec<Range<usize>>,
pub imports: Vec<Range<usize>>,
pub logging: Vec<Range<usize>>,
pub tests: Vec<Range<usize>>,
pub error_line: Option<usize>,
}
pub fn classify(lang: Lang, src: &str) -> Option<Classified> {
let mut parser = Parser::new();
parser.set_language(&lang.grammar()).ok()?;
let tree = parser.parse(parsable(lang, src).as_ref(), None)?;
let mut classified = Classified::default();
let mut cursor = tree.walk();
'walk: loop {
let node = cursor.node();
let test_block = (lang == Lang::Rust)
.then(|| rust_test_block(node, src))
.flatten();
let layer = if let Some(block) = &test_block {
classified.tests.push(block.clone());
None
} else if lang.is_comment(node.kind()) || (lang == Lang::Python && is_docstring(node)) {
Some(&mut classified.comments)
} else if is_import(lang, node) {
Some(&mut classified.imports)
} else if is_logging(lang, node, src) {
Some(&mut classified.logging)
} else {
None
};
if let Some(layer) = layer {
layer.push(node.byte_range());
} else if test_block.is_none() && cursor.goto_first_child() {
continue;
}
while !cursor.goto_next_sibling() {
if !cursor.goto_parent() {
break 'walk;
}
}
}
classified.error_line = first_error(tree.root_node()).map(|node| node.start_position().row + 1);
Some(classified)
}
fn is_docstring(node: Node) -> bool {
if node.kind() != "expression_statement" || node.named_child_count() != 1 {
return false;
}
let Some(expr) = node.named_child(0) else {
return false;
};
if !matches!(expr.kind(), "string" | "concatenated_string") {
return false;
}
let Some(parent) = node.parent() else {
return false;
};
let is_body = match parent.kind() {
"module" => true,
"block" => parent.parent().is_some_and(|owner| {
matches!(owner.kind(), "function_definition" | "class_definition")
}),
_ => false,
};
if !is_body {
return false;
}
let mut cursor = parent.walk();
let first_statement = parent
.named_children(&mut cursor)
.find(|child| child.kind() != "comment");
first_statement.is_some_and(|first| first.id() == node.id())
}
fn is_import(lang: Lang, node: Node) -> bool {
match lang {
Lang::Swift | Lang::Go => node.kind() == "import_declaration",
Lang::TypeScript | Lang::Tsx | Lang::JavaScript => {
node.kind() == "import_statement"
|| (node.kind() == "export_statement"
&& node.child_by_field_name("source").is_some())
}
Lang::Python => matches!(
node.kind(),
"import_statement" | "import_from_statement" | "future_import_statement"
),
Lang::Rust => matches!(node.kind(), "use_declaration" | "extern_crate_declaration"),
Lang::Kotlin => node.kind() == "import",
}
}
fn rust_test_block(node: Node, src: &str) -> Option<Range<usize>> {
if !matches!(node.kind(), "mod_item" | "function_item") {
return None;
}
let mut start = None;
let mut is_test = false;
let mut sibling = node.prev_sibling();
while let Some(previous) = sibling {
match previous.kind() {
"attribute_item" => {
start = Some(previous.start_byte());
is_test |= src
.get(previous.byte_range())
.is_some_and(is_test_attribute);
}
"line_comment" | "block_comment" => {}
_ => break,
}
sibling = previous.prev_sibling();
}
is_test.then(|| start.unwrap_or(node.start_byte())..node.end_byte())
}
fn is_test_attribute(attribute: &str) -> bool {
let inner = attribute
.trim()
.trim_start_matches("#[")
.trim_end_matches(']')
.trim();
let words = |text: &str| -> Vec<String> {
text.split(|c: char| !(c.is_alphanumeric() || c == '_'))
.filter(|word| !word.is_empty())
.map(str::to_string)
.collect()
};
if inner.starts_with("cfg") {
return words(inner).iter().any(|word| word == "test");
}
let path = inner.split('(').next().unwrap_or(inner);
let last = path.rsplit("::").next().unwrap_or(path).trim();
last == "test" || last == "rstest" || last.starts_with("test_")
}
fn is_logging(lang: Lang, node: Node, src: &str) -> bool {
let call = match lang {
Lang::Swift | Lang::Kotlin => (node.kind() == "call_expression"
&& node
.parent()
.is_some_and(|p| matches!(p.kind(), "statements" | "source_file" | "block")))
.then_some(node),
_ => (node.kind() == "expression_statement")
.then(|| node.named_child(0))
.flatten()
.filter(|child| {
matches!(
child.kind(),
"call_expression" | "call" | "macro_invocation"
)
}),
};
let Some(call) = call else {
return false;
};
let callee = match (lang, call.kind()) {
(Lang::Swift | Lang::Kotlin, _) => call.named_child(0),
(_, "macro_invocation") => call.child_by_field_name("macro"),
_ => call.child_by_field_name("function"),
};
callee
.and_then(|callee| src.get(callee.byte_range()))
.is_some_and(|callee| is_logging_callee(lang, callee))
}
fn is_logging_callee(lang: Lang, callee: &str) -> bool {
let callee: String = callee.chars().filter(|c| !c.is_whitespace()).collect();
let segments: Vec<&str> = callee
.split(['.', '?', '!', ':'])
.filter(|s| !s.is_empty())
.collect();
let Some((&method, receivers)) = segments.split_last() else {
return false;
};
if receivers.is_empty() {
return match lang {
Lang::Swift => matches!(method, "print" | "debugPrint" | "NSLog" | "os_log" | "dump"),
Lang::Python => method == "print",
Lang::Go | Lang::Kotlin => matches!(method, "print" | "println"),
Lang::Rust => matches!(
method,
"println"
| "eprintln"
| "print"
| "eprint"
| "dbg"
| "trace"
| "debug"
| "info"
| "warn"
| "error"
),
_ => false,
};
}
let receiver = receivers
.last()
.map(|receiver| receiver.to_ascii_lowercase())
.unwrap_or_default();
let level = method.to_ascii_lowercase();
match lang {
Lang::Go if receiver == "fmt" => matches!(method, "Print" | "Println" | "Printf"),
Lang::Go if receiver == "log" => method.starts_with("Print"),
Lang::Go if level.starts_with("fatal") || level.starts_with("panic") => false,
Lang::Kotlin if receiver == "log" || receiver == "timber" => {
matches!(method, "d" | "i" | "w" | "e" | "v" | "wtf") || is_level(&level)
}
_ => is_level(&level) && receivers.iter().any(|receiver| is_logger(receiver)),
}
}
fn is_level(level: &str) -> bool {
const LEVELS: [&str; 13] = [
"log",
"debug",
"info",
"notice",
"warn",
"warning",
"error",
"fault",
"trace",
"critical",
"exception",
"verbose",
"fatal",
];
LEVELS.contains(&level)
|| level
.strip_suffix(['f', 'w'])
.is_some_and(|stem| LEVELS.contains(&stem))
}
fn is_logger(receiver: &str) -> bool {
let receiver = receiver.to_ascii_lowercase();
matches!(
receiver.as_str(),
"console" | "logging" | "log" | "tracing" | "slog" | "timber"
) || receiver.ends_with("logger")
}
fn first_error(root: Node) -> Option<Node> {
if !root.has_error() {
return None;
}
let mut node = root;
'descend: loop {
let mut cursor = node.walk();
for child in node.children(&mut cursor) {
if child.has_error() || child.is_missing() {
node = child;
continue 'descend;
}
}
return Some(node);
}
}