Skip to main content

codetether_rlm/oracle/tree_sitter_oracle/
oracle.rs

1use anyhow::{Result, anyhow};
2
3/// Tree-sitter based oracle for validating structural queries.
4pub struct TreeSitterOracle {
5    pub(crate) source: String,
6    pub(crate) tree: Option<tree_sitter::Tree>,
7    pub(crate) parser: Option<tree_sitter::Parser>,
8}
9
10impl TreeSitterOracle {
11    /// Create a new tree-sitter oracle for the given source.
12    pub fn new(source: String) -> Self {
13        Self {
14            source,
15            tree: None,
16            parser: None,
17        }
18    }
19
20    pub(crate) fn ensure_parser(&mut self) -> Result<()> {
21        if self.parser.is_some() {
22            return Ok(());
23        }
24
25        let mut parser = tree_sitter::Parser::new();
26        parser.set_language(&tree_sitter_rust::LANGUAGE.into())?;
27        self.parser = Some(parser);
28        Ok(())
29    }
30
31    pub(crate) fn parse(&mut self) -> Result<&tree_sitter::Tree> {
32        self.ensure_parser()?;
33        if self.tree.is_none() {
34            let parser = self
35                .parser
36                .as_mut()
37                .ok_or_else(|| anyhow!("Parser not initialized"))?;
38            self.tree = Some(parse_source(parser, &self.source)?);
39        }
40        self.tree
41            .as_ref()
42            .ok_or_else(|| anyhow!("Failed to retain parsed tree"))
43    }
44}
45
46fn parse_source(parser: &mut tree_sitter::Parser, source: &str) -> Result<tree_sitter::Tree> {
47    parser
48        .parse(source, None)
49        .ok_or_else(|| anyhow!("Failed to parse source"))
50}