Skip to main content

rust_pdf_box/cos/
document.rs

1use crate::cos::object::{COSDictionary, COSObject, ObjectId};
2use crate::error::{PdfError, PdfResult};
3use std::collections::{BTreeMap, BTreeSet};
4
5/// Low-level PDF document holding all COS objects and the trailer dictionary.
6#[derive(Debug, Clone)]
7pub struct COSDocument {
8    pub objects: BTreeMap<ObjectId, COSObject>,
9    pub trailer: COSDictionary,
10    pub version: String,
11    pub max_object_number: u32,
12    pub dirty_objects: BTreeSet<ObjectId>,
13}
14
15impl Default for COSDocument {
16    fn default() -> Self {
17        Self::new()
18    }
19}
20
21impl COSDocument {
22    pub fn new() -> Self {
23        Self {
24            objects: BTreeMap::new(),
25            trailer: COSDictionary::new(),
26            version: "1.7".to_string(),
27            max_object_number: 0,
28            dirty_objects: BTreeSet::new(),
29        }
30    }
31
32    /// Allocates a new, unique object id with generation 0.
33    pub fn allocate_id(&mut self) -> ObjectId {
34        self.max_object_number += 1;
35        let id = ObjectId::new(self.max_object_number, 0);
36        self.dirty_objects.insert(id);
37        id
38    }
39
40    /// Stores a new object and returns its allocated ObjectId.
41    pub fn add_object<T: Into<COSObject>>(&mut self, obj: T) -> ObjectId {
42        let id = self.allocate_id();
43        self.objects.insert(id, obj.into());
44        id
45    }
46
47    /// Sets or replaces an object with a specific ObjectId and marks it dirty.
48    pub fn set_object<T: Into<COSObject>>(&mut self, id: ObjectId, obj: T) {
49        if id.num > self.max_object_number {
50            self.max_object_number = id.num;
51        }
52        self.objects.insert(id, obj.into());
53        self.dirty_objects.insert(id);
54    }
55
56    /// Retrieves an object by its ObjectId.
57    pub fn get_object(&self, id: ObjectId) -> Option<&COSObject> {
58        self.objects.get(&id)
59    }
60
61    /// Retrieves a mutable reference to an object by its ObjectId and marks it dirty.
62    pub fn get_object_mut(&mut self, id: ObjectId) -> Option<&mut COSObject> {
63        self.dirty_objects.insert(id);
64        self.objects.get_mut(&id)
65    }
66
67    /// Resolves an object: if it's a Reference, fetches the target object; otherwise returns the object.
68    pub fn resolve<'a>(&'a self, obj: &'a COSObject) -> &'a COSObject {
69        match obj {
70            COSObject::Reference(id) => self.objects.get(id).unwrap_or(&COSObject::Null),
71            other => other,
72        }
73    }
74
75    /// Resolves a reference or returns error if not found.
76    pub fn get_resolved(&self, id: ObjectId) -> PdfResult<&COSObject> {
77        self.objects.get(&id).ok_or(PdfError::InvalidReference(id))
78    }
79
80    /// Gets the Document Catalog dictionary (/Root).
81    pub fn get_catalog(&self) -> Option<&COSDictionary> {
82        let root = self.trailer.get("Root")?;
83        let resolved = self.resolve(root);
84        resolved.as_dict()
85    }
86
87    /// Gets mutable Document Catalog dictionary.
88    pub fn get_catalog_mut(&mut self) -> Option<&mut COSDictionary> {
89        let root_id = self.trailer.get_reference("Root")?;
90        self.dirty_objects.insert(root_id);
91        self.objects.get_mut(&root_id)?.as_dict_mut()
92    }
93
94    /// Gets Document Information dictionary (/Info).
95    pub fn get_info(&self) -> Option<&COSDictionary> {
96        let info = self.trailer.get("Info")?;
97        let resolved = self.resolve(info);
98        resolved.as_dict()
99    }
100
101    /// Clears the dirty tracking set (e.g. after saving).
102    pub fn clear_dirty(&mut self) {
103        self.dirty_objects.clear();
104    }
105}