fix(optimize): walk shared resource dictionaries and name-tree kids once (#5956)

* fix(optimize): walk shared resource dictionaries and name-tree kids once

* test(optimize): restore the file around the appended tests

* address review: import ObjRef in the optimize audit module

crates/optimize/src/audit.rs

@@ -9,7 +9,7 @@
 
 use std::collections::{HashMap, HashSet};
 
-use pdfcraft_cos::{Dict, Document, MAX_DECODED, Object, XrefEntry};
+use pdfcraft_cos::{Dict, Document, MAX_DECODED, ObjRef, Object, XrefEntry};
 
 /// The categories, in Acrobat's order.
 #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)]
@@ -91,6 +91,9 @@ const MAX_FORM_DEPTH: usize = 64;
 struct Audit<'a> {
     doc: &'a Document,
     owner: HashMap<u32, SpaceCategory>,
+    /// Resource dictionaries (by reference) already walked: one shared by many pages, forms or
+    /// appearances is walked once, as claiming is the same each time.
+    walked: HashSet<ObjRef>,
 }
 
 impl Audit<'_> {
@@ -124,6 +127,18 @@ impl Audit<'_> {
         o.map(|x| self.doc.resolve(x)).and_then(|x| x.as_dict().cloned()).unwrap_or_default()
     }
 
+    /// The resource dictionary `holder` points at (if it does), by kind: [`Self::resources`],
+    /// once per dictionary given by reference.
+    fn resources_of(&mut self, holder: Option<&Object>, depth: usize) {
+        if let Some(Object::Ref(r)) = holder
+            && !self.walked.insert(*r)
+        {
+            return;
+        }
+        let res = self.resolve_dict(holder);
+        self.resources(&res, depth);
+    }
+
     /// A resource dictionary, by kind; `depth` counts the form XObjects it is nested in.
     fn resources(&mut self, res: &Dict, depth: usize) {
         let xobjects = self.resolve_dict(res.get(b"XObject"));
@@ -135,8 +150,7 @@ impl Audit<'_> {
                     // A form: its own resources first, by kind, then the rest of it. Owning it
                     // first stops a form that draws itself.
                     self.owner.insert(r.num, SpaceCategory::XObjectForms);
-                    let inner = self.resolve_dict(s.dict.get(b"Resources"));
-                    self.resources(&inner, depth + 1);
+                    self.resources_of(s.dict.get(b"Resources"), depth + 1);
                     self.claim(&Object::Dict(s.dict.clone()), SpaceCategory::XObjectForms);
                 }
                 _ => self.claim(x, SpaceCategory::XObjectForms),
@@ -172,8 +186,7 @@ impl Audit<'_> {
                 _ => Vec::new(),
             };
             for dict in streams {
-                let res = self.resolve_dict(dict.get(b"Resources"));
-                self.resources(&res, 1);
+                self.resources_of(dict.get(b"Resources"), 1);
             }
         }
     }
