pdf_interpret/

cache.rs

use crate::util::hash128;
use kurbo::Affine;
use pdf_syntax::object::{Array, Dict, MaybeRef, Name, Null, ObjRef, Object, Stream};
use std::any::Any;
use std::collections::hash_map::Entry;
use std::collections::{HashMap, VecDeque};
use std::sync::{Arc, Mutex};

type CacheMap = HashMap<u128, Option<Box<dyn Any + Send + Sync>>>;

/// Maximum number of decoded images retained by the document-level image cache.
///
/// The cache is bounded by BOTH entry count (this constant) and total decoded
/// bytes ([`IMAGE_CACHE_MAX_BYTES`]); an insertion evicts least-recently-used
/// entries until both bounds hold. The count cap catches "many small images"
/// (logos, headers/footers repeated across pages); the byte cap catches "few
/// but huge images" that would otherwise pass the count cap while consuming
/// hundreds of MB. The cache is per-`PdfDocument` and lives only as long as
/// that document.
const IMAGE_CACHE_CAPACITY: usize = 32;

/// Maximum total decoded-image bytes retained by the document-level image cache.
///
/// 256 MiB is a conservative working-set ceiling for image-heavy documents: it
/// comfortably holds typical repeated assets while bounding peak transient
/// memory for pathological inputs (e.g. a handful of full-page scans). An entry
/// larger than the whole budget is returned to the caller but not retained.
const IMAGE_CACHE_MAX_BYTES: usize = 256 * 1024 * 1024;

struct ImageCache {
    /// Maps a cache key to the decoded value and its accounted byte weight.
    map: HashMap<u128, (Box<dyn Any + Send + Sync>, usize)>,
    /// Least-recently-used order; front = oldest, back = most recently used.
    order: VecDeque<u128>,
    /// Sum of the byte weights of all entries currently in `map`.
    total_bytes: usize,
}

/// A cache to store decoded images and other objects to avoid parsing/decoding them repeatedly.
#[derive(Clone)]
pub struct Cache {
    generic: Arc<Mutex<CacheMap>>,
    images: Arc<Mutex<ImageCache>>,
}

impl Default for Cache {
    fn default() -> Self {
        Self::new()
    }
}

impl Cache {
    /// Create a new, empty cache.
    pub fn new() -> Self {
        Self {
            generic: Arc::new(Mutex::new(HashMap::new())),
            images: Arc::new(Mutex::new(ImageCache {
                map: HashMap::new(),
                order: VecDeque::new(),
                total_bytes: 0,
            })),
        }
    }

