Skip to main content

vtcode_core/code/code_completion/cache/
mod.rs

1use super::engine::CompletionSuggestion;
2use hashbrown::HashMap;
3use std::sync::Arc;
4use std::time::{Duration, Instant};
5
6/// Completion cache for performance optimization
7pub struct CompletionCache {
8    cache: HashMap<String, CacheEntry>,
9    max_entries: usize,
10    ttl: Duration,
11}
12
13#[derive(Debug, Clone)]
14struct CacheEntry {
15    suggestions: Arc<Vec<CompletionSuggestion>>,
16    created_at: Instant,
17    access_count: usize,
18}
19
20impl CompletionCache {
21    pub fn new() -> Self {
22        Self {
23            cache: HashMap::new(),
24            max_entries: 1000,
25            ttl: Duration::from_secs(300), // 5 minutes
26        }
27    }
28
29    /// Get cached suggestions for a context
30    pub fn get(&mut self, context_key: &str) -> Option<Vec<CompletionSuggestion>> {
31        if let Some(entry) = self.cache.get_mut(context_key) {
32            if entry.created_at.elapsed() < self.ttl {
33                entry.access_count += 1;
34                return Some((*entry.suggestions).clone());
35            } else {
36                self.cache.remove(context_key);
37            }
38        }
39        None
40    }
41
42    /// Get a shared Arc reference to the suggestions, avoiding a Vec clone.
43    pub fn get_shared(&mut self, context_key: &str) -> Option<Arc<Vec<CompletionSuggestion>>> {
44        if let Some(entry) = self.cache.get_mut(context_key) {
45            if entry.created_at.elapsed() < self.ttl {
46                entry.access_count += 1;
47                return Some(Arc::clone(&entry.suggestions));
48            } else {
49                self.cache.remove(context_key);
50            }
51        }
52        None
53    }
54
55    /// Cache suggestions for a context
56    pub fn put(&mut self, context_key: String, suggestions: Vec<CompletionSuggestion>) {
57        // Clean up expired entries
58        self.cleanup_expired();
59
60        // Remove least recently used entries if at capacity
61        if self.cache.len() >= self.max_entries {
62            self.evict_lru();
63        }
64
65        let entry = CacheEntry {
66            suggestions: Arc::new(suggestions),
67            created_at: Instant::now(),
68            access_count: 1,
69        };
70
71        self.cache.insert(context_key, entry);
72    }
73
74    fn cleanup_expired(&mut self) {
75        let now = Instant::now();
76        self.cache
77            .retain(|_, entry| now.duration_since(entry.created_at) < self.ttl);
78    }
79
80    fn evict_lru(&mut self) {
81        let lru_key = self
82            .cache
83            .iter()
84            .min_by_key(|(_, entry)| entry.access_count)
85            .map(|(key, _)| key.clone());
86
87        if let Some(key) = lru_key {
88            self.cache.remove(&key);
89        }
90    }
91}
92
93impl Default for CompletionCache {
94    fn default() -> Self {
95        Self::new()
96    }
97}