treehouse/src/vfs/image_size_cache.rs

use std::{fmt, io::Cursor};

use anyhow::Context;
use dashmap::DashMap;
use tracing::{info_span, instrument, warn};

use crate::config;

use super::{Content, Dir, ImageSize, Query, VPath, VPathBuf, query};

pub struct ImageSizeCache<T> {
    inner: T,
    cache: DashMap<VPathBuf, ImageSize>,
}

impl<T> ImageSizeCache<T> {
    pub fn new(inner: T) -> Self {
        Self {
            inner,
            cache: DashMap::new(),
        }
    }
}

impl<T> ImageSizeCache<T>
where
    T: Dir,
{
    fn compute_image_size(&self, path: &VPath) -> anyhow::Result<Option<ImageSize>> {
        if path.extension().is_some_and(config::is_image_file)
            && let Some(content) = query::<Content>(&self.inner, path)
        {
            if path.extension() == Some("svg") {
                return Ok(svg_size(&content.string()?));
            } else {
                let _span = info_span!("raster_image_size").entered();
                let reader = image::ImageReader::new(Cursor::new(content.bytes()))
                    .with_guessed_format()
                    .context("cannot guess image format")?;
                let (width, height) = reader.into_dimensions()?;
                return Ok(Some(ImageSize { width, height }));
            }
        }

        Ok(None)
    }

    #[instrument("ImageSizeCache::image_size", skip(self))]
    fn image_size(&self, path: &VPath) -> Option<ImageSize> {
        self.cache.get(path).map(|x| *x).or_else(|| {
            let _span = info_span!("cache_miss").entered();

            let image_size = self
                .compute_image_size(path)
                .inspect_err(|err| warn!(%path, ?err, "compute_image_size failure"))
                .ok()
                .flatten();
            if let Some(image_size) = image_size {
                self.cache.insert(path.to_owned(), image_size);
            }
            image_size
        })
    }
}

impl<T> Dir for ImageSizeCache<T>
where
    T: Dir,
{
    fn query(&self, path: &VPath, query: &mut Query) {
        query.try_provide(|| self.image_size(path));

        self.inner.query(path, query);
    }
}

impl<T> fmt::Debug for ImageSizeCache<T>
where
    T: fmt::Debug,
{
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "ImageSizeCache({:?})", self.inner)
    }
}

/// Quickly determine the size of an SVG without parsing it into a DOM.
///
/// This method is a tentative check; the point is to return an image size that's _good enough
/// default_ rather than what the size is going to be in the user's web browser.
#[instrument(skip(svg))]
fn svg_size(svg: &str) -> Option<ImageSize> {
    let mut tokenizer = xmlparser::Tokenizer::from(svg);

    fn parse_view_box(s: &str) -> Option<[u32; 4]> {
        let mut iter = s.split_whitespace();
        let min_x = iter.next()?.parse().ok()?;
        let min_y = iter.next()?.parse().ok()?;
        let width = iter.next()?.parse().ok()?;
        let height = iter.next()?.parse().ok()?;
        Some([min_x, min_y, width, height])
    }

    let mut in_svg = false;
    let mut width: Option<u32> = None;
    let mut height: Option<u32> = None;
    let mut view_box: Option<[u32; 4]> = None;
    while let Some(Ok(token)) = tokenizer.next() {
        if let xmlparser::Token::ElementStart { local, .. } = &token
            && local == "svg"
        {
            in_svg = true;
            continue;
        }

        if in_svg {
            // If another element starts, we're no longer in the root <svg>.
            if let xmlparser::Token::ElementStart { .. } = &token {
                break;
            }

            if let xmlparser::Token::Attribute { local, value, .. } = &token {
                match local.as_str() {
                    "width" => width = value.parse().ok(),
                    "height" => height = value.parse().ok(),
                    "viewBox" => {
                        view_box = parse_view_box(value);
                    }
                    _ => (),
                }
                continue;
            }
        }
    }

    match (width, height, view_box) {
        (Some(width), Some(height), _) | (_, _, Some([_, _, width, height])) => {
            Some(ImageSize { width, height })
        }
        _ => None,
    }
}