collective
Browse / github.com/storytold/pdfcraft / patch
perf(cos): write_full keeps the objects it parsed while numbering
cos: write_full parses every object twice (26-33% of the save work) ## Problem `write_full` (crates/cos/src/writer.rs) walks the object graph to assign numbers, keeps only the `ObjRef`, and then calls `reread` for each object when it writes. Because the reader cache holds at most 128 objects, every object is parsed twice per full save (a Save As, or any save of a reconstructed or encrypted file). ## Numbers Same 120k-object file, full save. Work done by `write_full` alone (scan phase of the probe subtracted), prototype keeps the traversal's `Arc<Object>` per object for the write pass: | | before | after prototype | change | |---|---|---|---| | classic tabl
- Status
- open
- Score
- 0.00
- Adopters
- 0 other operators
- Author
- vangel76 (imported from pull request #4876, open)
- Base
e9fc890a862ea2d8c895ea32ddecda98580483cd- Size
- +194 −20 in 6 files
- Review
- pending
- Shared
- 2026-10-28
- Id
ccbc3484e48596e4340adb89
How it went
No other operator has reported using it yet.
Diff
raw · crates/cos/src/document.rs crates/cos/src/document/objects.rs crates/cos/src/document/state.rs crates/cos/src/document/tests.rs crates/cos/src/object.rs crates/cos/src/writer.rs
@@ -23,10 +23,11 @@ pub(crate) use reader::find; #[cfg(test)] use reader::rfind; pub(crate) use state::ObjectReader;+pub(crate) use state::owned_bytes; use state::{ObjStm, Slot}; pub use state::{ReadState, SourceIdentity}; #[cfg(test)]-use state::{SECTION_READS, objstm_member, owned_bytes};+use state::{SECTION_READS, objstm_member}; use xref::XrefIndex; pub use xref::{Revision, XrefEntry}; @@ -1,5 +1,5 @@ use super::reader::find;-use super::state::{LENGTHS, LoadGuard, ObjStm, ObjectReader, ReadState, Slot, SourceIdentity, objstm_member, owned_bytes};+use super::state::{CACHE_INSERTS, LENGTHS, LoadGuard, ObjStm, ObjectReader, ReadState, Slot, SourceIdentity, objstm_member, owned_bytes}; use super::xref::{Revision, XrefEntry}; use super::*; @@ -181,6 +181,7 @@ impl Document { // The editing cache is bounded like the scan reader's: re-parsing a file-backed view // is cheap, and returned objects outlive eviction. self.cache.lock().map_err(|_| CosError::Poisoned)?.insert(num, obj.clone());+ CACHE_INSERTS.with(|c| c.set(c.get().wrapping_add(1))); self.evict_editor_cache_if_over_budget(num, &obj)?; Ok(obj) }@@ -307,7 +308,7 @@ impl Document { let mut document = self.clone(); document.cache = Arc::default(); document.objstms = Arc::default();- ObjectReader { document, decoded_count: 0 }+ ObjectReader { document, decoded_count: 0, counted_objects: 0, counted_bytes: 0 } } fn load(&self, num: u32, depth: u32) -> Result<Object, CosError> {@@ -122,6 +122,10 @@ impl std::fmt::Debug for ReadState { pub(crate) struct ObjectReader { pub(crate) document: Document, pub(super) decoded_count: usize,+ /// Objects in the parsed-object cache and the bytes of their own, as last measured. Kept+ /// as a running total so a read does not measure the whole cache again.+ pub(super) counted_objects: usize,+ pub(super) counted_bytes: usize, } impl ObjectReader {@@ -133,7 +137,9 @@ impl ObjectReader { } pub(crate) fn try_get(&mut self, reference: ObjRef) -> Result<Arc<Object>, CosError> {+ let inserts_before = CACHE_INSERTS.with(|c| c.get()); let object = self.document.try_get(reference.num);+ let inserted = CACHE_INSERTS.with(|c| c.get()).wrapping_sub(inserts_before); // A missed eviction here only costs a re-parse, so never fail the read on a poisoned // lock. The editing-cache evictors propagate instead; they guard user edits. if let Ok(mut streams) = self.document.objstms.lock() {@@ -165,14 +171,41 @@ impl ObjectReader { // have cached large objects. if let Ok(mut cache) = self.document.cache.lock() { let data = &self.document.data;- if cache.len() >= Self::OBJECTS || cache.values().fold(0usize, |n, o| n.saturating_add(owned_bytes(o, data, 0))) > Self::STREAM_BYTES {+ // A hit inserted nothing; a read that inserted exactly this object (no nested+ // /Length or object stream load) added only it. Cache length and pointer equality+ // alone cannot tell that from a read that evicted and replaced entries (a large+ // nested object makes the editor cache keep only the newcomer), so count the+ // inserts. Anything else, such as a failed load, a /Length that cached further+ // objects or a cache changed from outside, measures the whole cache again.+ let grew_by_this_object = inserted == 1+ && cache.len() == self.counted_objects.saturating_add(1)+ && matches!((&object, cache.get(&reference.num)), (Ok(read), Some(cached)) if Arc::ptr_eq(read, cached));+ if object.is_ok() && inserted == 0 && cache.len() == self.counted_objects {+ // Unchanged: a hit, or an object held by the overlay.+ } else if let (true, Ok(read)) = (grew_by_this_object, &object) {+ self.counted_bytes = self.counted_bytes.saturating_add(owned_bytes(read, data, 0));+ self.counted_objects = cache.len();+ } else {+ self.counted_bytes = cache.values().fold(0usize, |n, o| n.saturating_add(owned_bytes(o, data, 0)));+ self.counted_objects = cache.len();+ }+ if cache.len() >= Self::OBJECTS || self.counted_bytes > Self::STREAM_BYTES { cache.clear();+ self.counted_objects = 0;+ self.counted_bytes = 0; } } object } } +thread_local! {+ /// Objects this thread has put into a document's parsed-object cache. A reader compares+ /// it before and after a read: a read that inserted exactly one object cannot have+ /// evicted or replaced anything else, whatever the cache length says.+ pub(super) static CACHE_INSERTS: std::cell::Cell<u64> = const { std::cell::Cell::new(0) };+}+ #[cfg(test)] thread_local! { /// Cross-reference sections read on this thread.@@ -183,21 +216,24 @@ thread_local! { /// is measured in full. const OWNED_BYTES_DEPTH: usize = 256; -/// Bytes `object` holds of its own: string and stream data that is not a view into `data`-/// (the input file). An object nested deeper than [`OWNED_BYTES_DEPTH`] cannot be measured-/// and counts as too large to keep.-pub(super) fn owned_bytes(object: &Object, data: &Arc<Vec<u8>>, depth: usize) -> usize {+/// Bytes `object` holds of its own: string, name and dictionary-key data, and stream data that+/// is not a view into `data` (the input file). A dictionary that indexes its keys holds each+/// one twice. An object nested deeper than [`OWNED_BYTES_DEPTH`] cannot be measured and counts+/// as too large to keep.+pub(crate) fn owned_bytes(object: &Object, data: &Arc<Vec<u8>>, depth: usize) -> usize { if depth > OWNED_BYTES_DEPTH { return usize::MAX; }+ let dict_bytes = |own: usize, d: &crate::object::Dict| {+ let copies = if d.is_indexed() { 2 } else { 1 };+ d.iter().fold(own, |n, (k, v)| n.saturating_add(k.len().saturating_mul(copies)).saturating_add(owned_bytes(v, data, depth + 1)))+ }; match object { Object::String(s) => s.bytes.len(),- Object::Stream(s) => {- let own = if s.raw.shares(data) { 0 } else { s.raw.len() };- s.dict.iter().fold(own, |n, (_, v)| n.saturating_add(owned_bytes(v, data, depth + 1)))- }+ Object::Name(n) => n.len(),+ Object::Stream(s) => dict_bytes(if s.raw.shares(data) { 0 } else { s.raw.len() }, &s.dict), Object::Array(a) => a.iter().fold(0usize, |n, v| n.saturating_add(owned_bytes(v, data, depth + 1))),- Object::Dict(d) => d.iter().fold(0usize, |n, (_, v)| n.saturating_add(owned_bytes(v, data, depth + 1))),+ Object::Dict(d) => dict_bytes(0, d), _ => 0, } }@@ -693,11 +693,43 @@ fn temporary_reader_keeps_file_views_and_evicts_data_of_its_own() { // Data of its own (decrypted streams, strings) counts: past the budget the cache is // cleared after the next read, including a failed one. let owned = Object::Stream(crate::Stream::from_raw(Dict::new(), vec![b'y'; ObjectReader::STREAM_BYTES + 1]));- reader.document.cache.lock().unwrap().insert(3, Arc::new(owned));+ reader.document.cache.lock().unwrap().insert(5, Arc::new(owned)); assert_eq!(*reader.get(ObjRef::new(99, 0)), Object::Null); assert!(reader.document.cache.lock().unwrap().is_empty(), "owned data over budget is evicted"); } +#[test]+fn temporary_reader_running_total_matches_the_cache() {+ let mut bodies = vec!["<< /Type /Catalog /Pages 2 0 R >>".to_owned(), "<< /Type /Pages /Kids [] /Count 0 >>".to_owned()];+ bodies.extend((0..300).map(|i| format!("<< /Value {i} /Name ({}) >>", "s".repeat(i % 40))));+ let refs: Vec<_> = bodies.iter().map(String::as_str).collect();+ let doc = Document::open(Arc::new(build(&refs, "/Root 1 0 R"))).unwrap();+ let mut reader = doc.object_reader();+ for num in (1..=302).chain([5, 301, 99_999, 7]) {+ let _ = reader.try_get(ObjRef::new(num, 0));+ let cache = reader.document.cache.lock().unwrap();+ let measured = cache.values().fold(0usize, |n, o| n.saturating_add(owned_bytes(o, &reader.document.data, 0)));+ assert_eq!((reader.counted_objects, reader.counted_bytes), (cache.len(), measured), "after object {num}");+ assert!(cache.len() < ObjectReader::OBJECTS);+ }+}++#[test]+fn temporary_reader_recounts_when_a_read_replaces_cache_entries() {+ // One cached object, then a read of a string over the editor cache's byte budget: the+ // editor eviction keeps only the newcomer, so the cache length stays at one. The+ // string is data of its own and must not stay cached past the reader's budget.+ let big = format!("({})", "a".repeat(Document::EDITOR_CACHE_BYTES + 1));+ let bytes = build(&["<< /Type /Catalog /Pages 2 0 R >>", "<< /Type /Pages /Kids [] /Count 0 >>", "<< /Small 1 >>", &big], "/Root 1 0 R");+ let doc = Document::open(Arc::new(bytes)).unwrap();+ let mut reader = doc.object_reader();+ let _ = reader.get(ObjRef::new(3, 0));+ assert_eq!((reader.document.cache.lock().unwrap().len(), reader.counted_objects), (1, 1));+ assert!(matches!(&*reader.get(ObjRef::new(4, 0)), Object::String(_)));+ assert!(reader.document.cache.lock().unwrap().is_empty(), "a replaced entry must be counted: the string is over budget");+ assert_eq!((reader.counted_objects, reader.counted_bytes), (0, 0));+}+ #[test] fn temporary_reader_and_full_save_do_not_fill_snapshot_caches() { let mut source = Document::new_empty();@@ -172,6 +172,11 @@ impl Dict { self.entries.iter().map(|(k, v)| (k, v)) } + /// Whether the keys are held a second time, in the lookup index.+ pub(crate) fn is_indexed(&self) -> bool {+ self.index.is_some()+ }+ pub fn len(&self) -> usize { self.entries.len() }@@ -15,7 +15,7 @@ use std::io::Write; use std::sync::Arc; use crate::CosError;-use crate::document::{Document, ObjectReader, XrefEntry};+use crate::document::{Document, ObjectReader, XrefEntry, owned_bytes}; use crate::object::{Dict, ObjRef, Object, PdfString, Stream}; /// Serialize one object. Deterministic: same object → same bytes.@@ -280,6 +280,10 @@ pub fn write_full(doc: &Document, opts: &SaveOptions) -> Result<Vec<u8>, CosErro let mut map: HashMap<ObjRef, u32> = HashMap::new(); // Record the type while traversing so classifying packed objects needs no second parse. let mut order: Vec<(ObjRef, bool)> = Vec::new();+ // The parsed object of each entry in `order`, kept for the write pass while the budget+ // allows it. Past the budget the write pass reads the object again.+ let mut kept: Vec<Option<Arc<Object>>> = Vec::new();+ let mut kept_budget = KeptObjects::default(); let mut queue: VecDeque<ObjRef> = VecDeque::new(); let visit_refs = |o: &Object, queue: &mut VecDeque<ObjRef>| collect_refs(o, &mut |r| queue.push_back(r)); for key in [&b"Root"[..], b"Info", b"Encrypt"] {@@ -305,6 +309,7 @@ pub fn write_full(doc: &Document, opts: &SaveOptions) -> Result<Vec<u8>, CosErro map.insert(r, output_number(order.len())?); order.push((r, matches!(*o, Object::Stream(_)))); visit_refs(&o, &mut queue);+ kept.push(kept_budget.admit(&o, reader.document.bytes()).then_some(o)); } let version = reader.document.version().chars().take(3).collect::<String>(); let version = if opts.object_streams && version.as_str() < "1.5" { "1.5".to_string() } else { version };@@ -329,14 +334,14 @@ pub fn write_full(doc: &Document, opts: &SaveOptions) -> Result<Vec<u8>, CosErro let mut rows: BTreeMap<u32, Row> = BTreeMap::new(); // Keep only references, not renamed copies of the entire object graph. Standalone // objects still precede object streams, preserving the established output order.- let mut packed: Vec<(u32, ObjRef)> = Vec::new();+ let mut packed: Vec<(u32, ObjRef, Option<Arc<Object>>)> = Vec::new(); for (i, (r, is_stream)) in order.iter().enumerate() { let num = output_number(i)?; let in_stream = opts.object_streams && !is_stream && !encryption_objects.contains(r) && root_outside != Some(*r); if in_stream {- packed.push((num, *r));+ packed.push((num, *r, kept[i].take())); } else {- let o = renumber(&*reread(&mut reader, *r, *is_stream)?, &map);+ let o = renumber(&*written_object(&mut reader, *r, *is_stream, kept[i].take())?, &map); rows.insert(num, Row::InFile(out.len() as u64, 0)); write_indirect(num, 0, &prepared(&reader.document, r.num, num, 0, &o), &mut out); }@@ -345,9 +350,9 @@ pub fn write_full(doc: &Document, opts: &SaveOptions) -> Result<Vec<u8>, CosErro // individually: the whole stream is (§7.5.7, §7.6.2). let mut next = output_number(order.len())?; let mut batch = ObjectStreamBatch::default();- for (num, reference) in packed {+ for (num, reference, parsed) in packed { let mut encoded = Vec::new();- serialize(&renumber(&*reread(&mut reader, reference, false)?, &map), &mut encoded);+ serialize(&renumber(&*written_object(&mut reader, reference, false, parsed)?, &map), &mut encoded); if batch.count > 0 && batch.body.len().saturating_add(encoded.len()).saturating_add(1) > OBJSTM_BYTES { batch.flush(&reader.document, &mut next, &mut rows, &mut out)?; }@@ -400,6 +405,56 @@ pub fn write_full(doc: &Document, opts: &SaveOptions) -> Result<Vec<u8>, CosErro Ok(out) } +/// What `write_full` keeps of the objects it parsed while numbering them: owned bytes (see+/// [`owned_bytes`]; streams read from the file are views and cost only their dictionary) and+/// tree nodes. Beyond either limit objects are read again when written, which bounds the memory+/// of a save of a huge document.+#[derive(Default)]+struct KeptObjects {+ bytes: usize,+ nodes: usize,+}++impl KeptObjects {+ const BYTES: usize = 64 * 1024 * 1024;+ const NODES: usize = 2_000_000;++ /// Whether `object` still fits the budget; if so it is counted against it.+ fn admit(&mut self, object: &Object, data: &Arc<Vec<u8>>) -> bool {+ let bytes = self.bytes.saturating_add(owned_bytes(object, data, 0));+ let nodes = self.nodes.saturating_add(node_count(object, 0));+ if bytes > Self::BYTES || nodes > Self::NODES {+ return false;+ }+ (self.bytes, self.nodes) = (bytes, nodes);+ true+ }+}++/// Values in `object`'s tree, counting the object itself. Past the parser's nesting limit+/// the count is `usize::MAX`, so nothing that deep is kept.+fn node_count(object: &Object, depth: usize) -> usize {+ if depth > 256 {+ return usize::MAX;+ }+ let children = match object {+ Object::Stream(s) => s.dict.iter().fold(0usize, |n, (_, v)| n.saturating_add(node_count(v, depth + 1))),+ Object::Array(a) => a.iter().fold(0usize, |n, v| n.saturating_add(node_count(v, depth + 1))),+ Object::Dict(d) => d.iter().fold(0usize, |n, (_, v)| n.saturating_add(node_count(v, depth + 1))),+ _ => 0,+ };+ children.saturating_add(1)+}++/// The object to write for `r`: the one parsed while numbering when it was kept, else `r` read+/// again.+fn written_object(reader: &mut ObjectReader, r: ObjRef, is_stream: bool, kept: Option<Arc<Object>>) -> Result<Arc<Object>, CosError> {+ match kept {+ Some(o) => Ok(o),+ None => reread(reader, r, is_stream),+ }+}+ /// Object `r` read again for writing. The traversal read it once and only kept its number, so /// a second read that fails, finds nothing or finds a different kind of object must fail the /// save rather than write `null` in its place.@@ -674,6 +729,50 @@ mod tests { Lexer::new(&out, 0).object().expect("parses") } + #[test]+ fn kept_objects_stop_at_their_byte_and_node_budget() {+ let data = Arc::new(Vec::new());+ let string = Object::String(PdfString::literal(vec![b'x'; 1000]));+ let array = Object::Array(vec![Object::Int(1), Object::Null]);+ let mut kept = KeptObjects::default();+ assert!(kept.admit(&string, &data) && kept.admit(&array, &data));+ assert_eq!((kept.bytes, kept.nodes), (1000, 4));+ let mut full = KeptObjects { bytes: KeptObjects::BYTES - 999, nodes: 0 };+ assert!(!full.admit(&string, &data), "one byte over the byte budget");+ assert_eq!((full.bytes, full.nodes), (KeptObjects::BYTES - 999, 0), "a refused object is not counted");+ assert!(full.admit(&Object::Int(1), &data));+ let mut full = KeptObjects { bytes: 0, nodes: KeptObjects::NODES - 2 };+ assert!(!full.admit(&array, &data), "a node past the node budget");+ assert!(full.admit(&Object::Int(1), &data));+ let mut deep = Object::Null;+ for _ in 0..300 {+ deep = Object::Array(vec![deep]);+ }+ assert!(!KeptObjects::default().admit(&deep, &data), "nesting past the limit is never kept");+ }++ #[test]+ fn kept_objects_count_names_and_dictionary_keys() {+ let data = Arc::new(Vec::new());+ let name = Object::Name(vec![b'n'; KeptObjects::BYTES + 1]);+ let mut kept = KeptObjects::default();+ assert!(!kept.admit(&name, &data), "a name over the byte budget is not kept");+ assert_eq!((kept.bytes, kept.nodes), (0, 0));+ let mut dict = Dict::new();+ dict.set(vec![b'k'; KeptObjects::BYTES + 1], Object::Null);+ assert!(!kept.admit(&Object::Dict(dict), &data), "nor a dictionary key");+ let mut small = Dict::new();+ small.set(b"Key".to_vec(), Object::name("Value"));+ assert!(kept.admit(&Object::Dict(small), &data));+ assert_eq!(kept.bytes, 3 + 5);+ // A dictionary large enough to index its keys holds each one twice.+ let mut big = Dict::new();+ for i in 0..64 {+ big.set(format!("K{i:03}").into_bytes(), Object::Null);+ }+ assert_eq!(owned_bytes(&Object::Dict(big), &data, 0), 64 * 4 * 2);+ }+ #[test] fn a_second_read_that_differs_fails_the_save_instead_of_writing_null() { let mut doc = Document::new_empty();