use crate::core::semantic::parser_pool::ParserPoolManager;
use crate::utils::error::ContextCreatorError;
use moka::future::Cache;
use std::collections::hash_map::DefaultHasher;
use std::hash::{Hash, Hasher};
use std::path::{Path, PathBuf};
use std::sync::Arc;
use std::time::Duration;
use tokio::time::timeout;
use tree_sitter::Tree;
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
struct CacheKey {
path: PathBuf,
content_hash: u64,
language: String,
}
impl CacheKey {
fn new(path: &Path, content: &str, language: &str) -> Self {
let mut hasher = DefaultHasher::new();
content.hash(&mut hasher);
let content_hash = hasher.finish();
Self {
path: path.to_path_buf(),
content_hash,
language: language.to_string(),
}
}
}
#[derive(Clone)]
struct CacheEntry {
tree: Arc<Tree>,
content: Arc<String>,
}
#[derive(Clone)]
pub struct AstCacheV2 {
cache: Cache<CacheKey, CacheEntry>,
parser_pool: Arc<ParserPoolManager>,
parse_timeout: Duration,
}
impl AstCacheV2 {
pub fn new(capacity: u64) -> Self {
let cache = Cache::builder()
.max_capacity(capacity)
.time_to_live(Duration::from_secs(3600)) .build();
Self {
cache,
parser_pool: Arc::new(ParserPoolManager::new()),
parse_timeout: Duration::from_secs(30), }
}
pub fn with_config(capacity: u64, ttl: Duration, parse_timeout: Duration) -> Self {
let cache = Cache::builder()
.max_capacity(capacity)
.time_to_live(ttl)
.build();
Self {
cache,
parser_pool: Arc::new(ParserPoolManager::new()),
parse_timeout,
}
}
pub async fn get_or_parse(
&self,
path: &Path,
content: &str,
language: &str,
) -> Result<Arc<Tree>, ContextCreatorError> {
let key = CacheKey::new(path, content, language);
let parser_pool = self.parser_pool.clone();
let content_clone = content.to_string();
let language_clone = language.to_string();
let path_clone = path.to_path_buf();
let timeout_duration = self.parse_timeout;
let entry = self
.cache
.try_get_with(key, async move {
let parse_result = timeout(timeout_duration, async {
let mut parser = parser_pool.get_parser(&language_clone).await?;
let tree = parser.parse(&content_clone, None).ok_or_else(|| {
ContextCreatorError::ParseError(format!(
"Failed to parse {} file: {}",
language_clone,
path_clone.display()
))
})?;
Ok::<Tree, ContextCreatorError>(tree)
})
.await;
match parse_result {
Ok(Ok(tree)) => Ok(CacheEntry {
tree: Arc::new(tree),
content: Arc::new(content_clone),
}),
Ok(Err(e)) => Err(e),
Err(_) => Err(ContextCreatorError::ParseError(format!(
"Parsing timed out after {:?} for file: {}",
timeout_duration,
path_clone.display()
))),
}
})
.await
.map_err(|e| {
ContextCreatorError::ParseError(format!("Failed to cache parse result: {e}"))
})?;
Ok(entry.tree.clone())
}
pub async fn get_content(
&self,
path: &Path,
content_hash: &str,
language: &str,
) -> Option<Arc<String>> {
let mut hasher = DefaultHasher::new();
content_hash.hash(&mut hasher);
let hash = hasher.finish();
let key = CacheKey {
path: path.to_path_buf(),
content_hash: hash,
language: language.to_string(),
};
self.cache
.get(&key)
.await
.map(|entry| entry.content.clone())
}
pub async fn clear(&self) {
self.cache.invalidate_all();
}
pub fn len(&self) -> u64 {
self.cache.entry_count()
}
pub fn is_empty(&self) -> bool {
self.cache.entry_count() == 0
}
pub fn stats(&self) -> CacheStats {
CacheStats {
hits: 0, misses: 0,
evictions: 0,
entry_count: self.cache.entry_count(),
}
}
}
#[derive(Debug, Clone)]
pub struct CacheStats {
pub hits: u64,
pub misses: u64,
pub evictions: u64,
pub entry_count: u64,
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn test_cache_operations() {
let cache = AstCacheV2::new(10);
assert_eq!(cache.len(), 0);
assert!(cache.is_empty());
let path = Path::new("test.rs");
let content = "fn main() {}";
let result = cache.get_or_parse(path, content, "rust").await;
assert!(result.is_ok());
tokio::time::sleep(tokio::time::Duration::from_millis(10)).await;
let result2 = cache.get_or_parse(path, content, "rust").await;
assert!(result2.is_ok());
cache.clear().await;
tokio::time::sleep(tokio::time::Duration::from_millis(10)).await;
}
#[tokio::test]
async fn test_cache_hit() {
let cache = AstCacheV2::new(10);
let path = Path::new("test.py");
let content = "def test(): pass";
let result1 = cache.get_or_parse(path, content, "python").await;
assert!(result1.is_ok());
let result2 = cache.get_or_parse(path, content, "python").await;
assert!(result2.is_ok());
assert!(Arc::ptr_eq(&result1.unwrap(), &result2.unwrap()));
}
#[tokio::test]
async fn test_cache_invalidation_on_content_change() {
let cache = AstCacheV2::new(10);
let path = Path::new("test.js");
let content1 = "function test() {}";
let content2 = "function test2() {}";
let result1 = cache.get_or_parse(path, content1, "javascript").await;
assert!(result1.is_ok());
let result2 = cache.get_or_parse(path, content2, "javascript").await;
assert!(result2.is_ok());
assert!(!Arc::ptr_eq(&result1.unwrap(), &result2.unwrap()));
}
#[tokio::test]
async fn test_concurrent_parsing() {
let cache = Arc::new(AstCacheV2::new(100));
let mut handles = vec![];
for _i in 0..10 {
let cache_clone = cache.clone();
let handle = tokio::spawn(async move {
let path = Path::new("concurrent.rs");
let content = "fn main() { println!(\"test\"); }";
cache_clone.get_or_parse(path, content, "rust").await
});
handles.push(handle);
}
for handle in handles {
assert!(handle.await.unwrap().is_ok());
}
tokio::time::sleep(tokio::time::Duration::from_millis(50)).await;
assert!(cache.len() <= 10); }
#[tokio::test]
async fn test_timeout_configuration() {
let cache =
AstCacheV2::with_config(10, Duration::from_secs(3600), Duration::from_millis(1));
let path = Path::new("test.rs");
let content = "fn main() {}";
let result = cache.get_or_parse(path, content, "rust").await;
assert!(result.is_ok());
}
}