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
34pub(crate) fn 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.
49pub(crate) type 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.
55pub(crate) fn materialize_inherited(
56    doc: &mut Document,
57    page_id: ObjectId,
58    shared: &mut SharedValues,
59) {
60    // Only asked about, never copied: which attributes the page has itself and where its parent is.
61    let (parent, has) = match doc.get_object(page_id).and_then(|o| o.as_dict()) {
62        Ok(page) => (
63            page.get(b"Parent").ok().and_then(|p| p.as_reference().ok()),
64            INHERITED.map(|key| page.has(key)),
65        ),
66        Err(_) => return,
67    };
68    let mut found: Vec<(&[u8], Object)> = Vec::new();
69    for (index, &key) in INHERITED.iter().enumerate() {
70        if has[index] {
71            continue;
72        }
73        let mut holder: Option<ObjectId> = None;
74        let mut cursor = parent;
75        let mut hops = 0;
76        while let Some(id) = cursor {
77            hops += 1;
78            if hops > 64 {
79                break;
80            }
81            let Ok(node) = doc.get_object(id).and_then(|o| o.as_dict()) else {
82                break;
83            };
84            if node.has(key) {
85                holder = Some(id);
86                break;
87            }
88            cursor = node.get(b"Parent").ok().and_then(|p| p.as_reference().ok());
89        }
90        let Some(holder) = holder else { continue };
91        let value = if let Some(value) = shared.get(&(holder, index)).cloned() {
92            value
93        } else {
94            let direct = doc
95                .get_object(holder)
96                .ok()
97                .and_then(|o| o.as_dict().ok())
98                .and_then(|node| node.get(key).ok())
99                .cloned();
100            let value = match direct {
101                None => continue,
102                // A reference or a number costs the same shared or copied.
103                Some(
104                    plain @ (Object::Reference(_)
105                    | Object::Integer(_)
106                    | Object::Real(_)
107                    | Object::Boolean(_)
108                    | Object::Null),
109                ) => plain,
110                Some(container) => Object::Reference(doc.add_object(container)),
111            };
112            shared.insert((holder, index), value.clone());
113            value
114        };
115        found.push((key, value));
116    }
117    if let Ok(Object::Dictionary(dict)) = doc.get_object_mut(page_id) {
118        for (key, value) in found {
119            dict.set(key, value);
120        }
121    }
122}
123
124/// Merges `inputs` in order into one document (one page tree, pages in input order).
125pub fn merge_pdfs(inputs: &[&[u8]]) -> Result<Vec<u8>, MergeError> {
126    if inputs.is_empty() {
127        return Err(MergeError::Empty);
128    }
129    let mut max_id: u32 = 1;
130    let mut page_ids: Vec<ObjectId> = Vec::new();
131    let mut objects: BTreeMap<ObjectId, Object> = BTreeMap::new();
132    for (index, bytes) in inputs.iter().enumerate() {
133        let mut doc = load(index, bytes)?;
134        let pages: Vec<ObjectId> = doc.get_pages().into_values().collect();
135        let mut shared = SharedValues::new();
136        for id in &pages {
137            materialize_inherited(&mut doc, *id, &mut shared);
138        }
139        // Renumber moves page ids too; recompute them from the renumbered document.
140        doc.renumber_objects_with(max_id);
141        max_id = doc.max_id + 1;
142        page_ids.extend(doc.get_pages().into_values());
143        objects.extend(doc.objects);
144    }
145    assemble(objects, page_ids)
146}
147
148/// Builds a new document containing `pages` (1-based, any order, repeats allowed) of `source`.
149pub fn select_pages(source: &[u8], pages: &[u32]) -> Result<Vec<u8>, MergeError> {
150    if pages.is_empty() {
151        return Err(MergeError::Empty);
152    }
153    let mut doc = load(0, source)?;
154    let map = doc.get_pages();
155    let mut shared = SharedValues::new();
156    for id in map.values().copied().collect::<Vec<_>>() {
157        materialize_inherited(&mut doc, id, &mut shared);
158    }
159    let mut picked = Vec::with_capacity(pages.len());
160    let mut objects = doc.objects.clone();
161    let mut next = doc.max_id;
162    for &p in pages {
163        let id = *map.get(&p).ok_or(MergeError::PageOutOfRange(p))?;
164        if picked.contains(&id) {
165            // Repeated page: clone the page dictionary under a fresh id (shares content streams).
166            next += 1;
167            let new_id = (next, 0);
168            let copy = objects
169                .get(&id)
170                .cloned()
171                .ok_or(MergeError::PageOutOfRange(p))?;
172            objects.insert(new_id, copy);
173            picked.push(new_id);
174        } else {
175            picked.push(id);
176        }
177    }
178    assemble(objects.into_iter().collect(), picked)
179}
180
181fn assemble(
182    objects: BTreeMap<ObjectId, Object>,
183    page_ids: Vec<ObjectId>,
184) -> Result<Vec<u8>, MergeError> {
185    let mut merged = Document::with_version("1.7");
186    merged.objects = objects;
187    merged.max_id = merged.objects.keys().map(|id| id.0).max().unwrap_or(1);
188    let pages_id = merged.new_object_id();
189    let kids: Vec<Object> = page_ids.iter().map(|id| Object::Reference(*id)).collect();
190    let mut pages_dict = Dictionary::new();
191    pages_dict.set("Type", Object::Name(b"Pages".to_vec()));
192    pages_dict.set("Count", Object::Integer(kids.len() as i64));
193    pages_dict.set("Kids", Object::Array(kids));
194    merged
195        .objects
196        .insert(pages_id, Object::Dictionary(pages_dict));
197    for id in &page_ids {
198        if let Ok(Object::Dictionary(dict)) = merged.get_object_mut(*id) {
199            dict.set("Parent", Object::Reference(pages_id));
200        }
201    }
202    let mut catalog = Dictionary::new();
203    catalog.set("Type", Object::Name(b"Catalog".to_vec()));
204    catalog.set("Pages", Object::Reference(pages_id));
205    let catalog_id = merged.add_object(Object::Dictionary(catalog));
206    merged.trailer.set("Root", catalog_id);
207    merged.prune_objects();
208    merged.renumber_objects();
209    merged.compress();
210    let mut out = Vec::new();
211    merged
212        .save_to(&mut out)
213        .map_err(|e| MergeError::Write(e.to_string()))?;
214    Ok(out)
215}
rightkit_pdf/merge.rs

rightkit_pdf/
merge.rs
