Skip to main content

rich_micro/
cache.rs

1//! Decoded, fitted images for drawing, cached (#584).
2//!
3//! An asset's image is read from its package, decoded under the package
4//! [`Limits`], fitted to exactly its cells at the terminal's cell size, and
5//! (for an animation) resampled: deduplicated and rate-limited under a
6//! memory budget. The result, a [`Prepared`], is kept in an in-memory cache
7//! bounded in bytes (the oldest entries go first). A still image fitted to
8//! its cells is also written to an on-disk cache of optimised variants, so
9//! the next process decodes a tiny PNG instead of the source. A variant is
10//! named by a SHA-256 of everything it depends on (the source bytes, the
11//! size in cells, the cell size and whether it is animated), so no package
12//! can name another asset's variant and replace what it draws.
13
14use std::collections::{HashMap, VecDeque};
15use std::path::{Path, PathBuf};
16use std::sync::{Arc, OnceLock};
17use std::time::Duration;
18
19use rich::Segment;
20use rich_art::graphics::{
21    AnimationFrame, CellPixels, FrameBudget, decode_animation, encode_png, fit_to_cells,
22};
23use rich_art::image::{self, DynamicImage};
24
25use sha2::{Digest, Sha256};
26
27use crate::model::{ImageFormat, ImageRef, MicroAsset};
28use crate::package::Limits;
29
30/// SHA-256 of `parts`, each length-prefixed so ("ab", "c") and ("a", "bc")
31/// differ, after a tag naming this key's layout. Stable across builds, so
32/// on-disk keys stay valid.
33pub(crate) fn digest(parts: &[&[u8]]) -> [u8; 32] {
34    let mut hasher = Sha256::new();
35    hasher.update(b"rs-rich-micro variant v2");
36    for part in parts {
37        hasher.update((part.len() as u64).to_le_bytes());
38        hasher.update(part);
39    }
40    hasher.finalize().into()
41}
42
43fn hex(bytes: &[u8]) -> String {
44    bytes.iter().map(|b| format!("{b:02x}")).collect()
45}
46
47/// An asset's image, ready to draw at one size.
48#[derive(Debug)]
49pub struct Prepared {
50    /// The first 64 bits of the SHA-256 of the source bytes, the size in
51    /// cells and the cell size: an id for the image within a session.
52    pub key: u64,
53    pub cols: usize,
54    pub rows: usize,
55    pub cell: CellPixels,
56    /// Fitted to exactly `cols × rows` cells; one frame for a still image.
57    pub frames: Vec<AnimationFrame>,
58    /// Per-frame protocol escapes, made on first use.
59    pub(crate) escapes: OnceLock<Vec<String>>,
60    /// Per-frame block cells, one line per row, made on first use.
61    pub(crate) blocks: OnceLock<Vec<Vec<Vec<Segment>>>>,
62}
63
64impl Prepared {
65    fn new(
66        key: u64,
67        cols: usize,
68        rows: usize,
69        cell: CellPixels,
70        frames: Vec<AnimationFrame>,
71    ) -> Self {
72        Prepared {
73            key,
74            cols,
75            rows,
76            cell,
77            frames,
78            escapes: OnceLock::new(),
79            blocks: OnceLock::new(),
80        }
81    }
82
83    /// Whether it has more than one frame.
84    pub fn animated(&self) -> bool {
85        self.frames.len() > 1
86    }
87
88    /// One full cycle.
89    pub fn duration(&self) -> Duration {
90        self.frames.iter().map(|frame| frame.delay).sum()
91    }
92
93    /// The frame that shows `elapsed` into playback (looping), and how long
94    /// until the next one.
95    pub fn frame_at(&self, elapsed: Duration) -> (usize, Option<Duration>) {
96        let total = self.duration();
97        if !self.animated() || total.is_zero() {
98            return (0, None);
99        }
100        let mut at = Duration::from_nanos((elapsed.as_nanos() % total.as_nanos()) as u64);
101        for (index, frame) in self.frames.iter().enumerate() {
102            if at < frame.delay {
103                return (index, Some(frame.delay - at));
104            }
105            at -= frame.delay;
106        }
107        (0, Some(self.frames[0].delay))
108    }
109
110    /// Bytes its frames take.
111    fn weight(&self) -> usize {
112        self.frames
113            .iter()
114            .map(|frame| frame.image.as_raw().len())
115            .sum()
116    }
117}
118
119/// Where an asset's images come from and what they are drawn as.
120#[derive(Clone, Debug, PartialEq, Eq, Hash)]
121struct Identity {
122    name: String,
123    origin: String,
124    image: String,
125    cols: usize,
126    rows: usize,
127    cell: CellPixels,
128    animate: bool,
129}
130
131/// The in-memory cache, bounded in bytes, plus an optional on-disk cache.
132#[derive(Debug)]
133pub struct ImageCache {
134    entries: HashMap<Identity, Arc<Prepared>>,
135    order: VecDeque<Identity>,
136    used: usize,
137    budget: usize,
138    disk: Option<PathBuf>,
139    /// Identities that failed to load, so they are not retried every frame.
140    failed: HashMap<Identity, ()>,
141}
142
143impl ImageCache {
144    /// Default bound on the decoded frames kept in memory: 32 MiB.
145    pub const DEFAULT_BUDGET: usize = 32 << 20;
146
147    pub fn new(budget: usize, disk: Option<PathBuf>) -> ImageCache {
148        ImageCache {
149            entries: HashMap::new(),
150            order: VecDeque::new(),
151            used: 0,
152            budget,
153            disk,
154            failed: HashMap::new(),
155        }
156    }
157
158    /// Bytes the cached frames take.
159    pub fn used(&self) -> usize {
160        self.used
161    }
162
163    /// Entries held.
164    pub fn len(&self) -> usize {
165        self.entries.len()
166    }
167
168    pub fn is_empty(&self) -> bool {
169        self.entries.is_empty()
170    }
171
172    /// `asset` ready to draw at `cell` pixels per cell, animated when
173    /// `animate` and it has an animation. `None` when it has no image, or its
174    /// image will not load or decode within `limits`.
175    pub fn prepare(
176        &mut self,
177        asset: &MicroAsset,
178        cell: CellPixels,
179        animate: bool,
180        limits: &Limits,
181    ) -> Option<Arc<Prepared>> {
182        let (image, animated) = pick(asset, animate)?;
183        let identity = Identity {
184            name: asset.name().to_string(),
185            origin: asset.origin().to_string(),
186            image: image.path.clone(),
187            cols: asset.cols(),
188            rows: asset.size().rows(),
189            cell,
190            animate: animated,
191        };
192        if let Some(found) = self.entries.get(&identity) {
193            return Some(Arc::clone(found));
194        }
195        if self.failed.contains_key(&identity) {
196            return None;
197        }
198        let prepared = self.load(asset, image, animated, &identity, limits);
199        match prepared {
200            Some(prepared) => {
201                let prepared = Arc::new(prepared);
202                self.insert(identity, Arc::clone(&prepared));
203                Some(prepared)
204            }
205            None => {
206                self.failed.insert(identity, ());
207                None
208            }
209        }
210    }
211
212    fn insert(&mut self, identity: Identity, prepared: Arc<Prepared>) {
213        self.used += prepared.weight();
214        self.order.push_back(identity.clone());
215        self.entries.insert(identity, prepared);
216        while self.used > self.budget && self.order.len() > 1 {
217            let Some(oldest) = self.order.pop_front() else {
218                break;
219            };
220            if let Some(gone) = self.entries.remove(&oldest) {
221                self.used -= gone.weight();
222            }
223        }
224    }
225
226    fn load(
227        &self,
228        asset: &MicroAsset,
229        image: &ImageRef,
230        animated: bool,
231        identity: &Identity,
232        limits: &Limits,
233    ) -> Option<Prepared> {
234        let bytes = asset.read_variant(image, limits).ok()?;
235        let (cols, rows) = (identity.cols, identity.rows);
236        let digest = digest(&[
237            &bytes,
238            &(cols as u64).to_le_bytes(),
239            &(rows as u64).to_le_bytes(),
240            &identity.cell.width.to_le_bytes(),
241            &identity.cell.height.to_le_bytes(),
242            &[u8::from(animated)],
243        ]);
244        let key = u64::from_le_bytes(digest[..8].try_into().expect("8 bytes"));
245        if !animated && let Some(fitted) = self.disk_get(&digest, cols, rows, identity.cell) {
246            return Some(Prepared::new(
247                key,
248                cols,
249                rows,
250                identity.cell,
251                vec![still(fitted)],
252            ));
253        }
254        let budget = FrameBudget {
255            max_bytes: usize::try_from(limits.max_decoded_bytes).unwrap_or(usize::MAX),
256            max_frames: limits.max_frames as usize,
257            ..FrameBudget::default()
258        };
259        let decoded = decode_animation(&bytes, budget).ok()?;
260        let frames = if animated {
261            let (width, height) = identity.cell.span(cols, rows);
262            rich_art::graphics::resample(&decoded, width, height, budget)
263        } else {
264            let first = decoded.into_iter().next()?;
265            let fitted = fit_to_cells(
266                &DynamicImage::ImageRgba8(first.image),
267                cols,
268                rows,
269                identity.cell,
270            );
271            self.disk_put(&digest, &fitted);
272            vec![still(fitted)]
273        };
274        if frames.is_empty() {
275            return None;
276        }
277        Some(Prepared::new(key, cols, rows, identity.cell, frames))
278    }
279
280    fn disk_path(&self, digest: &[u8; 32]) -> Option<PathBuf> {
281        Some(self.disk.as_ref()?.join(format!("{}.png", hex(digest))))
282    }
283
284    fn disk_get(
285        &self,
286        digest: &[u8; 32],
287        cols: usize,
288        rows: usize,
289        cell: CellPixels,
290    ) -> Option<image::RgbaImage> {
291        let path = self.disk_path(digest)?;
292        let bytes = std::fs::read(&path).ok()?;
293        let image = image::load_from_memory_with_format(&bytes, image::ImageFormat::Png)
294            .ok()?
295            .to_rgba8();
296        // A file that is not what the key says is ignored.
297        (image.dimensions() == cell.span(cols, rows)).then_some(image)
298    }
299
300    fn disk_put(&self, digest: &[u8; 32], image: &image::RgbaImage) {
301        let Some(path) = self.disk_path(digest) else {
302            return;
303        };
304        if path.exists() {
305            return;
306        }
307        let Some(png) = encode_png(image) else {
308            return;
309        };
310        let _ = write_atomically(&path, &png);
311    }
312}
313
314fn still(image: image::RgbaImage) -> AnimationFrame {
315    AnimationFrame {
316        image,
317        delay: Duration::from_millis(100),
318    }
319}
320
321/// Write `bytes` to `path` through a temporary file in its directory.
322fn write_atomically(path: &Path, bytes: &[u8]) -> std::io::Result<()> {
323    let dir = path.parent().ok_or(std::io::ErrorKind::InvalidInput)?;
324    std::fs::create_dir_all(dir)?;
325    let temporary = dir.join(format!(
326        ".{}.{}.tmp",
327        path.file_name().and_then(|n| n.to_str()).unwrap_or("micro"),
328        std::process::id()
329    ));
330    std::fs::write(&temporary, bytes)?;
331    std::fs::rename(&temporary, path).inspect_err(|_| {
332        let _ = std::fs::remove_file(&temporary);
333    })
334}
335
336/// The image to draw `asset` from: its animation when `animate` and it has
337/// one this build decodes (GIF, APNG, WebP), else its still image, else the first
338/// image of its animation or frame sequence. And whether it is animated.
339fn pick(asset: &MicroAsset, animate: bool) -> Option<(&ImageRef, bool)> {
340    let variants = asset.variants();
341    let decodable = |image: &&ImageRef| {
342        matches!(
343            image.info.format,
344            ImageFormat::Png | ImageFormat::Apng | ImageFormat::Gif | ImageFormat::WebP
345        )
346    };
347    if animate
348        && let Some(animation) = variants.animation.as_ref().filter(decodable)
349        && animation.info.frames > 1
350    {
351        return Some((animation, true));
352    }
353    variants
354        .static_image
355        .as_ref()
356        .filter(decodable)
357        .or(variants.animation.as_ref().filter(decodable))
358        .or(variants.frames.iter().find(decodable))
359        .map(|image| (image, false))
360}
361
362/// The user cache directory for micro assets: `RICH_CACHE_DIR/micro`, else
363/// `$XDG_CACHE_HOME/rich/micro`, `~/Library/Caches/rich/micro` on macOS,
364/// `%LOCALAPPDATA%\rich\cache\micro` on Windows, `~/.cache/rich/micro`
365/// elsewhere.
366pub fn user_cache_dir() -> Option<PathBuf> {
367    let var = |name: &str| std::env::var_os(name).filter(|v| !v.is_empty());
368    if let Some(dir) = var("RICH_CACHE_DIR") {
369        return Some(PathBuf::from(dir).join("micro"));
370    }
371    if cfg!(windows) {
372        return var("LOCALAPPDATA")
373            .map(|d| PathBuf::from(d).join("rich").join("cache").join("micro"));
374    }
375    if let Some(dir) = var("XDG_CACHE_HOME") {
376        return Some(PathBuf::from(dir).join("rich").join("micro"));
377    }
378    let home = PathBuf::from(var("HOME")?);
379    if cfg!(target_os = "macos") {
380        Some(
381            home.join("Library")
382                .join("Caches")
383                .join("rich")
384                .join("micro"),
385        )
386    } else {
387        Some(home.join(".cache").join("rich").join("micro"))
388    }
389}
390
391#[cfg(test)]
392mod tests {
393    use super::*;
394
395    #[test]
396    fn digests_separate_parts() {
397        assert_ne!(digest(&[b"ab", b"c"]), digest(&[b"a", b"bc"]));
398        assert_ne!(digest(&[b"ab", b""]), digest(&[b"ab"]));
399        assert_eq!(digest(&[b"x"]), digest(&[b"x"]));
400        assert_eq!(hex(&[0, 0xab]), "00ab");
401    }
402
403    #[test]
404    fn frames_are_time_indexed_and_loop() {
405        let frame = |ms| AnimationFrame {
406            image: image::RgbaImage::new(1, 1),
407            delay: Duration::from_millis(ms),
408        };
409        let prepared = Prepared::new(1, 2, 1, CellPixels::DEFAULT, vec![frame(100), frame(50)]);
410        assert_eq!(
411            prepared.frame_at(Duration::ZERO),
412            (0, Some(Duration::from_millis(100)))
413        );
414        assert_eq!(
415            prepared.frame_at(Duration::from_millis(120)),
416            (1, Some(Duration::from_millis(30)))
417        );
418        assert_eq!(prepared.frame_at(Duration::from_millis(150)).0, 0);
419        let still = Prepared::new(1, 2, 1, CellPixels::DEFAULT, vec![frame(100)]);
420        assert_eq!(still.frame_at(Duration::from_secs(9)), (0, None));
421    }
422}