Skip to main content

codetether_rlm/oracle/tree_sitter_oracle/
query.rs

1use anyhow::Result;
2use streaming_iterator::StreamingIterator;
3
4use super::{AstMatch, AstQueryResult, TreeSitterOracle, query_match};
5
6impl TreeSitterOracle {
7    /// Execute a tree-sitter S-expression query.
8    pub fn query(&mut self, query_str: &str) -> Result<AstQueryResult> {
9        self.parse()?;
10        let tree = self.tree.as_ref().expect("tree parsed");
11        let query = tree_sitter::Query::new(&tree_sitter_rust::LANGUAGE.into(), query_str)?;
12        let matches = collect_matches(&query, tree.root_node(), self.source.as_bytes())?;
13
14        Ok(AstQueryResult {
15            query_type: query_str.to_string(),
16            matches,
17        })
18    }
19}
20
21fn collect_matches(
22    query: &tree_sitter::Query,
23    root: tree_sitter::Node<'_>,
24    source: &[u8],
25) -> Result<Vec<AstMatch>> {
26    let mut cursor = tree_sitter::QueryCursor::new();
27    let mut stream = cursor.matches(query, root, source);
28    let mut results = Vec::new();
29    while let Some(match_) = stream.next() {
30        results.push(query_match::build(query, match_, source)?);
31    }
32    Ok(results)
33}