codetether_rlm/oracle/tree_sitter_oracle/
query.rs1use anyhow::Result;
2use streaming_iterator::StreamingIterator;
3
4use super::{AstMatch, AstQueryResult, TreeSitterOracle, query_match};
5
6impl TreeSitterOracle {
7 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}