garbage_code_hunter/treesitter/

query.rs

use streaming_iterator::StreamingIterator;

use crate::analyzer::Severity;
use crate::language::Language;

use super::engine::ParsedFile;

/// A tree-sitter query based rule definition.
///
/// This is the primary building block for tree-sitter rules.
/// Each rule specifies a query pattern (in tree-sitter query syntax),
/// the languages it applies to, and a handler that converts
/// query matches into `CodeIssue`s.
pub struct QueryRule {
    /// Unique rule identifier (e.g. "single-letter-variable").
    pub name: &'static str,

    /// Languages this rule applies to.
    pub languages: &'static [Language],

    /// The tree-sitter query pattern string.
    /// Uses standard tree-sitter query syntax with named captures.
    pub pattern: &'static str,

    /// Default severity when match is found.
    pub severity: Severity,

    /// Custom handler to produce issues from a match.
    /// If None, a default handler is used (one issue per match at the capture node).
    pub handler: Option<QueryHandler>,

    /// Whether to skip test files.
    pub skips_test_files: bool,
}

/// Function signature for custom query match handlers.
pub type QueryHandler =
    fn(file: &ParsedFile, captures: &[QueryCapture], match_index: usize) -> Vec<IssueCandidate>;

/// A single named capture from a tree-sitter query match.
#[derive(Debug, Clone)]
pub struct QueryCapture<'a> {
    /// The capture name from the query pattern (e.g. "ident" from `(identifier) @ident`).
    pub name: String,

    /// The matched syntax node.
    pub node: tree_sitter::Node<'a>,

    /// Source text of the matched node.
    pub text: &'a str,
}

/// An issue candidate produced by a query rule handler.
#[derive(Debug, Clone)]
pub struct IssueCandidate {
    pub line: usize,
    pub column: usize,
    pub message: String,
    pub severity: Severity,
}

/// Execute a tree-sitter query against a parsed file and collect captures.
///
/// Returns a list of capture groups, one per query match.
/// Each group contains all named captures for that match.
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)
}

/// Run a `QueryRule` against a parsed file and produce issues.
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 {
            // Default handler: use the first capture's location
            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;

    /// Objective: Verify basic query matching works for Rust source
    /// Invariants: Query for identifiers should find all identifier nodes
    #[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");
        // Should find: main, x
        assert_eq!(captures.len(), 2, "Should find 2 identifiers: main, x");
    }

    /// Objective: Verify single-letter variable detection via query
    /// Invariants: Pattern matching single-character identifiers should catch them
    #[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");

        // Match `let` bindings with single-letter patterns
        let pattern = "
            (let_declaration
                pattern: (identifier) @var
                (#match? @var \"^[a-z]$\")
            )
        ";
        let captures = collect_captures(&file, pattern).expect("Query should succeed");

        // The `let` pattern: `let a = ...; let bb = ...; let ccc = ...;`
        // Only `a` should match (single letter)
        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'");
        }
    }

    /// Objective: Verify invalid query returns an error
    /// Invariants: Malformed query pattern should not panic
    #[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");
        // Unknown node type in query should be an error
        assert!(result.is_err(), "Query with unknown node type should error");
    }

    /// Objective: Verify QueryRule default handler produces issues
    /// Invariants: A QueryRule with no custom handler should still produce issues
    #[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");
    }
}