use std::collections::hash_map::Entry;
use std::collections::HashMap;
use tree_sitter::{Parser, Query};
use crate::language::SupportedLanguage;
pub struct ParserPool {
parsers: HashMap<SupportedLanguage, Parser>,
queries: HashMap<SupportedLanguage, Query>,
}
impl Default for ParserPool {
fn default() -> Self {
Self::new()
}
}
impl ParserPool {
pub fn new() -> Self {
Self {
parsers: HashMap::new(),
queries: HashMap::new(),
}
}
pub fn get_or_create_query(&mut self, lang: SupportedLanguage) -> Option<&Query> {
if let Entry::Vacant(e) = self.queries.entry(lang) {
let ts_language = lang.ts_language();
let query_source = lang.highlights_query();
let query = Query::new(&ts_language, query_source).ok()?;
e.insert(query);
}
self.queries.get(&lang)
}
pub fn get_or_create(&mut self, ext: &str) -> Option<&mut Parser> {
let lang = SupportedLanguage::from_extension(ext)?;
if let Entry::Vacant(e) = self.parsers.entry(lang) {
let ts_language = lang.ts_language();
let mut parser = Parser::new();
if parser.set_language(&ts_language).is_err() {
return None;
}
e.insert(parser);
}
self.parsers.get_mut(&lang)
}
pub fn supports_extension(ext: &str) -> bool {
SupportedLanguage::is_supported(ext)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_parser_pool_rust() {
let mut pool = ParserPool::new();
let parser = pool.get_or_create("rs");
assert!(parser.is_some(), "Should create Rust parser");
let parser2 = pool.get_or_create("rs");
assert!(parser2.is_some(), "Should reuse Rust parser");
}
#[test]
fn test_parser_pool_typescript() {
let mut pool = ParserPool::new();
assert!(pool.get_or_create("ts").is_some());
assert!(pool.get_or_create("tsx").is_some());
}
#[test]
fn test_parser_pool_javascript() {
let mut pool = ParserPool::new();
assert!(pool.get_or_create("js").is_some());
assert!(pool.get_or_create("jsx").is_some());
}
#[test]
fn test_parser_pool_go() {
let mut pool = ParserPool::new();
assert!(pool.get_or_create("go").is_some());
}
#[test]
fn test_parser_pool_python() {
let mut pool = ParserPool::new();
assert!(pool.get_or_create("py").is_some());
}
#[test]
fn test_parser_pool_ruby() {
let mut pool = ParserPool::new();
assert!(pool.get_or_create("rb").is_some());
assert!(pool.get_or_create("rake").is_some());
assert!(pool.get_or_create("gemspec").is_some());
}
#[test]
fn test_parser_pool_zig() {
let mut pool = ParserPool::new();
assert!(pool.get_or_create("zig").is_some());
}
#[test]
fn test_parser_pool_c() {
let mut pool = ParserPool::new();
assert!(pool.get_or_create("c").is_some());
assert!(pool.get_or_create("h").is_some());
}
#[test]
fn test_parser_pool_cpp() {
let mut pool = ParserPool::new();
assert!(pool.get_or_create("cpp").is_some());
assert!(pool.get_or_create("cc").is_some());
assert!(pool.get_or_create("cxx").is_some());
assert!(pool.get_or_create("hpp").is_some());
assert!(pool.get_or_create("hxx").is_some());
}
#[test]
fn test_parser_pool_java() {
let mut pool = ParserPool::new();
assert!(pool.get_or_create("java").is_some());
}
#[test]
fn test_parser_pool_csharp() {
let mut pool = ParserPool::new();
assert!(pool.get_or_create("cs").is_some());
}
#[test]
fn test_parser_pool_lua() {
let mut pool = ParserPool::new();
assert!(pool.get_or_create("lua").is_some());
}
#[test]
fn test_parser_pool_bash() {
let mut pool = ParserPool::new();
assert!(pool.get_or_create("sh").is_some());
assert!(pool.get_or_create("bash").is_some());
assert!(pool.get_or_create("zsh").is_some());
}
#[test]
fn test_parser_pool_php() {
let mut pool = ParserPool::new();
assert!(pool.get_or_create("php").is_some());
}
#[test]
fn test_parser_pool_swift() {
let mut pool = ParserPool::new();
assert!(pool.get_or_create("swift").is_some());
}
#[test]
fn test_parser_pool_haskell() {
let mut pool = ParserPool::new();
assert!(pool.get_or_create("hs").is_some());
assert!(pool.get_or_create("lhs").is_some());
}
#[test]
fn test_parser_pool_moonbit() {
let mut pool = ParserPool::new();
assert!(pool.get_or_create("mbt").is_some());
}
#[test]
fn test_parser_pool_svelte() {
let mut pool = ParserPool::new();
assert!(pool.get_or_create("svelte").is_some());
}
#[test]
fn test_parser_pool_vue() {
let mut pool = ParserPool::new();
assert!(pool.get_or_create("vue").is_some());
}
#[test]
fn test_parser_pool_markdown() {
let mut pool = ParserPool::new();
assert!(pool.get_or_create("md").is_some());
assert!(pool.get_or_create("markdown").is_some());
}
#[test]
fn test_parser_pool_markdown_inline() {
let mut pool = ParserPool::new();
assert!(pool.get_or_create("md_inline").is_some());
}
#[test]
fn test_parser_pool_unsupported() {
let mut pool = ParserPool::new();
assert!(pool.get_or_create("yaml").is_none());
assert!(pool.get_or_create("toml").is_none());
}
#[test]
fn test_supports_extension() {
assert!(ParserPool::supports_extension("rs"));
assert!(ParserPool::supports_extension("ts"));
assert!(ParserPool::supports_extension("tsx"));
assert!(ParserPool::supports_extension("js"));
assert!(ParserPool::supports_extension("jsx"));
assert!(ParserPool::supports_extension("go"));
assert!(ParserPool::supports_extension("py"));
assert!(ParserPool::supports_extension("lua"));
assert!(ParserPool::supports_extension("sh"));
assert!(ParserPool::supports_extension("php"));
assert!(ParserPool::supports_extension("swift"));
assert!(ParserPool::supports_extension("hs"));
assert!(ParserPool::supports_extension("svelte"));
assert!(ParserPool::supports_extension("vue"));
assert!(ParserPool::supports_extension("md"));
assert!(!ParserPool::supports_extension("yaml"));
}
#[test]
fn test_parser_can_parse() {
let mut pool = ParserPool::new();
let parser = pool.get_or_create("rs").unwrap();
let code = "fn main() { println!(\"Hello\"); }";
let tree = parser.parse(code, None);
assert!(tree.is_some(), "Should parse Rust code");
}
#[test]
fn test_query_cache_creates_query() {
let mut pool = ParserPool::new();
let query = pool.get_or_create_query(SupportedLanguage::Rust);
assert!(query.is_some(), "Should create Rust highlight query");
}
#[test]
fn test_query_cache_reuses_query() {
let mut pool = ParserPool::new();
let query1 = pool.get_or_create_query(SupportedLanguage::TypeScript);
assert!(query1.is_some());
let query2 = pool.get_or_create_query(SupportedLanguage::TypeScript);
assert!(query2.is_some());
assert!(
pool.queries.contains_key(&SupportedLanguage::TypeScript),
"Query should be cached"
);
}
#[test]
fn test_query_cache_multiple_languages() {
let mut pool = ParserPool::new();
assert!(pool.get_or_create_query(SupportedLanguage::Rust).is_some());
assert!(pool
.get_or_create_query(SupportedLanguage::TypeScript)
.is_some());
assert!(pool
.get_or_create_query(SupportedLanguage::JavaScript)
.is_some());
assert!(pool.get_or_create_query(SupportedLanguage::Css).is_some());
assert_eq!(pool.queries.len(), 4);
}
#[test]
fn test_query_has_capture_names() {
let mut pool = ParserPool::new();
let query = pool.get_or_create_query(SupportedLanguage::Rust).unwrap();
let capture_names = query.capture_names();
assert!(!capture_names.is_empty(), "Query should have capture names");
}
#[test]
fn test_query_cache_markdown() {
let mut pool = ParserPool::new();
let query = pool.get_or_create_query(SupportedLanguage::Markdown);
assert!(
query.is_some(),
"Should compile Markdown block highlight query"
);
}
#[test]
fn test_query_cache_markdown_inline() {
let mut pool = ParserPool::new();
let query = pool.get_or_create_query(SupportedLanguage::MarkdownInline);
assert!(
query.is_some(),
"Should compile MarkdownInline highlight query"
);
}
#[test]
fn test_markdown_parser_can_parse() {
let mut pool = ParserPool::new();
let parser = pool.get_or_create("md").unwrap();
let code = "# Hello\n\nSome **bold** text.\n";
let tree = parser.parse(code, None);
assert!(tree.is_some(), "Should parse Markdown code");
}
#[test]
fn test_supports_extension_markdown() {
assert!(ParserPool::supports_extension("md"));
assert!(ParserPool::supports_extension("markdown"));
assert!(ParserPool::supports_extension("md_inline"));
}
}