Skip to main content

tree_sitter_language_pack/
query.rs

1use std::borrow::Cow;
2
3use crate::Error;
4use crate::node::{NodeInfo, node_info_from_node};
5use tree_sitter::StreamingIterator;
6
7/// A single match from a tree-sitter query, with captured nodes.
8#[derive(Debug, Clone)]
9pub struct QueryMatch {
10    /// The pattern index that matched (position in the query string).
11    pub pattern_index: usize,
12    /// Captures: list of (capture_name, node_info) pairs.
13    pub captures: Vec<(Cow<'static, str>, NodeInfo)>,
14}
15
16/// Execute a tree-sitter query pattern against a parsed tree.
17///
18/// The `query_source` is an S-expression pattern like:
19/// ```text
20/// (function_definition name: (identifier) @name)
21/// ```
22///
23/// Returns all matches with their captured nodes.
24///
25/// # Arguments
26///
27/// * `tree` - The parsed syntax tree to query.
28/// * `language` - Language name (used to compile the query pattern).
29/// * `query_source` - The tree-sitter query pattern string.
30/// * `source` - The original source code bytes (needed for capture resolution).
31///
32/// # Examples
33///
34/// ```no_run
35/// let tree = tree_sitter_language_pack::parse::parse_string("python", b"def hello(): pass").unwrap();
36/// let matches = tree_sitter_language_pack::query::run_query(
37///     &tree,
38///     "python",
39///     "(function_definition name: (identifier) @fn_name)",
40///     b"def hello(): pass",
41/// ).unwrap();
42/// assert!(!matches.is_empty());
43/// ```
44pub fn run_query(
45    tree: &tree_sitter::Tree,
46    language: &str,
47    query_source: &str,
48    source: &[u8],
49) -> Result<Vec<QueryMatch>, Error> {
50    let lang = crate::get_language(language)?;
51    let query = tree_sitter::Query::new(&lang, query_source).map_err(|e| Error::QueryError(format!("{e}")))?;
52    let capture_names: Vec<Cow<'static, str>> = query
53        .capture_names()
54        .iter()
55        .map(|s| Cow::Owned(s.to_string()))
56        .collect();
57
58    let mut cursor = tree_sitter::QueryCursor::new();
59    let mut matches = cursor.matches(&query, tree.root_node(), source);
60
61    // Tree-sitter 0.26+ evaluates standard text predicates (`#eq?`, `#not-eq?`,
62    // `#match?`, `#not-match?`, `#any-of?`, `#not-any-of?`) internally via
63    // `satisfies_text_predicates()` during `QueryCursor::matches()` iteration.
64    // The `general_predicates()` method only returns predicates with operators
65    // that tree-sitter does NOT recognize (i.e., custom predicates). Since we
66    // don't define any custom predicates, no additional filtering is needed.
67    let mut results = Vec::new();
68    while let Some(m) = matches.next() {
69        let captures = m
70            .captures
71            .iter()
72            .map(|c| {
73                let name = capture_names[c.index as usize].clone();
74                let info = node_info_from_node(c.node);
75                (name, info)
76            })
77            .collect();
78        results.push(QueryMatch {
79            pattern_index: m.pattern_index,
80            captures,
81        });
82    }
83    Ok(results)
84}
85
86#[cfg(test)]
87mod tests {
88    use super::*;
89
90    #[test]
91    fn test_run_query_invalid_language() {
92        // Create a dummy tree from any available language
93        let langs = crate::available_languages();
94        if langs.is_empty() {
95            return;
96        }
97        let tree = crate::parse::parse_string(&langs[0], b"x").unwrap();
98        let result = run_query(&tree, "nonexistent_xyz", "(identifier) @id", b"x");
99        assert!(result.is_err());
100    }
101
102    #[test]
103    fn test_run_query_invalid_pattern() {
104        let langs = crate::available_languages();
105        if langs.is_empty() {
106            return;
107        }
108        let first = &langs[0];
109        let tree = crate::parse::parse_string(first, b"x").unwrap();
110        let result = run_query(&tree, first, "((((invalid syntax", b"x");
111        assert!(result.is_err());
112    }
113
114    #[test]
115    fn test_run_query_no_matches() {
116        let langs = crate::available_languages();
117        if langs.is_empty() {
118            return;
119        }
120        let first = &langs[0];
121        let tree = crate::parse::parse_string(first, b"x").unwrap();
122        // Query for a node type that is unlikely to exist for a single "x"
123        let result = run_query(&tree, first, "(function_definition) @fn", b"x");
124        // This might error if the grammar doesn't have function_definition,
125        // or return empty matches. Either is acceptable.
126        if let Ok(matches) = result {
127            assert!(matches.is_empty());
128        }
129        // Query compilation error is fine for some grammars
130    }
131}