| 1 | //! Decoded-image LRU with a byte budget. |
| 2 | |
| 3 | use std::collections::HashMap; |
| 4 | use std::path::{Path, PathBuf}; |
| 5 | use std::sync::Arc; |
| 6 | |
| 7 | use crate::decode::DecodedImage; |
| 8 | |
| 9 | struct Slot { |
| 10 | img: Arc<DecodedImage>, |
| 11 | used: u64, |
| 12 | } |
| 13 | |
| 14 | pub struct ImageCache { |
| 15 | map: HashMap<PathBuf, Slot>, |
| 16 | budget: usize, |
| 17 | bytes: usize, |
| 18 | clock: u64, |
| 19 | } |
| 20 | |
| 21 | impl ImageCache { |
| 22 | pub fn new(budget_bytes: usize) -> Self { |
| 23 | Self { |
| 24 | map: HashMap::new(), |
| 25 | budget: budget_bytes, |
| 26 | bytes: 0, |
| 27 | clock: 0, |
| 28 | } |
| 29 | } |
| 30 | |
| 31 | /// Takes effect at the next insert. |
| 32 | pub fn set_budget(&mut self, budget_bytes: usize) { |
| 33 | self.budget = budget_bytes; |
| 34 | } |
| 35 | |
| 36 | pub fn get(&mut self, path: &Path) -> Option<Arc<DecodedImage>> { |
| 37 | self.clock += 1; |
| 38 | let slot = self.map.get_mut(path)?; |
| 39 | slot.used = self.clock; |
| 40 | Some(slot.img.clone()) |
| 41 | } |
| 42 | |
| 43 | #[cfg(test)] |
| 44 | pub fn contains(&self, path: &Path) -> bool { |
| 45 | self.map.contains_key(path) |
| 46 | } |
| 47 | |
| 48 | /// Insert, then evict least-recently-used entries over budget, sparing `keep`. |
| 49 | pub fn insert(&mut self, path: PathBuf, img: Arc<DecodedImage>, keep: &[&Path]) { |
| 50 | self.clock += 1; |
| 51 | self.bytes += img.bytes(); |
| 52 | if let Some(old) = self.map.insert( |
| 53 | path, |
| 54 | Slot { |
| 55 | img, |
| 56 | used: self.clock, |
| 57 | }, |
| 58 | ) { |
| 59 | self.bytes -= old.img.bytes(); |
| 60 | } |
| 61 | while self.bytes > self.budget { |
| 62 | let victim = self |
| 63 | .map |
| 64 | .iter() |
| 65 | .filter(|(p, _)| !keep.contains(&p.as_path())) |
| 66 | .min_by_key(|(_, s)| s.used) |
| 67 | .map(|(p, _)| p.clone()); |
| 68 | let Some(victim) = victim else { break }; |
| 69 | self.remove(&victim); |
| 70 | } |
| 71 | } |
| 72 | |
| 73 | pub fn remove(&mut self, path: &Path) { |
| 74 | if let Some(old) = self.map.remove(path) { |
| 75 | self.bytes -= old.img.bytes(); |
| 76 | } |
| 77 | } |
| 78 | } |
| 79 | |
| 80 | #[cfg(test)] |
| 81 | mod tests { |
| 82 | use super::*; |
| 83 | |
| 84 | fn img(bytes: usize) -> Arc<DecodedImage> { |
| 85 | Arc::new(DecodedImage::still([1, 1], vec![0; bytes], "PNG")) |
| 86 | } |
| 87 | |
| 88 | #[test] |
| 89 | fn evicts_least_recent_but_spares_kept() { |
| 90 | let mut c = ImageCache::new(250); |
| 91 | c.insert("a".into(), img(100), &[]); |
| 92 | c.insert("b".into(), img(100), &[]); |
| 93 | c.get(Path::new("a")); |
| 94 | c.insert("c".into(), img(100), &[]); |
| 95 | assert!(c.contains(Path::new("a")) && !c.contains(Path::new("b"))); |
| 96 | c.insert("d".into(), img(100), &[Path::new("a")]); |
| 97 | assert!(c.contains(Path::new("a")) && !c.contains(Path::new("c"))); |
| 98 | } |
| 99 | } |