    pub(crate) fn get_or_insert_with<T: Clone + Send + Sync + 'static>(
        &self,
        id: u128,
        f: impl FnOnce() -> Option<T>,
    ) -> Option<T> {
        let mut locked = self.generic.lock().unwrap();

        // We can't use `get_or_insert_with` here, because if the closure makes another access to the
        // cache, we end up with a deadlock.
        match locked.entry(id) {
            Entry::Occupied(o) => o
                .get()
                .as_ref()
                .and_then(|val| val.downcast_ref::<T>().cloned()),
            Entry::Vacant(_) => {
                drop(locked);
                let val = f();
                self.generic.lock().unwrap().insert(
                    id,
                    val.clone()
                        .map(|val| Box::new(val) as Box<dyn Any + Send + Sync>),
                );

                val
            }
        }
    }

    /// Look up a parsed font program by `id`, or build it with `f` and keep it
    /// for the rest of the document.
    ///
    /// Unlike [`Self::get_or_insert_with`], a `None` is not remembered: a
    /// program that fails to parse is tried again on the next page, so every
    /// warning the failure raises reaches the sink of the page that hit it,
    /// exactly as without the cache.
    pub(crate) fn get_or_insert_some<T: Clone + Send + Sync + 'static>(
        &self,
        id: u128,
        f: impl FnOnce() -> Option<T>,
    ) -> Option<T> {
        if let Some(Some(val)) = self.generic.lock().unwrap().get(&id)
            && let Some(val) = val.downcast_ref::<T>()
        {
            return Some(val.clone());
        }
        let val = f()?;
        self.generic.lock().unwrap().insert(
            id,
            Some(Box::new(val.clone()) as Box<dyn Any + Send + Sync>),
        );
        Some(val)
    }

    /// The key a font program stream is kept under; `tag` names the parser
    /// that reads it, since one stream can be tried by several.
    ///
    /// THE BYTES, NOT THE ADDRESS. `InterpreterSettings::shared_cache` is
    /// public and a `Cache` is `Clone`, so one cache can serve two
    /// documents, and object 5 with the same dictionary is a different font
    /// in each. The key is therefore a hash of the stream's dictionary and
    /// its raw bytes, which no other font program shares. Hashing a font of a
    /// few hundred kilobytes costs tens of microseconds per page it is set
    /// on, against the milliseconds the parse it replaces took.
    pub(crate) fn font_stream_key(tag: &[u8], stream: &Stream<'_>) -> u128 {
        hash128(&(tag, stream.cache_key(), stream.raw_data().as_ref()))
    }

    /// The key a font program handed over as bytes (by the font resolver) is
    /// kept under: the bytes themselves, so two resolvers that answer the
    /// same query differently never share an entry.
    pub(crate) fn font_data_key(tag: &[u8], index: u32, data: &[u8]) -> u128 {
        hash128(&(tag, index, data))
    }

    /// Look up a decoded image by `id`, or compute it with `f` and cache it.
    ///
    /// `weight` returns the byte size of a freshly-computed value; it is used to
    /// keep the cache within [`IMAGE_CACHE_MAX_BYTES`]. The cache is also bounded
    /// to [`IMAGE_CACHE_CAPACITY`] entries. After an insertion, least-recently-
    /// used entries are evicted until both bounds hold (a single value larger
    /// than the whole budget is returned but not retained).
    pub(crate) fn get_or_insert_image<T: Clone + Send + Sync + 'static>(
        &self,
        id: u128,
        f: impl FnOnce() -> Option<T>,
        weight: impl FnOnce(&T) -> usize,
    ) -> Option<T> {
        let mut locked = self.images.lock().unwrap();
        if locked.map.contains_key(&id) {
            if let Some(pos) = locked.order.iter().position(|&x| x == id) {
                locked.order.remove(pos);
            }
            locked.order.push_back(id);
            let (val, _) = locked.map.get(&id).unwrap();
            return val.downcast_ref::<T>().cloned();
        }

        drop(locked);
        let val = f();
        if let Some(ref v) = val {
            let w = weight(v);
            let mut locked = self.images.lock().unwrap();
            // A repeated key would otherwise leak its old weight; subtract it.
            if let Some((_, old_w)) = locked
                .map
                .insert(id, (Box::new(v.clone()) as Box<dyn Any + Send + Sync>, w))
            {
                locked.total_bytes = locked.total_bytes.saturating_sub(old_w);
                if let Some(pos) = locked.order.iter().position(|&x| x == id) {
                    locked.order.remove(pos);
                }
            }
            locked.total_bytes += w;
            locked.order.push_back(id);

            // Evict least-recently-used entries until both the count and the
            // byte budget are satisfied.
            while !locked.order.is_empty()
                && (locked.order.len() > IMAGE_CACHE_CAPACITY
                    || locked.total_bytes > IMAGE_CACHE_MAX_BYTES)
            {
                if let Some(oldest_id) = locked.order.pop_front()
                    && let Some((_, evicted_w)) = locked.map.remove(&oldest_id)
                {
                    locked.total_bytes = locked.total_bytes.saturating_sub(evicted_w);
                }
            }
        }
        val
    }

    /// Current image-cache occupancy as `(entry_count, total_bytes)`.
    ///
    /// Exposed for measurement and tests; both values are always within
    /// [`IMAGE_CACHE_CAPACITY`] and [`IMAGE_CACHE_MAX_BYTES`] respectively after
    /// any completed `get_or_insert_image` call.
    #[cfg(test)]
    pub(crate) fn image_cache_stats(&self) -> (usize, usize) {
        let locked = self.images.lock().unwrap();
        (locked.map.len(), locked.total_bytes)
    }
}

