1use 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
30pub(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#[derive(Debug)]
49pub struct Prepared {
50 pub key: u64,
53 pub cols: usize,
54 pub rows: usize,
55 pub cell: CellPixels,
56 pub frames: Vec<AnimationFrame>,
58 pub(crate) escapes: OnceLock<Vec<String>>,
60 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 pub fn animated(&self) -> bool {
85 self.frames.len() > 1
86 }
87
88 pub fn duration(&self) -> Duration {
90 self.frames.iter().map(|frame| frame.delay).sum()
91 }
92
93 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 fn weight(&self) -> usize {
112 self.frames
113 .iter()
114 .map(|frame| frame.image.as_raw().len())
115 .sum()
116 }
117}
118
119#[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#[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 failed: HashMap<Identity, ()>,
141}
142
143impl ImageCache {
144 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 pub fn used(&self) -> usize {
160 self.used
161 }
162
163 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 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 (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
321fn 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
336fn 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
362pub 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}