@@ -216,19 +229,17 @@ fn object_stream(doc: &Document, num: u32) -> Option<u32> {
 /// Audit the space a document uses; `file_len` is the size of the file on disk (the rest is
 /// overhead).
 pub fn audit_space(doc: &Document, file_len: u64) -> Vec<SpaceUse> {
-    let mut a = Audit { doc, owner: HashMap::new() };
+    let mut a = Audit { doc, owner: HashMap::new(), walked: HashSet::new() };
     let cat = doc.root().and_then(|r| doc.get(r).as_dict().cloned()).unwrap_or_default();
     let pages = pdfcraft_model::pages(doc);
     // Order matters: the most specific owners first. Resources by kind before anything that
     // uses them: a content stream's dictionary may point at what it draws (PdfCraft's /PCAdded
     // records an added image), and form fields and comments draw with fonts and images too.
     for p in &pages {
-        let res = a.resolve_dict(p.dict.get(b"Resources"));
-        a.resources(&res, 0);
+        a.resources_of(p.dict.get(b"Resources"), 0);
     }
     let form = a.resolve_dict(cat.get(b"AcroForm"));
-    let form_resources = a.resolve_dict(form.get(b"DR"));
-    a.resources(&form_resources, 0);
+    a.resources_of(form.get(b"DR"), 0);
     for p in &pages {
         for x in annotations(doc, &p.dict) {
             a.appearance_resources(&x);

crates/optimize/src/links.rs

@@ -52,8 +52,12 @@ fn tree_entries(doc: &Document, node: &Object) -> Vec<(Vec<u8>, Object)> {
                 }
             }
         }
-        if let Some(kids) = d.get(b"Kids").map(|x| doc.resolve(x)).and_then(|x| x.as_array().cloned()) {
-            queue.extend(kids.into_iter().map(|k| (k, depth + 1)));
+        // A kids array shared by reference is read once, and only kids not yet seen are queued:
+        // nodes sharing one array can't multiply the queue.
+        if let Some(kids) =
+            d.get(b"Kids").filter(|k| !matches!(k, Object::Ref(r) if !seen.insert(*r))).map(|x| doc.resolve(x)).and_then(|x| x.as_array().cloned())
+        {
+            queue.extend(kids.into_iter().filter(|k| !matches!(k, Object::Ref(r) if seen.contains(r))).map(|k| (k, depth + 1)));
         }
     }
     out

crates/optimize/tests/optimize.rs

@@ -5613,6 +5613,70 @@ fn a_picture_whose_mask_is_not_a_stencil_is_left_alone() {
     }
 }
 
+/// A resource dictionary shared by reference between thousands of pages is walked once by the
+/// space audit, not once per page.
+#[test]
+fn space_audit_walks_a_shared_resource_dictionary_once() {
+    let mut doc = Document::new_empty();
+    let img = image(&mut doc, 4, 4, 3, false, None);
+    let mut x = Dict::new();
+    for i in 0..5_000 {
+        x.set(format!("Im{i}").into_bytes(), Object::Ref(img));
+    }
+    let mut res = Dict::new();
+    res.set(b"XObject".to_vec(), Object::Dict(x));
+    let shared = doc.add(Object::Dict(res));
+    let pages = doc.root().and_then(|r| doc.get(r).as_dict().and_then(|d| d.reference(b"Pages"))).unwrap();
+    let mut kids = Vec::new();
+    for _ in 0..2_000 {
+        let mut p = Dict::new();
+        p.set(b"Type".to_vec(), Object::name("Page"));
+        p.set(b"Parent".to_vec(), Object::Ref(pages));
+        p.set(b"MediaBox".to_vec(), Object::Array(vec![0.into(), 0.into(), 612.into(), 792.into()]));
+        p.set(b"Resources".to_vec(), Object::Ref(shared));
+        kids.push(Object::Ref(doc.add(Object::Dict(p))));
+    }
+    doc.update_dict(pages, |d| {
+        d.set(b"Count".to_vec(), Object::Int(kids.len() as i64));
+        d.set(b"Kids".to_vec(), Object::Array(kids));
+    })
+    .unwrap();
+    let started = std::time::Instant::now();
+    let audit = audit_space(&doc, 1 << 20);
+    assert!(started.elapsed() < std::time::Duration::from_secs(5), "{:?}", started.elapsed());
+    assert!(audit.iter().any(|u| u.category == SpaceCategory::Images && u.bytes > 0), "{audit:?}");
+}
+
+/// A name tree whose nodes all share one kids array (by reference) is read once per node, not
+/// once per node per kid: the queue can't grow to the nodes times the kids.
+#[test]
+fn a_name_tree_sharing_one_kids_array_is_walked_once() {
+    let mut doc = Document::new_empty();
+    let p = page(&mut doc, &[], "BT ET");
+    let s = |t: &str| Object::String(pdfcraft_cos::PdfString::literal(t.as_bytes().to_vec()));
+    let fit = Object::Array(vec![Object::Ref(p), Object::name("Fit")]);
+    let kids = doc.add(Object::Array(Vec::new()));
+    let mut nodes = Vec::new();
+    for i in 0..3_000 {
+        let mut node = Dict::new();
+        node.set(b"Kids".to_vec(), Object::Ref(kids));
+        node.set(b"Names".to_vec(), Object::Array(vec![s(&format!("d{i}")), fit.clone()]));
+        nodes.push(Object::Ref(doc.add(Object::Dict(node))));
+    }
+    doc.set(kids, Object::Array(nodes));
+    let mut tree = Dict::new();
+    tree.set(b"Kids".to_vec(), Object::Ref(kids));
+    let tree = doc.add(Object::Dict(tree));
+    let mut names = Dict::new();
+    names.set(b"Dests".to_vec(), Object::Ref(tree));
+    let root = doc.root().unwrap();
+    doc.update_dict(root, |c| c.set(b"Names".to_vec(), Object::Dict(names))).unwrap();
+    let started = std::time::Instant::now();
+    let report = optimize(&mut doc, &Settings::default()).unwrap();
+    assert!(started.elapsed() < std::time::Duration::from_secs(10), "{:?}", started.elapsed());
+    assert_eq!(report.unreferenced_dests, 3_000, "every node's destination was found, once");
+}
+
 /// Pictures run on several threads: what comes out doesn't depend on how many.
 #[test]
 fn the_same_file_comes_out_with_any_number_of_threads() {