Skip to main content

rust_pdf_box/manipulation/
merger.rs

1use crate::cos::object::{COSArray, COSDictionary, COSObject, COSStream, ObjectId};
2use crate::error::PdfResult;
3use crate::pd::document::PDDocument;
4use crate::pd::page::PDPage;
5use std::collections::BTreeMap;
6
7/// Merges multiple PDF documents into a single document with safe object ID remapping.
8pub struct PDFMerger;
9
10impl PDFMerger {
11    /// Merges multiple documents in order and returns the unified PDDocument.
12    pub fn merge(docs: &[&PDDocument]) -> PdfResult<PDDocument> {
13        let mut target_doc = PDDocument::new();
14        for &source_doc in docs {
15            Self::append_document(&mut target_doc, source_doc)?;
16        }
17        Ok(target_doc)
18    }
19
20    /// Appends all pages and referenced objects from source_doc into target_doc.
21    pub fn append_document(target: &mut PDDocument, source: &PDDocument) -> PdfResult<()> {
22        let mut id_map: BTreeMap<ObjectId, ObjectId> = BTreeMap::new();
23
24        // 1. Allocate target IDs for all objects in source
25        for &src_id in source.cos_doc.objects.keys() {
26            let new_id = target.cos_doc.allocate_id();
27            id_map.insert(src_id, new_id);
28        }
29
30        // 2. Clone and remap objects into target
31        for (&src_id, src_obj) in &source.cos_doc.objects {
32            let target_id = id_map[&src_id];
33            let remapped_obj = Self::remap_object(src_obj, &id_map);
34            target.cos_doc.set_object(target_id, remapped_obj);
35        }
36
37        // 3. Add remapped pages to target
38        for src_page in source.pages() {
39            let mut new_page_dict = if let Some(dict_obj) = src_page
40                .page_id
41                .and_then(|id| source.cos_doc.get_object(id))
42            {
43                if let Some(d) = dict_obj.as_dict() {
44                    Self::remap_dictionary(d, &id_map)
45                } else {
46                    Self::remap_dictionary(&src_page.dict, &id_map)
47                }
48            } else {
49                Self::remap_dictionary(&src_page.dict, &id_map)
50            };
51
52            // Set new parent to target pages root
53            new_page_dict.set("Parent", target.catalog.pages_id);
54
55            let page_obj_id = target.cos_doc.allocate_id();
56            let new_page = PDPage::from_dict(new_page_dict.clone(), page_obj_id);
57            target.cos_doc.set_object(page_obj_id, new_page_dict);
58            target.pages.push(new_page);
59        }
60
61        target.sync();
62        Ok(())
63    }
64
65    fn remap_object(obj: &COSObject, id_map: &BTreeMap<ObjectId, ObjectId>) -> COSObject {
66        match obj {
67            COSObject::Reference(id) => {
68                if let Some(&new_id) = id_map.get(id) {
69                    COSObject::Reference(new_id)
70                } else {
71                    COSObject::Reference(*id)
72                }
73            }
74            COSObject::Array(arr) => {
75                let mut new_arr = COSArray::new();
76                for item in arr.iter() {
77                    new_arr.push(Self::remap_object(item, id_map));
78                }
79                COSObject::Array(new_arr)
80            }
81            COSObject::Dictionary(d) => COSObject::Dictionary(Self::remap_dictionary(d, id_map)),
82            COSObject::Stream(s) => {
83                let new_dict = Self::remap_dictionary(&s.dict, id_map);
84                COSObject::Stream(COSStream::new(new_dict, s.data.clone()))
85            }
86            other => other.clone(),
87        }
88    }
89
90    fn remap_dictionary(
91        dict: &COSDictionary,
92        id_map: &BTreeMap<ObjectId, ObjectId>,
93    ) -> COSDictionary {
94        let mut new_dict = COSDictionary::new();
95        for (k, v) in &dict.entries {
96            new_dict.set(k.clone(), Self::remap_object(v, id_map));
97        }
98        new_dict
99    }
100}