agentroot_core/index/ast_chunker/
parser.rs

1//! Tree-sitter parser wrapper
2
3use super::language::Language;
4use crate::error::{Error, Result};
5use tree_sitter::{Language as TsLanguage, Parser, Tree};
6
7/// Parse source code into a tree-sitter AST
8pub fn parse(source: &str, language: Language) -> Result<Tree> {
9    let mut parser = Parser::new();
10    let ts_language = get_tree_sitter_language(language);
11    parser
12        .set_language(&ts_language)
13        .map_err(|e| Error::Parse(e.to_string()))?;
14    parser
15        .parse(source, None)
16        .ok_or_else(|| Error::Parse("Failed to parse source".to_string()))
17}
18
19/// Get the tree-sitter language for a Language enum variant.
20/// This is infallible since all Language variants have corresponding tree-sitter languages.
21fn get_tree_sitter_language(language: Language) -> TsLanguage {
22    match language {
23        Language::Rust => tree_sitter_rust::LANGUAGE.into(),
24        Language::Python => tree_sitter_python::LANGUAGE.into(),
25        Language::JavaScript => tree_sitter_javascript::LANGUAGE.into(),
26        Language::TypeScript => tree_sitter_typescript::LANGUAGE_TYPESCRIPT.into(),
27        Language::TypeScriptTsx => tree_sitter_typescript::LANGUAGE_TSX.into(),
28        Language::Go => tree_sitter_go::LANGUAGE.into(),
29    }
30}
31
32#[cfg(test)]
33mod tests {
34    use super::*;
35
36    #[test]
37    fn test_parse_rust() {
38        let source = "fn main() { println!(\"Hello\"); }";
39        let tree = parse(source, Language::Rust).unwrap();
40        assert_eq!(tree.root_node().kind(), "source_file");
41    }
42
43    #[test]
44    fn test_parse_python() {
45        let source = "def main():\n    print('Hello')";
46        let tree = parse(source, Language::Python).unwrap();
47        assert_eq!(tree.root_node().kind(), "module");
48    }
49
50    #[test]
51    fn test_parse_javascript() {
52        let source = "function main() { console.log('Hello'); }";
53        let tree = parse(source, Language::JavaScript).unwrap();
54        assert_eq!(tree.root_node().kind(), "program");
55    }
56
57    #[test]
58    fn test_parse_typescript() {
59        let source = "function main(): void { console.log('Hello'); }";
60        let tree = parse(source, Language::TypeScript).unwrap();
61        assert_eq!(tree.root_node().kind(), "program");
62    }
63
64    #[test]
65    fn test_parse_go() {
66        let source = "package main\n\nfunc main() { fmt.Println(\"Hello\") }";
67        let tree = parse(source, Language::Go).unwrap();
68        assert_eq!(tree.root_node().kind(), "source_file");
69    }
70}
agentroot_core/index/ast_chunker/parser.rs

agentroot_core/index/ast_chunker/
parser.rs