/// A trait for objects that can generate a unique cache key.
pub trait CacheKey {
    /// Returns the cache key for this object.
    fn cache_key(&self) -> u128;
}

impl<T: CacheKey, U: CacheKey> CacheKey for (T, U) {
    fn cache_key(&self) -> u128 {
        hash128(&(self.0.cache_key(), self.1.cache_key()))
    }
}

impl CacheKey for Dict<'_> {
    fn cache_key(&self) -> u128 {
        hash128(self.data())
    }
}

impl CacheKey for Stream<'_> {
    fn cache_key(&self) -> u128 {
        self.dict().cache_key()
    }
}

impl CacheKey for Null {
    fn cache_key(&self) -> u128 {
        hash128(self)
    }
}

impl CacheKey for bool {
    fn cache_key(&self) -> u128 {
        hash128(self)
    }
}

impl CacheKey for pdf_syntax::object::Number {
    fn cache_key(&self) -> u128 {
        hash128(&self.as_f64().to_bits())
    }
}

impl CacheKey for pdf_syntax::object::String {
    fn cache_key(&self) -> u128 {
        hash128(self.as_ref())
    }
}

impl CacheKey for Name {
    fn cache_key(&self) -> u128 {
        hash128(self)
    }
}

impl CacheKey for Array<'_> {
    fn cache_key(&self) -> u128 {
        hash128(self.data())
    }
}

impl CacheKey for Object<'_> {
    fn cache_key(&self) -> u128 {
        match self {
            Object::Null(n) => n.cache_key(),
            Object::Boolean(b) => b.cache_key(),
            Object::Number(n) => n.cache_key(),
            Object::String(s) => s.cache_key(),
            Object::Name(n) => n.cache_key(),
            Object::Dict(d) => d.cache_key(),
            Object::Array(a) => a.cache_key(),
            Object::Stream(s) => s.cache_key(),
        }
    }
}

impl CacheKey for ObjRef {
    fn cache_key(&self) -> u128 {
        hash128(self)
    }
}

impl<T: CacheKey> CacheKey for MaybeRef<T> {
    fn cache_key(&self) -> u128 {
        match self {
            Self::Ref(r) => r.cache_key(),
            Self::NotRef(o) => o.cache_key(),
        }
    }
}

impl CacheKey for Affine {
    fn cache_key(&self) -> u128 {
        let c = self.as_coeffs();
        hash128(&[
            c[0].to_bits(),
            c[1].to_bits(),
            c[2].to_bits(),
            c[3].to_bits(),
            c[4].to_bits(),
            c[5].to_bits(),
        ])
    }
}

impl CacheKey for u128 {
    fn cache_key(&self) -> u128 {
        hash128(self)
    }
}

