use crate::code_tools::ToolError;
use agcodex_ast::Language;
use dashmap::DashMap;
use std::collections::HashMap;
use std::sync::Arc;
use thiserror::Error;
use tree_sitter::Query;
#[derive(Debug, Error)]
pub enum QueryError {
#[error("unsupported language: {language}")]
UnsupportedLanguage { language: String },
#[error("invalid query type: {query_type} for language {language}")]
InvalidQueryType {
query_type: String,
language: String,
},
#[error("query compilation failed: {details}")]
CompilationFailed { details: String },
#[error("template not found: {template_name}")]
TemplateNotFound { template_name: String },
#[error("query execution failed: {reason}")]
ExecutionFailed { reason: String },
}
impl From<QueryError> for ToolError {
fn from(err: QueryError) -> Self {
ToolError::InvalidQuery(err.to_string())
}
}
#[derive(Debug, Clone, Copy, Hash, PartialEq, Eq)]
pub enum QueryType {
Functions,
Classes,
Imports,
Symbols,
Methods,
Constructors,
Signatures,
Modules,
}
#[derive(Debug, Clone)]
pub struct CompiledQuery {
pub query: Arc<Query>,
pub language: Language,
pub query_type: QueryType,
pub description: String,
pub capture_names: Vec<String>,
}
#[derive(Debug, Clone, Hash, PartialEq, Eq)]
struct CacheKey {
language: Language,
query_type: QueryType,
variant: Option<String>, }
#[derive(Debug)]
pub struct QueryCache {
cache: DashMap<CacheKey, Arc<CompiledQuery>>,
stats: DashMap<String, u64>,
}
impl QueryCache {
pub fn new() -> Self {
Self {
cache: DashMap::new(),
stats: DashMap::new(),
}
}
pub fn get(
&self,
language: Language,
query_type: QueryType,
variant: Option<&str>,
) -> Option<Arc<CompiledQuery>> {
let key = CacheKey {
language,
query_type,
variant: variant.map(|s| s.to_string()),
};
if let Some(query) = self.cache.get(&key) {
self.increment_stat("cache_hits");
Some(query.clone())
} else {
self.increment_stat("cache_misses");
None
}
}
pub fn insert(&self, compiled: Arc<CompiledQuery>, variant: Option<&str>) {
let key = CacheKey {
language: compiled.language,
query_type: compiled.query_type,
variant: variant.map(|s| s.to_string()),
};
self.cache.insert(key, compiled);
self.increment_stat("cache_insertions");
}
pub fn stats(&self) -> HashMap<String, u64> {
self.stats
.iter()
.map(|entry| (entry.key().clone(), *entry.value()))
.collect()
}
pub fn clear(&self) {
self.cache.clear();
self.increment_stat("cache_clears");
}
pub fn size(&self) -> usize {
self.cache.len()
}
fn increment_stat(&self, stat_name: &str) {
*self.stats.entry(stat_name.to_string()).or_insert(0) += 1;
}
}
impl Default for QueryCache {
fn default() -> Self {
Self::new()
}
}
#[derive(Debug, Clone)]
pub struct QueryTemplates {
templates: HashMap<(Language, QueryType), &'static str>,
}
impl QueryTemplates {
pub fn new() -> Self {
let mut templates = HashMap::new();
Self::add_rust_templates(&mut templates);
Self::add_python_templates(&mut templates);
Self::add_javascript_templates(&mut templates);
Self::add_go_templates(&mut templates);
Self::add_java_templates(&mut templates);
Self::add_c_cpp_templates(&mut templates);
Self { templates }
}
pub fn get(&self, language: Language, query_type: &QueryType) -> Option<&'static str> {
self.templates.get(&(language, *query_type)).copied()
}
pub fn supported_queries(&self, language: Language) -> Vec<QueryType> {
self.templates
.keys()
.filter(|(lang, _)| *lang == language)
.map(|(_, query_type)| *query_type)
.collect()
}
fn add_rust_templates(templates: &mut HashMap<(Language, QueryType), &'static str>) {
templates.insert(
(Language::Rust, QueryType::Functions),
r#"
[
(function_item
name: (identifier) @name
parameters: (parameters) @params
body: (block)? @body) @function
(associated_item
(function_item
name: (identifier) @name
parameters: (parameters) @params
body: (block)? @body) @function)
]
"#,
);
templates.insert(
(Language::Rust, QueryType::Classes),
r#"
[
(struct_item
name: (type_identifier) @name
body: (field_declaration_list)? @body) @struct
(enum_item
name: (type_identifier) @name
body: (enum_variant_list) @body) @enum
(union_item
name: (type_identifier) @name
body: (field_declaration_list) @body) @union
(trait_item
name: (type_identifier) @name
body: (declaration_list) @body) @trait
(impl_item
trait: (type_identifier)? @trait_name
type: (type_identifier) @type_name
body: (declaration_list) @body) @impl
]
"#,
);
templates.insert(
(Language::Rust, QueryType::Imports),
r#"
[
(use_declaration
argument: (use_clause) @clause) @import
(extern_crate_declaration
name: (identifier) @name
rename: (identifier)? @rename) @extern_crate
(mod_item
name: (identifier) @name) @module
]
"#,
);
templates.insert(
(Language::Rust, QueryType::Symbols),
r#"
[
(let_declaration
pattern: (identifier) @name
type: (_)? @type
value: (_)? @value) @variable
(const_item
name: (identifier) @name
type: (_) @type
value: (_) @value) @constant
(static_item
name: (identifier) @name
type: (_) @type
value: (_) @value) @static
(type_item
name: (type_identifier) @name
type: (_) @type) @type_alias
]
"#,
);
templates.insert(
(Language::Rust, QueryType::Methods),
r#"
(impl_item
body: (declaration_list
(function_item
name: (identifier) @name
parameters: (parameters) @params
body: (block)? @body) @method))
"#,
);
}
fn add_python_templates(templates: &mut HashMap<(Language, QueryType), &'static str>) {
templates.insert(
(Language::Python, QueryType::Functions),
r#"
[
(function_definition
name: (identifier) @name
parameters: (parameters) @params
body: (block) @body) @function
(async_function_definition
name: (identifier) @name
parameters: (parameters) @params
body: (block) @body) @async_function
]
"#,
);
templates.insert(
(Language::Python, QueryType::Classes),
r#"
(class_definition
name: (identifier) @name
superclasses: (argument_list)? @superclasses
body: (block) @body) @class
"#,
);
templates.insert(
(Language::Python, QueryType::Imports),
r#"
[
(import_statement
name: (dotted_name) @module) @import
(import_from_statement
module_name: (dotted_name)? @module
name: (dotted_name) @name) @import_from
(import_from_statement
module_name: (dotted_name)? @module
name: (aliased_import) @name) @import_from_alias
]
"#,
);
templates.insert(
(Language::Python, QueryType::Symbols),
r#"
[
(assignment
left: (identifier) @name
right: (_) @value) @variable
(assignment
left: (pattern_list) @names
right: (_) @value) @multiple_assignment
(augmented_assignment
left: (identifier) @name
right: (_) @value) @augmented_variable
]
"#,
);
templates.insert(
(Language::Python, QueryType::Methods),
r#"
(class_definition
body: (block
(function_definition
name: (identifier) @name
parameters: (parameters) @params
body: (block) @body) @method))
"#,
);
}
fn add_javascript_templates(templates: &mut HashMap<(Language, QueryType), &'static str>) {
let languages = [Language::JavaScript, Language::TypeScript];
for &lang in &languages {
templates.insert(
(lang, QueryType::Functions),
r#"
[
(function_declaration
name: (identifier) @name
parameters: (formal_parameters) @params
body: (statement_block) @body) @function
(function_expression
name: (identifier)? @name
parameters: (formal_parameters) @params
body: (statement_block) @body) @function_expr
(arrow_function
parameters: (formal_parameters) @params
body: (_) @body) @arrow_function
(generator_function_declaration
name: (identifier) @name
parameters: (formal_parameters) @params
body: (statement_block) @body) @generator
]
"#,
);
templates.insert(
(lang, QueryType::Classes),
r#"
[
(class_declaration
name: (identifier) @name
superclass: (class_heritage)? @superclass
body: (class_body) @body) @class
(class_expression
name: (identifier)? @name
superclass: (class_heritage)? @superclass
body: (class_body) @body) @class_expr
]
"#,
);
templates.insert(
(lang, QueryType::Imports),
r#"
[
(import_statement
source: (string) @source
(import_clause
(named_imports) @named)?) @import
(import_statement
source: (string) @source
(import_clause
(namespace_import) @namespace)?) @import_namespace
(export_statement) @export
]
"#,
);
templates.insert(
(lang, QueryType::Methods),
r#"
(class_body
(method_definition
name: (property_name) @name
parameters: (formal_parameters) @params
body: (statement_block) @body) @method)
"#,
);
templates.insert(
(lang, QueryType::Symbols),
r#"
[
(variable_declaration
(variable_declarator
name: (identifier) @name
value: (_)? @value)) @variable
(lexical_declaration
(variable_declarator
name: (identifier) @name
value: (_)? @value)) @lexical_variable
]
"#,
);
}
}
fn add_go_templates(templates: &mut HashMap<(Language, QueryType), &'static str>) {
templates.insert(
(Language::Go, QueryType::Functions),
r#"
[
(function_declaration
name: (identifier) @name
parameters: (parameter_list) @params
result: (_)? @return_type
body: (block) @body) @function
(method_declaration
receiver: (parameter_list) @receiver
name: (identifier) @name
parameters: (parameter_list) @params
result: (_)? @return_type
body: (block) @body) @method
]
"#,
);
templates.insert(
(Language::Go, QueryType::Classes),
r#"
[
(type_declaration
(type_spec
name: (type_identifier) @name
type: (struct_type) @struct_body)) @struct
(type_declaration
(type_spec
name: (type_identifier) @name
type: (interface_type) @interface_body)) @interface
]
"#,
);
templates.insert(
(Language::Go, QueryType::Imports),
r#"
[
(import_declaration
(import_spec
name: (package_identifier)? @alias
path: (interpreted_string_literal) @path)) @import
(package_clause
(package_identifier) @package_name) @package
]
"#,
);
templates.insert(
(Language::Go, QueryType::Symbols),
r#"
[
(var_declaration
(var_spec
name: (identifier) @name
type: (_)? @type
value: (_)? @value)) @variable
(const_declaration
(const_spec
name: (identifier) @name
type: (_)? @type
value: (_) @value)) @constant
]
"#,
);
}
fn add_java_templates(templates: &mut HashMap<(Language, QueryType), &'static str>) {
templates.insert(
(Language::Java, QueryType::Functions),
r#"
(method_declaration
(modifiers)? @modifiers
type: (_) @return_type
name: (identifier) @name
parameters: (formal_parameters) @params
body: (block)? @body) @method
"#,
);
templates.insert(
(Language::Java, QueryType::Classes),
r#"
[
(class_declaration
(modifiers)? @modifiers
name: (identifier) @name
superclass: (superclass)? @superclass
interfaces: (super_interfaces)? @interfaces
body: (class_body) @body) @class
(interface_declaration
(modifiers)? @modifiers
name: (identifier) @name
extends: (extends_interfaces)? @extends
body: (interface_body) @body) @interface
(enum_declaration
(modifiers)? @modifiers
name: (identifier) @name
interfaces: (super_interfaces)? @interfaces
body: (enum_body) @body) @enum
]
"#,
);
templates.insert(
(Language::Java, QueryType::Imports),
r#"
[
(import_declaration
(scoped_identifier) @import_path) @import
(import_declaration
(asterisk) @wildcard) @wildcard_import
(package_declaration
(scoped_identifier) @package_name) @package
]
"#,
);
templates.insert(
(Language::Java, QueryType::Symbols),
r#"
[
(field_declaration
(modifiers)? @modifiers
type: (_) @type
(variable_declarator
name: (identifier) @name
value: (_)? @value)) @field
(local_variable_declaration
type: (_) @type
(variable_declarator
name: (identifier) @name
value: (_)? @value)) @local_variable
]
"#,
);
templates.insert(
(Language::Java, QueryType::Constructors),
r#"
(constructor_declaration
(modifiers)? @modifiers
name: (identifier) @name
parameters: (formal_parameters) @params
body: (constructor_body) @body) @constructor
"#,
);
}
fn add_c_cpp_templates(templates: &mut HashMap<(Language, QueryType), &'static str>) {
let languages = [Language::C, Language::Cpp];
for &lang in &languages {
templates.insert(
(lang, QueryType::Functions),
r#"
[
(function_definition
type: (_) @return_type
declarator: (function_declarator
declarator: (_) @name
parameters: (parameter_list) @params)
body: (compound_statement) @body) @function
(declaration
type: (_) @return_type
declarator: (function_declarator
declarator: (_) @name
parameters: (parameter_list) @params)) @function_declaration
]
"#,
);
templates.insert(
(lang, QueryType::Classes),
r#"
[
(struct_specifier
name: (type_identifier) @name
body: (field_declaration_list) @body) @struct
(union_specifier
name: (type_identifier) @name
body: (field_declaration_list) @body) @union
(enum_specifier
name: (type_identifier) @name
body: (enumerator_list) @body) @enum
]
"#,
);
templates.insert(
(lang, QueryType::Imports),
r#"
[
(preproc_include
path: (string_literal) @path) @include
(preproc_include
path: (system_lib_string) @path) @include_system
]
"#,
);
templates.insert(
(lang, QueryType::Symbols),
r#"
[
(declaration
type: (_) @type
declarator: (identifier) @name) @variable
(init_declarator
declarator: (identifier) @name
value: (_) @value) @initialized_variable
(preproc_def
name: (identifier) @name
value: (_)? @value) @macro
]
"#,
);
}
templates.insert(
(Language::Cpp, QueryType::Classes),
r#"
[
(class_specifier
name: (type_identifier) @name
base: (base_class_clause)? @base_classes
body: (field_declaration_list) @body) @class
(struct_specifier
name: (type_identifier) @name
base: (base_class_clause)? @base_classes
body: (field_declaration_list) @body) @struct
(union_specifier
name: (type_identifier) @name
body: (field_declaration_list) @body) @union
(enum_specifier
name: (type_identifier) @name
body: (enumerator_list) @body) @enum
]
"#,
);
templates.insert(
(Language::Cpp, QueryType::Imports),
r#"
[
(preproc_include
path: (string_literal) @path) @include
(preproc_include
path: (system_lib_string) @path) @include_system
(using_declaration
(qualified_identifier) @name) @using
(namespace_definition
name: (identifier) @name
body: (declaration_list) @body) @namespace
]
"#,
);
}
}
impl Default for QueryTemplates {
fn default() -> Self {
Self::new()
}
}
#[derive(Debug, Clone)]
pub struct QueryBuilder {
templates: Arc<QueryTemplates>,
cache: Arc<QueryCache>,
}
impl QueryBuilder {
pub fn new() -> Self {
Self {
templates: Arc::new(QueryTemplates::new()),
cache: Arc::new(QueryCache::new()),
}
}
pub fn with_cache(cache: Arc<QueryCache>) -> Self {
Self {
templates: Arc::new(QueryTemplates::new()),
cache,
}
}
pub fn build_query(
&self,
language: Language,
query_type: QueryType,
custom_template: Option<&str>,
) -> Result<Arc<CompiledQuery>, QueryError> {
if let Some(cached) = self.cache.get(language, query_type, None) {
return Ok(cached);
}
let template = if let Some(custom) = custom_template {
custom
} else {
self.templates.get(language, &query_type).ok_or_else(|| {
QueryError::InvalidQueryType {
query_type: format!("{:?}", query_type),
language: language.name().to_string(),
}
})?
};
let ts_language = language.parser();
let query =
Query::new(&ts_language, template).map_err(|e| QueryError::CompilationFailed {
details: format!("Tree-sitter compilation error: {}", e),
})?;
let capture_names = query
.capture_names()
.iter()
.map(|&s| s.to_string())
.collect();
let compiled = Arc::new(CompiledQuery {
query: Arc::new(query),
language,
query_type,
description: format!("{:?} queries for {}", query_type, language.name()),
capture_names,
});
self.cache.insert(compiled.clone(), None);
Ok(compiled)
}
pub fn build_custom_query(
&self,
language: Language,
query_type: QueryType,
template: &str,
variant: &str,
) -> Result<Arc<CompiledQuery>, QueryError> {
if let Some(cached) = self.cache.get(language, query_type, Some(variant)) {
return Ok(cached);
}
let ts_language = language.parser();
let query =
Query::new(&ts_language, template).map_err(|e| QueryError::CompilationFailed {
details: format!("Custom query compilation error: {}", e),
})?;
let capture_names = query
.capture_names()
.iter()
.map(|&s| s.to_string())
.collect();
let compiled = Arc::new(CompiledQuery {
query: Arc::new(query),
language,
query_type,
description: format!(
"Custom {:?} query ({}) for {}",
query_type,
variant,
language.name()
),
capture_names,
});
self.cache.insert(compiled.clone(), Some(variant));
Ok(compiled)
}
pub fn supported_queries(&self, language: Language) -> Vec<QueryType> {
self.templates.supported_queries(language)
}
pub fn cache_stats(&self) -> HashMap<String, u64> {
self.cache.stats()
}
pub fn clear_cache(&self) {
self.cache.clear();
}
}
impl Default for QueryBuilder {
fn default() -> Self {
Self::new()
}
}
#[derive(Debug, Clone)]
pub struct QueryLibrary {
builder: QueryBuilder,
cache: Arc<QueryCache>,
templates: Arc<QueryTemplates>,
}
impl QueryLibrary {
pub fn new() -> Self {
let cache = Arc::new(QueryCache::new());
Self {
builder: QueryBuilder::with_cache(cache.clone()),
cache,
templates: Arc::new(QueryTemplates::new()),
}
}
pub fn get_query(
&self,
language: Language,
query_type: QueryType,
) -> Result<Arc<CompiledQuery>, QueryError> {
self.builder.build_query(language, query_type, None)
}
pub fn get_custom_query(
&self,
language: Language,
query_type: QueryType,
template: &str,
variant: &str,
) -> Result<Arc<CompiledQuery>, QueryError> {
self.builder
.build_custom_query(language, query_type, template, variant)
}
pub fn supports_query(&self, language: Language, query_type: &QueryType) -> bool {
self.templates.get(language, query_type).is_some()
}
pub fn supported_languages(&self) -> Vec<Language> {
vec![
Language::Rust,
Language::Python,
Language::JavaScript,
Language::TypeScript,
Language::Go,
Language::Java,
Language::C,
Language::Cpp,
]
}
pub fn supported_query_types(&self, language: Language) -> Vec<QueryType> {
self.builder.supported_queries(language)
}
pub fn stats(&self) -> QueryLibraryStats {
let cache_stats = self.cache.stats();
QueryLibraryStats {
cache_size: self.cache.size(),
cache_hits: cache_stats.get("cache_hits").copied().unwrap_or(0),
cache_misses: cache_stats.get("cache_misses").copied().unwrap_or(0),
total_queries: cache_stats.get("cache_insertions").copied().unwrap_or(0),
supported_languages: self.supported_languages().len(),
}
}
pub fn clear_cache(&self) {
self.cache.clear();
}
pub fn precompile_language(&self, language: Language) -> Result<usize, QueryError> {
let query_types = self.supported_query_types(language);
let mut compiled_count = 0;
for query_type in query_types {
match self.get_query(language, query_type) {
Ok(_) => compiled_count += 1,
Err(e) => {
eprintln!(
"Warning: Failed to precompile {:?} for {}: {}",
query_type,
language.name(),
e
);
}
}
}
Ok(compiled_count)
}
pub fn precompile_all(&self) -> Result<usize, QueryError> {
let languages = self.supported_languages();
let mut total_compiled = 0;
for language in languages {
match self.precompile_language(language) {
Ok(count) => total_compiled += count,
Err(e) => {
eprintln!(
"Warning: Failed to precompile queries for {}: {}",
language.name(),
e
);
}
}
}
Ok(total_compiled)
}
}
impl Default for QueryLibrary {
fn default() -> Self {
Self::new()
}
}
#[derive(Debug, Clone)]
pub struct QueryLibraryStats {
pub cache_size: usize,
pub cache_hits: u64,
pub cache_misses: u64,
pub total_queries: u64,
pub supported_languages: usize,
}
impl QueryLibraryStats {
pub fn hit_rate(&self) -> f64 {
let total_requests = self.cache_hits + self.cache_misses;
if total_requests == 0 {
0.0
} else {
self.cache_hits as f64 / total_requests as f64
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_query_cache_basic_operations() {
let cache = QueryCache::new();
assert_eq!(cache.size(), 0);
assert!(
cache
.get(Language::Rust, QueryType::Functions, None)
.is_none()
);
let stats = cache.stats();
assert_eq!(stats.get("cache_misses"), Some(&1));
}
#[test]
fn test_query_templates_rust() {
let templates = QueryTemplates::new();
assert!(
templates
.get(Language::Rust, &QueryType::Functions)
.is_some()
);
assert!(templates.get(Language::Rust, &QueryType::Classes).is_some());
assert!(templates.get(Language::Rust, &QueryType::Imports).is_some());
let supported = templates.supported_queries(Language::Rust);
assert!(supported.contains(&QueryType::Functions));
assert!(supported.contains(&QueryType::Classes));
assert!(supported.contains(&QueryType::Imports));
}
#[test]
fn test_query_library_creation() {
let library = QueryLibrary::new();
let languages = library.supported_languages();
assert!(languages.contains(&Language::Rust));
assert!(languages.contains(&Language::Python));
assert!(languages.contains(&Language::JavaScript));
let stats = library.stats();
assert_eq!(stats.cache_size, 0);
assert_eq!(stats.cache_hits, 0);
assert_eq!(stats.cache_misses, 0);
}
#[test]
fn test_query_library_support_check() {
let library = QueryLibrary::new();
assert!(library.supports_query(Language::Rust, &QueryType::Functions));
assert!(library.supports_query(Language::Python, &QueryType::Classes));
assert!(library.supports_query(Language::JavaScript, &QueryType::Imports));
let rust_queries = library.supported_query_types(Language::Rust);
assert!(!rust_queries.is_empty());
assert!(rust_queries.contains(&QueryType::Functions));
}
#[test]
fn test_stats_hit_rate_calculation() {
let stats = QueryLibraryStats {
cache_size: 10,
cache_hits: 8,
cache_misses: 2,
total_queries: 10,
supported_languages: 5,
};
assert_eq!(stats.hit_rate(), 0.8);
let empty_stats = QueryLibraryStats {
cache_size: 0,
cache_hits: 0,
cache_misses: 0,
total_queries: 0,
supported_languages: 0,
};
assert_eq!(empty_stats.hit_rate(), 0.0);
}
}