use streaming_iterator::StreamingIterator;
use crate::analyzer::Severity;
use crate::language::Language;
use super::engine::ParsedFile;
pub struct QueryRule {
pub name: &'static str,
pub languages: &'static [Language],
pub pattern: &'static str,
pub severity: Severity,
pub handler: Option<QueryHandler>,
pub skips_test_files: bool,
}
pub type QueryHandler =
fn(file: &ParsedFile, captures: &[QueryCapture], match_index: usize) -> Vec<IssueCandidate>;
#[derive(Debug, Clone)]
pub struct QueryCapture<'a> {
pub name: String,
pub node: tree_sitter::Node<'a>,
pub text: &'a str,
}
#[derive(Debug, Clone)]
pub struct IssueCandidate {
pub line: usize,
pub column: usize,
pub message: String,
pub severity: Severity,
}
pub fn collect_captures<'a>(
file: &'a ParsedFile,
pattern: &str,
) -> Result<Vec<Vec<QueryCapture<'a>>>, String> {
let lang = file.language;
let grammar = super::parsers::get_grammar(lang).ok_or_else(|| {
format!(
"No tree-sitter grammar available for {}",
lang.display_name()
)
})?;
let query = tree_sitter::Query::new(&grammar, pattern)
.map_err(|e| format!("Failed to create query: {}", e))?;
let mut cursor = tree_sitter::QueryCursor::new();
let root = file.root_node();
let mut matches = cursor.matches(&query, root, file.content.as_bytes());
let capture_names: Vec<String> = query
.capture_names()
.iter()
.map(|s| s.to_string())
.collect();
let mut result = Vec::new();
while let Some(match_) = matches.next() {
let captures: Vec<QueryCapture> = match_
.captures
.iter()
.map(|capture| {
let name_idx = capture.index as usize;
let name = capture_names
.get(name_idx)
.cloned()
.unwrap_or_else(|| "unknown".to_string());
let node = capture.node;
let start = node.start_byte();
let end = node.end_byte();
QueryCapture {
name,
node,
text: &file.content[start..end],
}
})
.collect();
result.push(captures);
}
Ok(result)
}
pub fn run_query_rule(file: &ParsedFile, rule: &QueryRule) -> Vec<IssueCandidate> {
let captures_group = match collect_captures(file, rule.pattern) {
Ok(groups) => groups,
Err(e) => {
tracing::warn!("Query rule '{}' error: {}", rule.name, e);
return vec![];
}
};
let mut results = Vec::new();
for (match_index, captures) in captures_group.iter().enumerate() {
if let Some(handler) = rule.handler {
results.extend(handler(file, captures, match_index));
} else {
if let Some(first) = captures.first() {
results.push(IssueCandidate {
line: first.node.start_position().row + 1,
column: first.node.start_position().column + 1,
message: format!("{} detected", rule.name),
severity: rule.severity.clone(),
});
}
}
}
results
}
#[cfg(test)]
mod tests {
use super::*;
use crate::treesitter::TreeSitterEngine;
#[test]
fn test_collect_captures_basic() {
let engine = TreeSitterEngine::new();
let code = "fn main() { let x = 42; }";
let file = engine
.parse_file(std::path::Path::new("test.rs"), code)
.expect("Should parse");
let captures = collect_captures(&file, "(identifier) @id").expect("Query should succeed");
assert!(!captures.is_empty(), "Should find at least one identifier");
assert_eq!(captures.len(), 2, "Should find 2 identifiers: main, x");
}
#[test]
fn test_single_letter_variable_query() {
let engine = TreeSitterEngine::new();
let code = "fn compute() { let a = 1; let bb = 2; let ccc = 3; }";
let file = engine
.parse_file(std::path::Path::new("test.rs"), code)
.expect("Should parse");
let pattern = "
(let_declaration
pattern: (identifier) @var
(#match? @var \"^[a-z]$\")
)
";
let captures = collect_captures(&file, pattern).expect("Query should succeed");
assert_eq!(
captures.len(),
1,
"Only 'a' should match single-letter pattern"
);
if let Some(first) = captures.first().and_then(|c| c.first()) {
assert_eq!(first.text, "a", "Should capture 'a'");
}
}
#[test]
fn test_invalid_query_returns_error() {
let engine = TreeSitterEngine::new();
let code = "fn main() {}";
let file = engine
.parse_file(std::path::Path::new("test.rs"), code)
.expect("Should parse");
let result = collect_captures(&file, "(nonexistent_node) @x");
assert!(result.is_err(), "Query with unknown node type should error");
}
#[test]
fn test_query_rule_default_handler() {
let engine = TreeSitterEngine::new();
let code = "fn main() { let x = 1; let y = 2; }";
let file = engine
.parse_file(std::path::Path::new("test.rs"), code)
.expect("Should parse");
let rule = QueryRule {
name: "single-letter-var",
languages: &[Language::Rust],
pattern: "
(let_declaration
pattern: (identifier) @var
(#match? @var \"^[a-z]$\")
)
",
severity: Severity::Spicy,
handler: None,
skips_test_files: false,
};
let issues = run_query_rule(&file, &rule);
assert_eq!(issues.len(), 2, "Should find 2 single-letter variables");
assert_eq!(issues[0].message, "single-letter-var detected");
}
}