1use std::collections::HashMap;
26use std::fs::Metadata;
27use std::path::{Path, PathBuf};
28use std::sync::{Arc, Mutex, OnceLock};
29use std::time::UNIX_EPOCH;
30
31const DEFAULT_BUDGET_MB: usize = 128;
33
34#[derive(Debug, Clone, Copy, PartialEq, Eq)]
37pub struct FileState {
38 pub mtime_ms: u64,
39 pub size_bytes: u64,
40}
41
42impl FileState {
43 pub fn from_metadata(meta: &Metadata) -> Option<Self> {
47 let mtime_ms = meta
48 .modified()
49 .ok()
50 .and_then(|t| t.duration_since(UNIX_EPOCH).ok())
51 .map(|d| d.as_millis() as u64)?;
52 Some(Self {
53 mtime_ms,
54 size_bytes: meta.len(),
55 })
56 }
57
58 pub fn from_path(path: &Path) -> Option<Self> {
60 Self::from_metadata(&path.metadata().ok()?)
61 }
62}
63
64struct Entry {
65 state: FileState,
66 content: Arc<str>,
67 last_used: u64,
69}
70
71struct Cache {
72 map: HashMap<PathBuf, Entry>,
73 total_bytes: usize,
74 budget_bytes: usize,
75 clock: u64,
76 hits: u64,
77 misses: u64,
78 inserts: u64,
79 evictions: u64,
80}
81
82impl Cache {
83 fn new(budget_bytes: usize) -> Self {
84 Self {
85 map: HashMap::new(),
86 total_bytes: 0,
87 budget_bytes,
88 clock: 0,
89 hits: 0,
90 misses: 0,
91 inserts: 0,
92 evictions: 0,
93 }
94 }
95
96 fn tick(&mut self) -> u64 {
97 self.clock += 1;
98 self.clock
99 }
100
101 fn remove_entry(&mut self, path: &Path) {
102 if let Some(old) = self.map.remove(path) {
103 self.total_bytes = self.total_bytes.saturating_sub(old.content.len());
104 }
105 }
106
107 fn evict_to_budget(&mut self) {
111 while self.total_bytes > self.budget_bytes && !self.map.is_empty() {
112 let Some(victim) = self
113 .map
114 .iter()
115 .min_by_key(|(_, e)| e.last_used)
116 .map(|(p, _)| p.clone())
117 else {
118 break;
119 };
120 self.remove_entry(&victim);
121 self.evictions += 1;
122 }
123 }
124}
125
126static CACHE: OnceLock<Mutex<Cache>> = OnceLock::new();
127
128fn budget_bytes() -> usize {
129 let mb = std::env::var("LEAN_CTX_CONTENT_CACHE_MB")
130 .ok()
131 .and_then(|v| v.trim().parse::<usize>().ok())
132 .unwrap_or(DEFAULT_BUDGET_MB);
133 mb.saturating_mul(1024 * 1024)
134}
135
136fn disabled() -> bool {
137 std::env::var("LEAN_CTX_DISABLE_CONTENT_CACHE")
140 .is_ok_and(|v| v == "1" || v.eq_ignore_ascii_case("true"))
141 || budget_bytes() == 0
142}
143
144fn cache() -> &'static Mutex<Cache> {
145 CACHE.get_or_init(|| Mutex::new(Cache::new(budget_bytes())))
146}
147
148fn lock() -> std::sync::MutexGuard<'static, Cache> {
149 cache()
150 .lock()
151 .unwrap_or_else(std::sync::PoisonError::into_inner)
152}
153
154pub fn get(path: &Path, current: FileState) -> Option<Arc<str>> {
159 if disabled() {
160 return None;
161 }
162 let mut c = lock();
163 let Some(entry) = c.map.get(path) else {
164 c.misses += 1;
165 return None;
166 };
167 let matches = entry.state == current;
168 if !matches {
169 c.remove_entry(path);
171 c.misses += 1;
172 return None;
173 }
174 let tick = c.tick();
175 c.hits += 1;
176 let entry = c.map.get_mut(path)?;
179 entry.last_used = tick;
180 Some(Arc::clone(&entry.content))
181}
182
183pub fn insert(path: &Path, state: FileState, content: Arc<str>) {
187 if disabled() || crate::core::memory_guard::is_under_pressure() {
188 return;
189 }
190 let len = content.len();
191 let mut c = lock();
192 if len > c.budget_bytes {
194 return;
195 }
196 c.remove_entry(path);
197 let tick = c.tick();
198 c.map.insert(
199 path.to_path_buf(),
200 Entry {
201 state,
202 content,
203 last_used: tick,
204 },
205 );
206 c.total_bytes += len;
207 c.inserts += 1;
208 if c.total_bytes > c.budget_bytes {
209 c.evict_to_budget();
210 }
211}
212
213pub fn get_or_read(path: &Path) -> Option<Arc<str>> {
219 let state = FileState::from_path(path)?;
220 if let Some(hit) = get(path, state) {
221 return Some(hit);
222 }
223 let content = std::fs::read_to_string(path).ok()?;
224 let arc: Arc<str> = Arc::from(content);
225 insert(path, state, Arc::clone(&arc));
226 Some(arc)
227}
228
229pub fn clear() {
232 if CACHE.get().is_none() {
233 return;
234 }
235 let mut c = lock();
236 c.map.clear();
237 c.total_bytes = 0;
238}
239
240pub fn trim_oldest_percent(percent: u8) {
244 if CACHE.get().is_none() {
245 return;
246 }
247 let mut c = lock();
248 if c.map.is_empty() {
249 return;
250 }
251 let pct = (percent.min(100)) as usize;
252 let target_evictions = c.map.len() * pct / 100;
253 for _ in 0..target_evictions {
254 let Some(victim) = c
255 .map
256 .iter()
257 .min_by_key(|(_, e)| e.last_used)
258 .map(|(p, _)| p.clone())
259 else {
260 break;
261 };
262 c.remove_entry(&victim);
263 c.evictions += 1;
264 }
265}
266
267pub fn memory_usage_bytes() -> usize {
269 if CACHE.get().is_none() {
270 return 0;
271 }
272 lock().total_bytes
273}
274
275#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
277pub struct CacheStats {
278 pub hits: u64,
279 pub misses: u64,
280 pub entries: usize,
281 pub bytes: usize,
282 pub inserts: u64,
283 pub evictions: u64,
284}
285
286pub fn stats() -> CacheStats {
287 if CACHE.get().is_none() {
288 return CacheStats::default();
289 }
290 let c = lock();
291 CacheStats {
292 hits: c.hits,
293 misses: c.misses,
294 entries: c.map.len(),
295 bytes: c.total_bytes,
296 inserts: c.inserts,
297 evictions: c.evictions,
298 }
299}
300
301#[cfg(test)]
302mod tests {
303 use super::*;
304
305 static TEST_LOCK: Mutex<()> = Mutex::new(());
308
309 fn fresh_cache(budget_bytes: usize) {
310 crate::test_env::remove_var("LEAN_CTX_CONTENT_CACHE_MB");
311 crate::test_env::remove_var("LEAN_CTX_DISABLE_CONTENT_CACHE");
312 let mut c = lock();
313 *c = Cache::new(budget_bytes);
314 }
315
316 fn write(dir: &Path, name: &str, body: &str) -> PathBuf {
317 let p = dir.join(name);
318 std::fs::write(&p, body).unwrap();
319 p
320 }
321
322 #[test]
323 fn hit_after_insert_with_matching_state() {
324 let _g = TEST_LOCK
325 .lock()
326 .unwrap_or_else(std::sync::PoisonError::into_inner);
327 fresh_cache(1024 * 1024);
328 let dir = tempfile::tempdir().unwrap();
329 let p = write(dir.path(), "a.rs", "fn main() {}\n");
330 let state = FileState::from_path(&p).unwrap();
331 assert!(get(&p, state).is_none(), "cold cache must miss");
332 insert(&p, state, Arc::from("fn main() {}\n"));
333 let got = get(&p, state).expect("warm cache must hit");
334 assert_eq!(&*got, "fn main() {}\n");
335 }
336
337 #[test]
338 fn mtime_or_size_change_invalidates() {
339 let _g = TEST_LOCK
340 .lock()
341 .unwrap_or_else(std::sync::PoisonError::into_inner);
342 fresh_cache(1024 * 1024);
343 let dir = tempfile::tempdir().unwrap();
344 let p = write(dir.path(), "a.rs", "v1\n");
345 let s1 = FileState::from_path(&p).unwrap();
346 insert(&p, s1, Arc::from("v1\n"));
347 assert!(get(&p, s1).is_some());
348
349 let s_bigger = FileState {
351 size_bytes: s1.size_bytes + 10,
352 ..s1
353 };
354 assert!(get(&p, s_bigger).is_none(), "size change must miss");
355 assert!(
356 get(&p, s1).is_none(),
357 "stale entry must be evicted on mismatch"
358 );
359
360 insert(&p, s1, Arc::from("v1\n"));
362 let s_newer = FileState {
363 mtime_ms: s1.mtime_ms + 1,
364 ..s1
365 };
366 assert!(get(&p, s_newer).is_none(), "mtime change must miss");
367 }
368
369 #[test]
370 fn get_or_read_populates_then_serves_from_cache() {
371 let _g = TEST_LOCK
372 .lock()
373 .unwrap_or_else(std::sync::PoisonError::into_inner);
374 fresh_cache(1024 * 1024);
375 let dir = tempfile::tempdir().unwrap();
376 let p = write(dir.path(), "a.rs", "hello world\n");
377
378 let before = stats();
379 let first = get_or_read(&p).unwrap();
380 assert_eq!(&*first, "hello world\n");
381 let after_first = stats();
382 assert_eq!(
383 after_first.inserts,
384 before.inserts + 1,
385 "first read inserts"
386 );
387
388 let second = get_or_read(&p).unwrap();
389 assert_eq!(&*second, "hello world\n");
390 let after_second = stats();
391 assert_eq!(
392 after_second.inserts, after_first.inserts,
393 "second read must NOT re-insert (served from cache)"
394 );
395 assert!(after_second.hits > after_first.hits, "second read is a hit");
396 }
397
398 #[test]
399 fn eviction_keeps_cache_within_budget() {
400 let _g = TEST_LOCK
401 .lock()
402 .unwrap_or_else(std::sync::PoisonError::into_inner);
403 fresh_cache(64);
405 let dir = tempfile::tempdir().unwrap();
406 let pa = write(dir.path(), "a", "aaaaaaaaaaaaaaaaaaaaaaaaaaaa"); let pb = write(dir.path(), "b", "bbbbbbbbbbbbbbbbbbbbbbbbbbbb");
408 let pc = write(dir.path(), "c", "cccccccccccccccccccccccccccc");
409 let sa = FileState::from_path(&pa).unwrap();
410 let sb = FileState::from_path(&pb).unwrap();
411 let sc = FileState::from_path(&pc).unwrap();
412
413 insert(&pa, sa, Arc::from("aaaaaaaaaaaaaaaaaaaaaaaaaaaa"));
414 let _ = get(&pa, sa);
416 insert(&pb, sb, Arc::from("bbbbbbbbbbbbbbbbbbbbbbbbbbbb"));
417 let _ = get(&pa, sa);
418 insert(&pc, sc, Arc::from("cccccccccccccccccccccccccccc"));
419
420 let st = stats();
421 assert!(st.bytes <= 64, "cache must respect byte budget: {st:?}");
422 assert!(st.evictions >= 1, "an eviction must have occurred: {st:?}");
423 assert!(get(&pa, sa).is_some(), "recently-used entry must survive");
424 }
425
426 #[test]
427 fn disabled_via_zero_budget_is_passthrough() {
428 let _g = TEST_LOCK
429 .lock()
430 .unwrap_or_else(std::sync::PoisonError::into_inner);
431 fresh_cache(1024 * 1024);
432 crate::test_env::set_var("LEAN_CTX_CONTENT_CACHE_MB", "0");
433 let dir = tempfile::tempdir().unwrap();
434 let p = write(dir.path(), "a.rs", "x\n");
435 let state = FileState::from_path(&p).unwrap();
436 insert(&p, state, Arc::from("x\n"));
437 assert!(get(&p, state).is_none(), "zero-budget cache is a no-op");
438 crate::test_env::remove_var("LEAN_CTX_CONTENT_CACHE_MB");
439 }
440}