rightkit_pdf/
merge.rs

1//! Page-level assembly: merge several PDFs into one and select/reorder pages.
2//!
3//! Donor: Legion's `merge_page_pdfs` (lopdf renumber-and-splice recipe). Hardened: attributes a page
4//! inherits from its `/Pages` ancestors (`Resources`, `MediaBox`, `CropBox`, `Rotate`) are
5//! materialized onto the page before the old page tree is discarded, and unreachable objects
6//! (old catalogs and page trees) are pruned.
7//!
8//! Materializing never copies a value per page. A value an ancestor holds directly (a `/Resources`
9//! dictionary, a box array) is turned into **one indirect object** the first time any page below that
10//! ancestor needs it, and every such page gets a reference to it, so ten thousand pages under one
11//! node with a large `/Resources` cost ten thousand references and one copy, not ten thousand
12//! copies. A value that is already a reference, or a number or boolean, is copied as is.
13
14use std::collections::BTreeMap;
15
16use lopdf::{Dictionary, Document, Object, ObjectId};
17
18#[derive(Debug, thiserror::Error)]
19pub enum MergeError {
20    #[error("no input documents")]
21    Empty,
22    #[error("input {index} could not be parsed: {message}")]
23    Parse { index: usize, message: String },
24    #[error("input {0} is encrypted")]
25    Encrypted(usize),
26    #[error("page {0} is out of range")]
27    PageOutOfRange(u32),
28    #[error("serialize failed: {0}")]
29    Write(String),
30}
31
32const INHERITED: [&[u8]; 4] = [b"Resources", b"MediaBox", b"CropBox", b"Rotate"];
33
34fn load(index: usize, bytes: &[u8]) -> Result<Document, MergeError> {
35    // The same bounded loader as inspection: no stream (an object stream included) may decode past
36    // `MAX_DECOMPRESSED_BYTES`.
37    let doc = crate::inspect::load_pdf(bytes).map_err(|e| MergeError::Parse {
38        index,
39        message: e.to_string(),
40    })?;
41    if doc.is_encrypted() {
42        return Err(MergeError::Encrypted(index));
43    }
44    Ok(doc)
45}
46
47/// The indirect object made for an inherited value, by the ancestor that holds it and which of
48/// [`INHERITED`] it is. A page that inherits the same value gets this and nothing else.
49type SharedValues = BTreeMap<(ObjectId, usize), Object>;
50
51/// Sets on `page_id` the attributes it inherits. A direct value on an ancestor is cloned **once**
52/// into an indirect object kept in `shared` (the clone is of a value the document already holds, so
53/// the total is bounded by the document); every page below that ancestor then gets the same
54/// reference.
55fn materialize_inherited(doc: &mut Document, page_id: ObjectId, shared: &mut SharedValues) {
56    // Only asked about, never copied: which attributes the page has itself and where its parent is.
57    let (parent, has) = match doc.get_object(page_id).and_then(|o| o.as_dict()) {
58        Ok(page) => (
59            page.get(b"Parent").ok().and_then(|p| p.as_reference().ok()),
60            INHERITED.map(|key| page.has(key)),
61        ),
62        Err(_) => return,
63    };
64    let mut found: Vec<(&[u8], Object)> = Vec::new();
65    for (index, &key) in INHERITED.iter().enumerate() {
66        if has[index] {
67            continue;
68        }
69        let mut holder: Option<ObjectId> = None;
70        let mut cursor = parent;
71        let mut hops = 0;
72        while let Some(id) = cursor {
73            hops += 1;
74            if hops > 64 {
75                break;
76            }
77            let Ok(node) = doc.get_object(id).and_then(|o| o.as_dict()) else {
78                break;
79            };
80            if node.has(key) {
81                holder = Some(id);
82                break;
83            }
84            cursor = node.get(b"Parent").ok().and_then(|p| p.as_reference().ok());
85        }
86        let Some(holder) = holder else { continue };
87        let value = if let Some(value) = shared.get(&(holder, index)).cloned() {
88            value
89        } else {
90            let direct = doc
91                .get_object(holder)
92                .ok()
93                .and_then(|o| o.as_dict().ok())
94                .and_then(|node| node.get(key).ok())
95                .cloned();
96            let value = match direct {
97                None => continue,
98                // A reference or a number costs the same shared or copied.
99                Some(
100                    plain @ (Object::Reference(_)
101                    | Object::Integer(_)
102                    | Object::Real(_)
103                    | Object::Boolean(_)
104                    | Object::Null),
105                ) => plain,
106                Some(container) => Object::Reference(doc.add_object(container)),
107            };
108            shared.insert((holder, index), value.clone());
109            value
110        };
111        found.push((key, value));
112    }
113    if let Ok(Object::Dictionary(dict)) = doc.get_object_mut(page_id) {
114        for (key, value) in found {
115            dict.set(key, value);
116        }
117    }
118}
119
120/// Merges `inputs` in order into one document (one page tree, pages in input order).
121pub fn merge_pdfs(inputs: &[&[u8]]) -> Result<Vec<u8>, MergeError> {
122    if inputs.is_empty() {
123        return Err(MergeError::Empty);
124    }
125    let mut max_id: u32 = 1;
126    let mut page_ids: Vec<ObjectId> = Vec::new();
127    let mut objects: BTreeMap<ObjectId, Object> = BTreeMap::new();
128    for (index, bytes) in inputs.iter().enumerate() {
129        let mut doc = load(index, bytes)?;
130        let pages: Vec<ObjectId> = doc.get_pages().into_values().collect();
131        let mut shared = SharedValues::new();
132        for id in &pages {
133            materialize_inherited(&mut doc, *id, &mut shared);
134        }
135        // Renumber moves page ids too; recompute them from the renumbered document.
136        doc.renumber_objects_with(max_id);
137        max_id = doc.max_id + 1;
138        page_ids.extend(doc.get_pages().into_values());
139        objects.extend(doc.objects);
140    }
141    assemble(objects, page_ids)
142}
143
144/// Builds a new document containing `pages` (1-based, any order, repeats allowed) of `source`.
145pub fn select_pages(source: &[u8], pages: &[u32]) -> Result<Vec<u8>, MergeError> {
146    if pages.is_empty() {
147        return Err(MergeError::Empty);
148    }
149    let mut doc = load(0, source)?;
150    let map = doc.get_pages();
151    let mut shared = SharedValues::new();
152    for id in map.values().copied().collect::<Vec<_>>() {
153        materialize_inherited(&mut doc, id, &mut shared);
154    }
155    let mut picked = Vec::with_capacity(pages.len());
156    let mut objects = doc.objects.clone();
157    let mut next = doc.max_id;
158    for &p in pages {
159        let id = *map.get(&p).ok_or(MergeError::PageOutOfRange(p))?;
160        if picked.contains(&id) {
161            // Repeated page: clone the page dictionary under a fresh id (shares content streams).
162            next += 1;
163            let new_id = (next, 0);
164            let copy = objects
165                .get(&id)
166                .cloned()
167                .ok_or(MergeError::PageOutOfRange(p))?;
168            objects.insert(new_id, copy);
169            picked.push(new_id);
170        } else {
171            picked.push(id);
172        }
173    }
174    assemble(objects.into_iter().collect(), picked)
175}
176
177fn assemble(
178    objects: BTreeMap<ObjectId, Object>,
179    page_ids: Vec<ObjectId>,
180) -> Result<Vec<u8>, MergeError> {
181    let mut merged = Document::with_version("1.7");
182    merged.objects = objects;
183    merged.max_id = merged.objects.keys().map(|id| id.0).max().unwrap_or(1);
184    let pages_id = merged.new_object_id();
185    let kids: Vec<Object> = page_ids.iter().map(|id| Object::Reference(*id)).collect();
186    let mut pages_dict = Dictionary::new();
187    pages_dict.set("Type", Object::Name(b"Pages".to_vec()));
188    pages_dict.set("Count", Object::Integer(kids.len() as i64));
189    pages_dict.set("Kids", Object::Array(kids));
190    merged
191        .objects
192        .insert(pages_id, Object::Dictionary(pages_dict));
193    for id in &page_ids {
194        if let Ok(Object::Dictionary(dict)) = merged.get_object_mut(*id) {
195            dict.set("Parent", Object::Reference(pages_id));
196        }
197    }
198    let mut catalog = Dictionary::new();
199    catalog.set("Type", Object::Name(b"Catalog".to_vec()));
200    catalog.set("Pages", Object::Reference(pages_id));
201    let catalog_id = merged.add_object(Object::Dictionary(catalog));
202    merged.trailer.set("Root", catalog_id);
203    merged.prune_objects();
204    merged.renumber_objects();
205    merged.compress();
206    let mut out = Vec::new();
207    merged
208        .save_to(&mut out)
209        .map_err(|e| MergeError::Write(e.to_string()))?;
210    Ok(out)
211}
rightkit_pdf/merge.rs

rightkit_pdf/
merge.rs