#[cfg(test)]
mod tests {
    /// Two documents, one `Cache`: object 5 is a stream with the same
    /// dictionary in each, but a different program. The keys must differ,
    /// or the second document would be handed the first one's font.
    #[test]
    fn a_font_stream_is_keyed_by_its_bytes_and_not_by_its_object_number() {
        use pdf_syntax::object::{ObjectIdentifier, Stream};

        fn pdf_with_object_5(program: &[u8]) -> Vec<u8> {
            let objects: Vec<Vec<u8>> = vec![
                b"<< /Type /Catalog /Pages 2 0 R >>".to_vec(),
                b"<< /Type /Pages /Kids [3 0 R] /Count 1 >>".to_vec(),
                b"<< /Type /Page /Parent 2 0 R /MediaBox [0 0 10 10] /Contents 4 0 R >>".to_vec(),
                b"<< /Length 0 >>\nstream\n\nendstream".to_vec(),
                [
                    format!("<< /Length {} >>\nstream\n", program.len()).into_bytes(),
                    program.to_vec(),
                    b"\nendstream".to_vec(),
                ]
                .concat(),
            ];
            let mut out = b"%PDF-1.7\n".to_vec();
            let mut offsets = Vec::new();
            for (i, body) in objects.iter().enumerate() {
                offsets.push(out.len());
                out.extend_from_slice(format!("{} 0 obj\n", i + 1).as_bytes());
                out.extend_from_slice(body);
                out.extend_from_slice(b"\nendobj\n");
            }
            let xref_pos = out.len();
            out.extend_from_slice(format!("xref\n0 {}\n", objects.len() + 1).as_bytes());
            out.extend_from_slice(b"0000000000 65535 f \n");
            for offset in offsets {
                out.extend_from_slice(format!("{offset:010} 00000 n \n").as_bytes());
            }
            out.extend_from_slice(
                format!(
                    "trailer\n<< /Size {} /Root 1 0 R >>\nstartxref\n{xref_pos}\n%%EOF\n",
                    objects.len() + 1
                )
                .as_bytes(),
            );
            out
        }

        // Same length, so the dictionaries are byte-identical too.
        let first = pdf_syntax::Pdf::new(pdf_with_object_5(b"FONT-PROGRAM-A")).expect("first");
        let second = pdf_syntax::Pdf::new(pdf_with_object_5(b"FONT-PROGRAM-B")).expect("second");
        let stream_of = |pdf: &pdf_syntax::Pdf| -> u128 {
            let stream = pdf
                .xref()
                .get::<Stream<'_>>(ObjectIdentifier::new(5, 0))
                .expect("object 5 is a stream");
            super::Cache::font_stream_key(b"type1", &stream)
        };
        let (a, b) = (stream_of(&first), stream_of(&second));
        assert_ne!(
            a, b,
            "two different programs under one object number share a key"
        );
        assert_eq!(a, stream_of(&first), "the key of one stream is stable");

        // And through the cache itself: the second document must not be
        // handed the first one's value.
        let cache = super::Cache::new();
        let got_a = cache.get_or_insert_some(a, || Some(1u8));
        let got_b = cache.get_or_insert_some(b, || Some(2u8));
        assert_eq!((got_a, got_b), (Some(1), Some(2)));
    }

    use super::*;

    #[test]
    fn test_generic_cache() {
        let cache = Cache::new();
        let val = cache.get_or_insert_with(12345, || Some("hello".to_string()));
        assert_eq!(val, Some("hello".to_string()));

        let cached = cache.get_or_insert_with(12345, || Some("world".to_string()));
        assert_eq!(cached, Some("hello".to_string()));
    }

    // A small weight; keeps these count-focused cases far below the byte budget.
    fn tiny(_: &i32) -> usize {
        4
    }

    #[test]
    fn test_image_cache_lru_eviction() {
        let cache = Cache::new();

        // 1. Insert 32 unique items
        for i in 0..32 {
            let val = cache.get_or_insert_image(i as u128, || Some(i), tiny);
            assert_eq!(val, Some(i));
        }

        // Verify all 32 items are still there
        for i in 0..32 {
            let cached = cache.get_or_insert_image::<i32>(i as u128, || None, tiny);
            assert_eq!(cached, Some(i));
        }

        // 2. Access key 0 to make it recently used (brings it to the back)
        let first_accessed = cache.get_or_insert_image::<i32>(0, || None, tiny);
        assert_eq!(first_accessed, Some(0));

        // 3. Insert a new item (key 32). This evicts the oldest entry, which is
        //    key 1 (key 0 was just touched).
        let val32 = cache.get_or_insert_image(32, || Some(32), tiny);
        assert_eq!(val32, Some(32));

        // Key 1 should be None (evicted)
        let evicted1 = cache.get_or_insert_image::<i32>(1, || None, tiny);
        assert_eq!(evicted1, None);

        // Key 0 should still be Some(0) (since we touched it)
        let kept0 = cache.get_or_insert_image::<i32>(0, || None, tiny);
        assert_eq!(kept0, Some(0));
    }

