use crate::cache::CachedContent;
use std::collections::HashMap;
use std::path::{Path, PathBuf};
use std::sync::{Arc, Mutex};
struct LruEntry {
content: CachedContent,
last_used: u64,
}
#[derive(Clone)]
pub struct LruCache {
inner: Arc<Mutex<LruInner>>,
}
struct LruInner {
map: HashMap<PathBuf, LruEntry>,
total_bytes: u64,
byte_budget: u64,
entry_budget: usize,
epoch: u64,
evicted_count: u64,
hook: Option<Box<dyn Fn(PathBuf) + Send + Sync>>,
}
impl LruCache {
pub fn new(byte_budget: u64, entry_budget: usize) -> Self {
Self {
inner: Arc::new(Mutex::new(LruInner {
map: HashMap::new(),
total_bytes: 0,
byte_budget,
entry_budget,
epoch: 0,
evicted_count: 0,
hook: None,
})),
}
}
pub fn on_evict<F: Fn(PathBuf) + Send + Sync + 'static>(&self, f: F) {
self.inner.lock().unwrap().hook = Some(Box::new(f));
}
pub fn size_bytes(&self) -> u64 {
self.inner.lock().unwrap().total_bytes
}
pub fn len(&self) -> usize {
self.inner.lock().unwrap().map.len()
}
pub fn is_empty(&self) -> bool {
self.len() == 0
}
pub fn evicted_count(&self) -> u64 {
self.inner.lock().unwrap().evicted_count
}
pub fn insert(&self, path: PathBuf, content: CachedContent) {
let size = content_size_of(&content);
let mut inner = self.inner.lock().unwrap();
inner.epoch += 1;
let tick = inner.epoch;
if let Some(old) = inner.map.insert(
path,
LruEntry {
content,
last_used: tick,
},
) {
inner.total_bytes = inner
.total_bytes
.saturating_sub(content_size_of(&old.content));
}
inner.total_bytes = inner.total_bytes.saturating_add(size);
inner.evict_if_needed();
}
pub fn insert_cold(&self, path: PathBuf, content: CachedContent) {
let size = content_size_of(&content);
let mut inner = self.inner.lock().unwrap();
if let Some(old) = inner.map.insert(
path,
LruEntry {
content,
last_used: 0,
},
) {
inner.total_bytes = inner
.total_bytes
.saturating_sub(content_size_of(&old.content));
}
inner.total_bytes = inner.total_bytes.saturating_add(size);
inner.evict_if_needed();
}
pub fn get(&self, path: &Path) -> Option<CachedContent> {
let mut inner = self.inner.lock().unwrap();
inner.epoch += 1;
let tick = inner.epoch;
let entry = inner.map.get_mut(path)?;
entry.last_used = tick;
Some(entry.content.clone())
}
pub fn remove(&self, path: &Path) -> Option<CachedContent> {
let mut inner = self.inner.lock().unwrap();
let entry = inner.map.remove(path)?;
inner.total_bytes = inner
.total_bytes
.saturating_sub(content_size_of(&entry.content));
Some(entry.content)
}
pub fn replace_all(&self, map: HashMap<PathBuf, CachedContent>) {
let mut inner = self.inner.lock().unwrap();
inner.map.clear();
inner.total_bytes = 0;
for (path, content) in map {
let size = content_size_of(&content);
inner.map.insert(
path,
LruEntry {
content,
last_used: 0,
},
);
inner.total_bytes = inner.total_bytes.saturating_add(size);
}
inner.evict_if_needed();
}
pub fn keys(&self) -> Vec<PathBuf> {
self.inner.lock().unwrap().map.keys().cloned().collect()
}
pub fn iter_all(&self) -> Vec<(PathBuf, CachedContent)> {
self.inner
.lock()
.unwrap()
.map
.iter()
.map(|(p, e)| (p.clone(), e.content.clone()))
.collect()
}
}
impl LruInner {
fn evict_if_needed(&mut self) {
while self.byte_budget > 0 && self.total_bytes > self.byte_budget && !self.map.is_empty() {
self.evict_lru_once();
}
while self.entry_budget > 0 && self.map.len() > self.entry_budget {
self.evict_lru_once();
}
}
fn evict_lru_once(&mut self) {
let victim = self
.map
.iter()
.min_by_key(|(_, e)| e.last_used)
.map(|(p, _)| p.clone());
let Some(path) = victim else { return };
if let Some(entry) = self.map.remove(&path) {
self.total_bytes = self
.total_bytes
.saturating_sub(content_size_of(&entry.content));
if let Some(hook) = &self.hook {
hook(path);
}
self.evicted_count += 1;
}
}
}
pub fn content_size_of(content: &CachedContent) -> u64 {
match content {
CachedContent::Text(s) => s.len() as u64,
CachedContent::Binary(b) => b.len() as u64,
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::sync::{Arc, Mutex};
fn text(s: &str) -> CachedContent {
CachedContent::Text(s.to_string())
}
#[test]
fn evicts_least_recently_used_when_over_byte_budget() {
let cache = LruCache::new(10, 0);
let seen = Arc::new(Mutex::new(Vec::new()));
let seen_hook = Arc::clone(&seen);
cache.on_evict(move |p| seen_hook.lock().unwrap().push(p));
let a = PathBuf::from("/a.md"); let b = PathBuf::from("/b.md"); cache.insert(a.clone(), text("aaaa"));
cache.insert(b.clone(), text("bbbbbb"));
assert_eq!(cache.size_bytes(), 10);
cache.get(&a);
cache.insert(PathBuf::from("/c.md"), text("cccccccc")); assert_eq!(cache.size_bytes(), 8);
assert!(cache.get(&b).is_none(), "b was the LRU victim");
assert!(cache.get(&a).is_none(), "a was evicted second");
assert!(cache.get(Path::new("/c.md")).is_some());
let evicted = seen.lock().unwrap();
assert_eq!(evicted.len(), 2);
assert_eq!(evicted[0], PathBuf::from("/b.md"));
assert_eq!(evicted[1], PathBuf::from("/a.md"));
}
#[test]
fn evicts_by_entry_budget() {
let cache = LruCache::new(0, 2);
cache.insert(PathBuf::from("/1.md"), text("one"));
cache.insert(PathBuf::from("/2.md"), text("two"));
cache.insert(PathBuf::from("/3.md"), text("three"));
assert_eq!(cache.len(), 2);
assert!(cache.get(Path::new("/1.md")).is_none());
assert!(cache.get(Path::new("/2.md")).is_some());
assert_eq!(cache.evicted_count(), 1);
}
#[test]
fn replace_all_resets_state() {
let cache = LruCache::new(0, 3);
cache.insert(PathBuf::from("/x.md"), text("xxxx"));
let mut fresh = HashMap::new();
fresh.insert(PathBuf::from("/y.md"), text("yy"));
cache.replace_all(fresh);
assert_eq!(cache.len(), 1);
assert!(cache.get(Path::new("/x.md")).is_none());
assert_eq!(cache.size_bytes(), 2);
}
#[test]
fn unlimited_budgets_keep_everything() {
let cache = LruCache::new(0, 0);
for i in 0..50 {
cache.insert(PathBuf::from(format!("/f{i}.md")), text("x"));
}
assert_eq!(cache.len(), 50);
assert_eq!(cache.evicted_count(), 0);
}
#[test]
fn insert_cold_is_evicted_before_warm_entries() {
let cache = LruCache::new(0, 2);
cache.insert_cold(PathBuf::from("/cold.md"), text("c"));
cache.insert(PathBuf::from("/warm.md"), text("w"));
cache.insert(PathBuf::from("/new.md"), text("n"));
assert!(cache.get(Path::new("/cold.md")).is_none());
assert!(cache.get(Path::new("/warm.md")).is_some());
}
#[test]
fn remove_does_not_fire_eviction_hook() {
let cache = LruCache::new(10, 0);
let evictions = Arc::new(Mutex::new(0u32));
let evictions_hook = Arc::clone(&evictions);
cache.on_evict(move |_| {
*evictions_hook.lock().unwrap() += 1;
});
cache.insert(PathBuf::from("/a.md"), text("aaaa"));
assert!(cache.remove(Path::new("/a.md")).is_some());
assert_eq!(*evictions.lock().unwrap(), 0);
assert_eq!(cache.evicted_count(), 0);
}
#[test]
fn insert_replaces_previous_size_accounting() {
let cache = LruCache::new(100, 0);
cache.insert(PathBuf::from("/a.md"), text("short"));
cache.insert(PathBuf::from("/a.md"), text("a much longer value"));
assert_eq!(cache.size_bytes(), "a much longer value".len() as u64);
assert_eq!(cache.len(), 1);
}
}