1use std::collections::{HashMap, HashSet};
4use std::path::PathBuf;
5use std::time::Duration;
6
7#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
9pub enum RefreshMode {
10 Manual,
12 Always,
14 Files,
16 #[default]
18 Auto,
19}
20
21#[derive(Debug, Clone, PartialEq, Eq, Hash)]
27pub enum MapCacheKey {
28 Files(FilesKey),
30 Auto(AutoKey),
32}
33
34#[derive(Debug, Clone, PartialEq, Eq, Hash)]
36pub struct FilesKey {
37 pub chat_fnames: Option<Vec<PathBuf>>,
39 pub other_fnames: Option<Vec<PathBuf>>,
41 pub max_tokens: usize,
43}
44
45#[derive(Debug, Clone, PartialEq, Eq, Hash)]
47pub struct AutoKey {
48 pub chat_fnames: Option<Vec<PathBuf>>,
50 pub other_fnames: Option<Vec<PathBuf>>,
52 pub max_tokens: usize,
54 pub mentioned_fnames: Option<Vec<String>>,
56 pub mentioned_idents: Option<Vec<String>>,
58 pub anchor_fnames: Option<Vec<PathBuf>>,
60 pub anchor_idents: Option<Vec<String>>,
62 pub anchor_scoped: Option<Vec<(PathBuf, String)>>,
64}
65
66pub struct AnchorCacheParams<'a> {
69 pub anchor_fnames: &'a [PathBuf],
71 pub anchor_idents: &'a HashSet<String>,
73 pub anchor_scoped: &'a [(PathBuf, String)],
75}
76
77impl MapCacheKey {
78 pub fn files(chat_fnames: &[PathBuf], other_fnames: &[PathBuf], max_tokens: usize) -> Self {
80 MapCacheKey::Files(FilesKey {
81 chat_fnames: non_empty_sorted(chat_fnames),
82 other_fnames: non_empty_sorted(other_fnames),
83 max_tokens,
84 })
85 }
86
87 pub fn auto(
89 chat_fnames: &[PathBuf],
90 other_fnames: &[PathBuf],
91 max_tokens: usize,
92 mentioned_fnames: &HashSet<String>,
93 mentioned_idents: &HashSet<String>,
94 anchors: AnchorCacheParams<'_>,
95 ) -> Self {
96 MapCacheKey::Auto(AutoKey {
97 chat_fnames: non_empty_sorted(chat_fnames),
98 other_fnames: non_empty_sorted(other_fnames),
99 max_tokens,
100 mentioned_fnames: non_empty_sorted_set(mentioned_fnames),
101 mentioned_idents: non_empty_sorted_set(mentioned_idents),
102 anchor_fnames: non_empty_sorted(anchors.anchor_fnames),
103 anchor_idents: non_empty_sorted_set(anchors.anchor_idents),
104 anchor_scoped: non_empty_sorted_pairs(anchors.anchor_scoped),
105 })
106 }
107}
108
109fn non_empty_sorted<T: Clone + Ord>(items: &[T]) -> Option<Vec<T>> {
111 if items.is_empty() {
112 None
113 } else {
114 let mut v: Vec<_> = items.to_vec();
115 v.sort();
116 Some(v)
117 }
118}
119
120fn non_empty_sorted_pairs<A: Clone + Ord, B: Clone + Ord>(items: &[(A, B)]) -> Option<Vec<(A, B)>> {
122 if items.is_empty() {
123 None
124 } else {
125 let mut v: Vec<_> = items.to_vec();
126 v.sort();
127 Some(v)
128 }
129}
130
131fn non_empty_sorted_set<T: Clone + Ord>(items: &HashSet<T>) -> Option<Vec<T>> {
133 if items.is_empty() {
134 None
135 } else {
136 let mut v: Vec<_> = items.iter().cloned().collect();
137 v.sort();
138 Some(v)
139 }
140}
141
142#[derive(Debug, Default)]
144pub struct MapCache {
145 cache: HashMap<MapCacheKey, String>,
147 last_map: Option<String>,
149 last_duration: Duration,
151}
152
153impl MapCache {
154 pub fn new() -> Self {
156 Self::default()
157 }
158
159 pub fn should_use_cache(&self, mode: RefreshMode, force_refresh: bool) -> bool {
163 if force_refresh {
164 return false;
165 }
166
167 match mode {
168 RefreshMode::Manual => self.last_map.is_some(),
169 RefreshMode::Always => false,
170 RefreshMode::Files => true,
171 RefreshMode::Auto => self.last_duration > Duration::from_secs_f64(1.0),
172 }
173 }
174
175 pub fn get(
179 &self,
180 key: &MapCacheKey,
181 mode: RefreshMode,
182 force_refresh: bool,
183 ) -> Option<&String> {
184 if !self.should_use_cache(mode, force_refresh) {
185 return None;
186 }
187
188 if mode == RefreshMode::Manual {
190 return self.last_map.as_ref();
191 }
192
193 self.cache.get(key)
194 }
195
196 pub fn get_last_map(&self) -> Option<&String> {
198 self.last_map.as_ref()
199 }
200
201 pub fn set(&mut self, key: MapCacheKey, value: String, duration: Duration) {
206 self.cache.insert(key, value.clone());
207 self.last_map = Some(value);
208 self.last_duration = duration;
209 }
210
211 pub fn last_duration(&self) -> Duration {
213 self.last_duration
214 }
215
216 pub fn clear(&mut self) {
218 self.cache.clear();
219 self.last_map = None;
220 self.last_duration = Duration::ZERO;
221 }
222}
223
224#[cfg(test)]
225mod tests {
226 use super::*;
227
228 #[test]
229 fn cache_key_files_mode() {
230 let key = MapCacheKey::files(
231 &[PathBuf::from("a.rs"), PathBuf::from("b.rs")],
232 &[PathBuf::from("c.rs")],
233 1024,
234 );
235
236 if let MapCacheKey::Files(k) = &key {
237 assert_eq!(k.max_tokens, 1024);
238 let chat = k.chat_fnames.as_ref().unwrap();
240 assert_eq!(chat[0], PathBuf::from("a.rs"));
241 assert_eq!(chat[1], PathBuf::from("b.rs"));
242 } else {
243 panic!("Expected Files key");
244 }
245 }
246
247 #[test]
248 fn cache_key_empty_is_none() {
249 let key = MapCacheKey::files(&[], &[PathBuf::from("c.rs")], 1024);
250
251 if let MapCacheKey::Files(k) = &key {
252 assert!(k.chat_fnames.is_none()); assert!(k.other_fnames.is_some());
254 } else {
255 panic!("Expected Files key");
256 }
257 }
258
259 #[test]
260 fn cache_set_and_get() {
261 let mut cache = MapCache::new();
262 let key = MapCacheKey::files(&[PathBuf::from("a.rs")], &[], 1024);
263
264 cache.set(key.clone(), "test map".to_string(), Duration::from_secs(2));
265
266 let result = cache.get(&key, RefreshMode::Files, false);
267 assert!(result.is_some());
268 assert_eq!(result.unwrap(), "test map");
269 }
270
271 #[test]
272 fn cache_manual_mode() {
273 let mut cache = MapCache::new();
274 let key = MapCacheKey::files(&[PathBuf::from("a.rs")], &[], 1024);
275
276 cache.set(key.clone(), "test map".to_string(), Duration::from_secs(1));
277
278 let other_key = MapCacheKey::files(&[PathBuf::from("b.rs")], &[], 2048);
280 let result = cache.get(&other_key, RefreshMode::Manual, false);
281 assert!(result.is_some());
282 assert_eq!(result.unwrap(), "test map");
283 }
284
285 #[test]
286 fn cache_always_mode() {
287 let mut cache = MapCache::new();
288 let key = MapCacheKey::files(&[PathBuf::from("a.rs")], &[], 1024);
289
290 cache.set(key.clone(), "test map".to_string(), Duration::from_secs(1));
291
292 let result = cache.get(&key, RefreshMode::Always, false);
294 assert!(result.is_none());
295 }
296
297 #[test]
298 fn cache_auto_mode_threshold() {
299 let mut cache = MapCache::new();
300 let key = MapCacheKey::auto(
301 &[PathBuf::from("a.rs")],
302 &[],
303 1024,
304 &HashSet::new(),
305 &HashSet::new(),
306 AnchorCacheParams {
307 anchor_fnames: &[],
308 anchor_idents: &HashSet::new(),
309 anchor_scoped: &[],
310 },
311 );
312
313 cache.set(
315 key.clone(),
316 "fast map".to_string(),
317 Duration::from_millis(500),
318 );
319 let result = cache.get(&key, RefreshMode::Auto, false);
320 assert!(result.is_none());
321
322 cache.set(key.clone(), "slow map".to_string(), Duration::from_secs(2));
324 let result = cache.get(&key, RefreshMode::Auto, false);
325 assert!(result.is_some());
326 }
327
328 #[test]
329 fn cache_force_refresh() {
330 let mut cache = MapCache::new();
331 let key = MapCacheKey::files(&[PathBuf::from("a.rs")], &[], 1024);
332
333 cache.set(key.clone(), "test map".to_string(), Duration::from_secs(2));
334
335 let result = cache.get(&key, RefreshMode::Files, true);
337 assert!(result.is_none());
338 }
339
340 #[test]
341 fn cache_always_writes() {
342 let mut cache = MapCache::new();
343 let key = MapCacheKey::files(&[PathBuf::from("a.rs")], &[], 1024);
344
345 cache.set(key.clone(), "test map".to_string(), Duration::from_secs(1));
347
348 let result = cache.get(&key, RefreshMode::Files, false);
350 assert!(result.is_some());
351 }
352}