    #[test]
    fn test_image_cache_evicts_by_byte_budget() {
        let cache = Cache::new();
        let big = 100 * 1024 * 1024; // 100 MiB each; three => 300 MiB > 256 MiB.

        cache.get_or_insert_image(1u128, || Some(1i32), |_| big);
        cache.get_or_insert_image(2u128, || Some(2i32), |_| big);
        cache.get_or_insert_image(3u128, || Some(3i32), |_| big);

        // Only three entries (far under the 32 count cap), but their bytes exceed
        // the budget, so the least-recently-used (key 1) is evicted.
        let (count, bytes) = cache.image_cache_stats();
        assert_eq!(count, 2, "byte budget should drop the cache to two entries");
        assert!(
            bytes <= IMAGE_CACHE_MAX_BYTES,
            "total bytes must stay within budget"
        );
        assert_eq!(
            cache.get_or_insert_image::<i32>(1, || None, |_| 0),
            None,
            "oldest entry evicted by byte budget"
        );
        assert_eq!(cache.get_or_insert_image::<i32>(2, || None, |_| 0), Some(2));
        assert_eq!(cache.get_or_insert_image::<i32>(3, || None, |_| 0), Some(3));
    }

    #[test]
    fn test_image_cache_byte_budget_respects_lru_touch() {
        let cache = Cache::new();
        let big = 100 * 1024 * 1024;

        cache.get_or_insert_image(1u128, || Some(1i32), |_| big);
        cache.get_or_insert_image(2u128, || Some(2i32), |_| big);
        // Touch key 1 so it becomes most-recently-used.
        assert_eq!(cache.get_or_insert_image::<i32>(1, || None, |_| 0), Some(1));
        // Inserting key 3 pushes bytes over budget; the LRU victim is now key 2.
        cache.get_or_insert_image(3u128, || Some(3i32), |_| big);

        assert_eq!(
            cache.get_or_insert_image::<i32>(2, || None, |_| 0),
            None,
            "least-recently-used key 2 evicted, not the touched key 1"
        );
        assert_eq!(cache.get_or_insert_image::<i32>(1, || None, |_| 0), Some(1));
        assert_eq!(cache.get_or_insert_image::<i32>(3, || None, |_| 0), Some(3));
    }

    #[test]
    fn test_image_cache_oversized_item_not_retained() {
        let cache = Cache::new();
        let huge = IMAGE_CACHE_MAX_BYTES + 1; // larger than the whole budget

        // The value is still returned to the caller for this call...
        assert_eq!(
            cache.get_or_insert_image(1u128, || Some(1i32), |_| huge),
            Some(1)
        );
        // ...but it is evicted immediately, so the cache retains nothing.
        let (count, bytes) = cache.image_cache_stats();
        assert_eq!(count, 0, "an item larger than the budget is not retained");
        assert_eq!(bytes, 0);
        assert_eq!(cache.get_or_insert_image::<i32>(1, || None, |_| 0), None);
    }

    #[test]
    fn test_image_cache_stats_stay_within_both_bounds() {
        let cache = Cache::new();
        // 100 small (1 KiB) items: the count cap bounds entries to 32; bytes
        // stay tiny. This is the measurement asserting both bounds hold.
        for i in 0..100u128 {
            cache.get_or_insert_image(i, || Some(i as i32), |_| 1024);
        }
        let (count, bytes) = cache.image_cache_stats();
        assert_eq!(count, IMAGE_CACHE_CAPACITY, "count cap holds at 32 entries");
        assert_eq!(
            bytes,
            IMAGE_CACHE_CAPACITY * 1024,
            "32 retained * 1 KiB each"
        );
        assert!(bytes <= IMAGE_CACHE_MAX_BYTES);
    }
}