1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
use super::{
    generate_loader_id, Bytes, BytesLoadResult, BytesLoader, BytesPoll, Context, Cow, HashMap,
    LoadError, Mutex,
};

/// Maps URI:s to [`Bytes`], e.g. found with `include_bytes!`.
///
/// By convention, the URI:s should be prefixed with `bytes://`.
#[derive(Default)]
pub struct DefaultBytesLoader {
    cache: Mutex<HashMap<Cow<'static, str>, Bytes>>,
}

impl DefaultBytesLoader {
    pub fn insert(&self, uri: impl Into<Cow<'static, str>>, bytes: impl Into<Bytes>) {
        self.cache
            .lock()
            .entry(uri.into())
            .or_insert_with_key(|_uri| {
                let bytes: Bytes = bytes.into();

                #[cfg(feature = "log")]
                log::trace!("loaded {} bytes for uri {_uri:?}", bytes.len());

                bytes
            });
    }
}

impl BytesLoader for DefaultBytesLoader {
    fn id(&self) -> &str {
        generate_loader_id!(DefaultBytesLoader)
    }

    fn load(&self, _: &Context, uri: &str) -> BytesLoadResult {
        // We accept uri:s that don't start with `bytes://` too… for now.
        match self.cache.lock().get(uri).cloned() {
            Some(bytes) => Ok(BytesPoll::Ready {
                size: None,
                bytes,
                mime: None,
            }),
            None => {
                if uri.starts_with("bytes://") {
                    Err(LoadError::Loading(
                        "Bytes not found. Did you forget to call Context::include_bytes?".into(),
                    ))
                } else {
                    Err(LoadError::NotSupported)
                }
            }
        }
    }

    fn forget(&self, uri: &str) {
        #[cfg(feature = "log")]
        log::trace!("forget {uri:?}");

        self.cache.lock().remove(uri);
    }

    fn forget_all(&self) {
        #[cfg(feature = "log")]
        log::trace!("forget all");

        self.cache.lock().clear();
    }

    fn byte_size(&self) -> usize {
        self.cache.lock().values().map(|bytes| bytes.len()).sum()
    }
}