address review: read a shared name-tree node or /Kids array again from a shallower path

crates/organize/src/attachments.rs

@@ -93,7 +93,7 @@ fn entries(doc: &Document) -> Vec<(Vec<u8>, Object)> {
         return Vec::new();
     };
     let mut out = Vec::new();
-    name_tree_entries(doc, &tree, 0, &mut std::collections::HashSet::new(), &mut out);
+    name_tree_entries(doc, &tree, 0, &mut crate::import::Seen::new(), &mut out);
     out
 }
 

crates/organize/src/import.rs

@@ -382,22 +382,38 @@ fn named_dest(src: &Document, name: &[u8]) -> Option<Object> {
     }
     let names = src.resolve(catalog.get(b"Names")?);
     let tree = names.as_dict()?.get(b"Dests")?.clone();
-    name_tree_lookup(src, &tree, name, 0, &mut HashSet::new())
+    name_tree_lookup(src, &tree, name, 0, &mut Seen::new())
 }
 
-/// The `/Kids` array of a name-tree node, unless it is an indirect array already read under
-/// another node: a cyclic or shared array of direct child dictionaries would otherwise be read
-/// again under each of them, exponentially in the depth.
-fn name_tree_kids(src: &Document, d: &Dict, seen: &mut HashSet<ObjRef>) -> Option<std::sync::Arc<Object>> {
+/// Name-tree objects read so far, each with the shallowest depth it was read at.
+pub(crate) type Seen = HashMap<ObjRef, u8>;
+
+/// Whether `r` is to be read at `depth`: never read, or read only deeper. A shared node or array
+/// first met deep (where the depth limit cut its subtree) is read again from a shallower path,
+/// so its entries aren't lost; each is read at most once per depth, so work stays bounded.
+fn first_at(seen: &mut Seen, r: ObjRef, depth: u8) -> bool {
+    match seen.get(&r) {
+        Some(&d) if d <= depth => false,
+        _ => {
+            seen.insert(r, depth);
+            true
+        }
+    }
+}
+
+/// The `/Kids` array of a name-tree node at `depth`, unless it is an indirect array already read
+/// there or shallower: a cyclic or shared array of direct child dictionaries would otherwise be
+/// read again under each of them, exponentially in the depth.
+fn name_tree_kids(src: &Document, d: &Dict, depth: u8, seen: &mut Seen) -> Option<std::sync::Arc<Object>> {
     match d.get(b"Kids")? {
-        Object::Ref(r) if !seen.insert(*r) => None,
+        Object::Ref(r) if !first_at(seen, *r, depth) => None,
         kids => Some(src.resolve(kids)),
     }
 }
 
 /// Name trees are read once per node: a node listed twice (or a cycle) is skipped.
-fn name_tree_lookup(src: &Document, node: &Object, key: &[u8], depth: u8, seen: &mut HashSet<ObjRef>) -> Option<Object> {
-    if depth > 32 || node.as_ref().is_some_and(|r| !seen.insert(r)) {
+fn name_tree_lookup(src: &Document, node: &Object, key: &[u8], depth: u8, seen: &mut Seen) -> Option<Object> {
+    if depth > 32 || node.as_ref().is_some_and(|r| !first_at(seen, r, depth)) {
         return None;
     }
     #[cfg(test)]
@@ -414,7 +430,7 @@ fn name_tree_lookup(src: &Document, node: &Object, key: &[u8], depth: u8, seen:
             }
         }
     }
-    let kids = name_tree_kids(src, d, seen);
+    let kids = name_tree_kids(src, d, depth, seen);
     kids.as_deref().and_then(Object::as_array).into_iter().flatten().find_map(|k| name_tree_lookup(src, k, key, depth + 1, seen))
 }
 
@@ -426,20 +442,13 @@ pub(crate) const MAX_NAME_TREE_KEY_BYTES: usize = 64 << 20;
 /// All (key, value) pairs of a name tree, in order, each node read once, up to
 /// [`MAX_NAME_TREE_ENTRIES`] entries and [`MAX_NAME_TREE_KEY_BYTES`] of keys (one key string
 /// named by many entries is copied once per entry).
-pub(crate) fn name_tree_entries(src: &Document, node: &Object, depth: u8, seen: &mut HashSet<ObjRef>, out: &mut Vec<(Vec<u8>, Object)>) {
+pub(crate) fn name_tree_entries(src: &Document, node: &Object, depth: u8, seen: &mut Seen, out: &mut Vec<(Vec<u8>, Object)>) {
     let mut key_bytes = out.iter().map(|(k, _)| k.len()).fold(0usize, usize::saturating_add);
     name_tree_entries_within(src, node, depth, seen, out, &mut key_bytes);
 }
 
-fn name_tree_entries_within(
-    src: &Document,
-    node: &Object,
-    depth: u8,
-    seen: &mut HashSet<ObjRef>,
-    out: &mut Vec<(Vec<u8>, Object)>,
-    key_bytes: &mut usize,
-) {
-    if depth > 32 || node.as_ref().is_some_and(|r| !seen.insert(r)) {
+fn name_tree_entries_within(src: &Document, node: &Object, depth: u8, seen: &mut Seen, out: &mut Vec<(Vec<u8>, Object)>, key_bytes: &mut usize) {
+    if depth > 32 || node.as_ref().is_some_and(|r| !first_at(seen, r, depth)) {
         return;
     }
     #[cfg(test)]
@@ -460,7 +469,7 @@ fn name_tree_entries_within(
             }
         }
     }
-    let kids = name_tree_kids(src, d, seen);
+    let kids = name_tree_kids(src, d, depth, seen);
     for k in kids.as_deref().and_then(Object::as_array).into_iter().flatten() {
         name_tree_entries_within(src, k, depth + 1, seen, out, key_bytes);
     }

crates/organize/src/import/carry.rs

@@ -1,5 +1,5 @@
 use super::copy::Copier;
-use super::{Dict, Document, HashMap, HashSet, ObjRef, Object, OrganizeError, name_tree_entries};
+use super::{Dict, Document, HashMap, ObjRef, Object, OrganizeError, name_tree_entries};
 
 /// Copy a source's document-level attachments (`/Names /EmbeddedFiles`) into `dst`.
 pub(super) fn collect_attachments(dst: &mut Document, src: &Document, out: &mut Vec<(Vec<u8>, Object)>) -> Result<(), OrganizeError> {
@@ -12,7 +12,7 @@ pub(super) fn collect_attachments(dst: &mut Document, src: &Document, out: &mut
         return Ok(());
     };
     let mut entries = Vec::new();
-    name_tree_entries(src, &tree, 0, &mut HashSet::new(), &mut entries);
+    name_tree_entries(src, &tree, 0, &mut crate::import::Seen::new(), &mut entries);
     let mut copier = Copier::new(src, HashMap::new(), HashMap::new());
     for (k, v) in entries {
         let v = copier.copy_value(dst, &v, None)?;

crates/organize/src/outline.rs

@@ -194,13 +194,23 @@ fn name_tree(doc: &Document) -> HashMap<Vec<u8>, Object> {
     let Some(tree) = names.as_dict().and_then(|n| n.get(b"Dests").cloned()) else { return out };
     // Nodes still to read, with their depth. Kids go on the stack in reverse, so they are read in order.
     let mut stack = vec![(tree, 0usize)];
-    let mut seen: HashSet<ObjRef> = HashSet::new();
+    // Each node and shared `/Kids` array with the shallowest depth it was read at: one first met
+    // deep, where the depth limit cut its subtree, is read again from a shallower path, so its
+    // entries aren't lost. Each is read at most once per depth, so the work stays bounded.
+    let mut seen: HashMap<ObjRef, usize> = HashMap::new();
+    let mut first_at = |r: ObjRef, depth: usize| match seen.get(&r) {
+        Some(&d) if d <= depth => false,
+        _ => {
+            seen.insert(r, depth);
+            true
+        }
+    };
     while let Some((raw, depth)) = stack.pop() {
         if depth > NAME_TREE_DEPTH {
             continue;
         }
         if let Object::Ref(r) = raw
-            && !seen.insert(r)
+            && !first_at(r, depth)
         {
             continue;
         }
@@ -219,7 +229,7 @@ fn name_tree(doc: &Document) -> HashMap<Vec<u8>, Object> {
         // An indirect `/Kids` array shared by several nodes is read under the first one only:
         // its entries are the same, and each is read once anyway.
         let kids = match d.get(b"Kids") {
-            Some(Object::Ref(r)) if !seen.insert(*r) => None,
+            Some(Object::Ref(r)) if !first_at(*r, depth) => None,
             Some(k) => Some(doc.resolve(k)),
             None => None,
         };
@@ -746,3 +756,40 @@ pub fn set_bookmark_open(doc: &mut Document, path: &[usize], open: bool) -> Resu
     }
     Ok(())
 }
+
+#[cfg(test)]
+mod name_tree_tests {
+    use super::{NAME_TREE_DEPTH, name_tree};
+    use pdfcraft_cos::{Dict, Document, Object, PdfString};
+
+    /// A `/Kids` array shared by a node at the depth limit and by the root: first read at the
+    /// limit (its entries are past it), it is read again from the root, so its name is found.
+    #[test]
+    fn a_shared_kids_array_first_met_at_the_limit_is_read_from_the_root() {
+        let mut doc = Document::new_empty();
+        let mut leaf = Dict::new();
+        leaf.set(b"Names".to_vec(), Object::Array(vec![Object::String(PdfString::literal(b"deep".to_vec())), Object::Int(7)]));
+        let shared = doc.add(Object::Array(vec![Object::Dict(leaf)]));
+        // A chain reaching the shared array at NAME_TREE_DEPTH, read first (kids go on the stack
+        // in reverse, so the first kid is read first).
+        let mut node = Dict::new();
+        node.set(b"Kids".to_vec(), Object::Ref(shared));
+        for _ in 0..NAME_TREE_DEPTH - 1 {
+            let mut up = Dict::new();
+            up.set(b"Kids".to_vec(), Object::Array(vec![Object::Dict(node)]));
+            node = up;
+        }
+        let mut near = Dict::new();
+        near.set(b"Kids".to_vec(), Object::Ref(shared));
+        let mut root = Dict::new();
+        root.set(b"Kids".to_vec(), Object::Array(vec![Object::Dict(node), Object::Dict(near)]));
+        let mut names = Dict::new();
+        names.set(b"Dests".to_vec(), Object::Dict(root));
+        let mut catalog = Dict::new();
+        catalog.set(b"Type".to_vec(), Object::name("Catalog"));
+        catalog.set(b"Names".to_vec(), Object::Dict(names));
+        let catalog = doc.add(Object::Dict(catalog));
+        doc.trailer_mut().set(b"Root".to_vec(), Object::Ref(catalog));
+        assert_eq!(name_tree(&doc).get(b"deep".as_slice()), Some(&Object::Int(7)));
+    }
+}

crates/organize/src/tests.rs

@@ -3843,6 +3843,31 @@ fn a_shared_name_tree_kids_array_is_read_once() {
     doc.set(kids, Object::Array(vec![inner.clone(), inner]));
     let mut out = Vec::new();
     crate::take_work();
-    crate::import::name_tree_entries(&doc, &Object::Dict(node), 0, &mut std::collections::HashSet::new(), &mut out);
+    crate::import::name_tree_entries(&doc, &Object::Dict(node), 0, &mut crate::import::Seen::new(), &mut out);
     assert!(crate::take_work() <= 4, "the shared array was read again");
 }
+
+/// The same for the import's name-tree walk: a shared `/Kids` array first met past the depth
+/// limit is read again from a shallower node.
+#[test]
+fn a_shared_name_tree_kids_array_first_met_deep_is_read_again_shallower() {
+    let mut doc = Document::new_empty();
+    let mut leaf = Dict::new();
+    leaf.set(b"Names".to_vec(), Object::Array(vec![Object::String(pdfcraft_cos::PdfString::literal(b"deep".to_vec())), Object::Int(7)]));
+    let shared = doc.add(Object::Array(vec![Object::Dict(leaf)]));
+    let mut node = Dict::new();
+    node.set(b"Kids".to_vec(), Object::Ref(shared));
+    // The innermost node sits at depth 32: it reads the array, whose entries are past the limit.
+    for _ in 0..31 {
+        let mut up = Dict::new();
+        up.set(b"Kids".to_vec(), Object::Array(vec![Object::Dict(node)]));
+        node = up;
+    }
+    let mut near = Dict::new();
+    near.set(b"Kids".to_vec(), Object::Ref(shared));
+    let mut root = Dict::new();
+    root.set(b"Kids".to_vec(), Object::Array(vec![Object::Dict(node), Object::Dict(near)]));
+    let mut out = Vec::new();
+    crate::import::name_tree_entries(&doc, &Object::Dict(root), 0, &mut crate::import::Seen::new(), &mut out);
+    assert_eq!(out, vec![(b"deep".to_vec(), Object::Int(7))]);
+}