Improve multi-threading support

Cargo.toml

@@ -35,6 +35,7 @@ paste = "1.0.15"
 pdf-writer = { git= "https://github.com/LaurenzV/pdf-writer", rev = "f95a19c" }
 proc-macro2 = "1.0.86"
 quote = "1.0.37"
+rayon = "1.10.0"
 resvg = "0.44.0"
 rustybuzz = "0.18.0"
 siphasher = "1.0.1"

README.md

@@ -23,6 +23,8 @@ krilla supports most features you would expect from a 2D graphics library, inclu
 - Excellent OpenType font support, supporting all major font types, including color fonts.
 - Linear, radial and sweep gradients, as well as patterns.
 - Embedding bitmap and SVG images.
+- Optional support for multi-threading via `rayon`, allowing for great speedups when creating
+compressed PDFs or PDF with lots of images.
 
 In addition to that, the library also supports the following PDF features:
 - Great subsetting for both, CFF-flavored and TTF-flavored fonts, ensuring small file sizes.

crates/krilla/Cargo.toml

@@ -13,14 +13,16 @@ exclude = ["src/tests"]
 [features]
 default = ["simple-text", "raster-images", "svg"]
 "comemo" = ["dep:comemo"]
+# Allow for multi-threaded PDF creation.
+"rayon" = ["dep:rayon"]
 # Allow for rendering SVG images and SVG-based glyphs.
 "svg" = ["dep:resvg", "dep:usvg", "dep:tiny-skia", "fontdb", "raster-images"]
 # Allow for rendering simple text without having to shape it yourself.
 "simple-text" = ["dep:rustybuzz"]
 # Add the convenience method for converting fontdb databases.
 "fontdb" = ["dep:fontdb"]
 # Allow for adding raster images to your document.
-"raster-images" = ["dep:zune-png", "dep:zune-jpeg", "dep:gif", "dep:image-webp"]
+"raster-images" = ["dep:zune-png", "dep:zune-jpeg", "dep:gif", "dep:image-webp", "dep:imagesize"]
 
 [dependencies]
 base64 = { workspace = true }
@@ -31,10 +33,11 @@ float-cmp = { workspace = true }
 fontdb = { workspace = true, optional = true }
 gif = { workspace = true, optional = true }
 image-webp = { workspace = true, optional = true }
-imagesize = { workspace = true }
+imagesize = { workspace = true, optional = true }
 miniz_oxide = { workspace = true }
 once_cell = { workspace = true }
 pdf-writer = { workspace = true }
+rayon = { workspace = true, optional = true }
 resvg = { workspace = true, optional = true }
 rustybuzz = { workspace = true, optional = true }
 siphasher = { workspace = true }

crates/krilla/src/chunk_container.rs

@@ -8,11 +8,27 @@ use pdf_writer::{Chunk, Finish, Name, Pdf, Ref, Str, TextStr};
 use std::collections::HashMap;
 use xmp_writer::{RenditionClass, XmpWriter};
 
-trait ChunkExt {
+trait ResExt {
+    fn res(&self) -> KrillaResult<&Chunk>;
+}
+
+trait WaitExt {
     fn wait(&self) -> &Chunk;
 }
 
-impl ChunkExt for Chunk {
+impl ResExt for Chunk {
+    fn res(&self) -> KrillaResult<&Chunk> {
+        Ok(self)
+    }
+}
+
+impl ResExt for KrillaResult<Chunk> {
+    fn res(&self) -> KrillaResult<&Chunk> {
+        self.as_ref().map_err(|e| e.clone())
+    }
+}
+
+impl WaitExt for Chunk {
     fn wait(&self) -> &Chunk {
         self
     }
@@ -28,7 +44,6 @@ pub struct ChunkContainer {
     pub(crate) destination_profiles: Option<(Ref, Chunk)>,
     pub(crate) struct_tree_root: Option<(Ref, Chunk)>,
 
-    pub(crate) pages: Vec<Chunk>,
     pub(crate) struct_elements: Vec<Chunk>,
     pub(crate) page_labels: Vec<Chunk>,
     pub(crate) annotations: Vec<Chunk>,
@@ -37,11 +52,12 @@ pub struct ChunkContainer {
     pub(crate) icc_profiles: Vec<Chunk>,
     pub(crate) destinations: Vec<Chunk>,
     pub(crate) ext_g_states: Vec<Chunk>,
-    pub(crate) images: Vec<Deferred<Chunk>>,
     pub(crate) masks: Vec<Chunk>,
     pub(crate) x_objects: Vec<Chunk>,
     pub(crate) shading_functions: Vec<Chunk>,
     pub(crate) patterns: Vec<Chunk>,
+    pub(crate) pages: Vec<Deferred<Chunk>>,
+    pub(crate) images: Vec<Deferred<KrillaResult<Chunk>>>,
 
     pub(crate) metadata: Option<Metadata>,
 }
@@ -85,7 +101,7 @@ impl ChunkContainer {
             ($remapper:expr, $remapped_ref:expr; $($field:expr),+) => {
                 $(
                     for chunk in $field {
-                        let chunk = chunk.wait();
+                        let chunk = chunk.wait().res()?;
                         chunks_len += chunk.len();
                         for ref_ in chunk.refs() {
                             debug_assert!(!remapper.contains_key(&ref_));
@@ -113,10 +129,10 @@ impl ChunkContainer {
         remap_field!(remapper, remapped_ref; &mut self.page_tree, &mut self.outline,
             &mut self.page_label_tree, &mut self.destination_profiles,
         &mut self.struct_tree_root);
-        remap_fields!(remapper, remapped_ref; &self.pages, &self.struct_elements, &self.page_labels,
+        remap_fields!(remapper, remapped_ref; &self.struct_elements, &self.page_labels,
             &self.annotations, &self.fonts, &self.color_spaces, &self.icc_profiles, &self.destinations,
-            &self.ext_g_states, &self.images, &self.masks, &self.x_objects, &self.shading_functions,
-            &self.patterns
+            &self.ext_g_states, &self.masks, &self.x_objects, &self.shading_functions,
+            &self.patterns, &self.pages, &self.images
         );
 
         macro_rules! write_field {
@@ -133,7 +149,7 @@ impl ChunkContainer {
             ($remapper:expr, $pdf:expr; $($field:expr),+) => {
                 $(
                     for chunk in $field {
-                        let chunk = chunk.wait();
+                        let chunk = chunk.wait().res()?;
                         chunk.renumber_into($pdf, |old| *$remapper.get(&old).unwrap());
                     }
                 )+
@@ -143,10 +159,10 @@ impl ChunkContainer {
         write_field!(remapper, &mut pdf; &self.page_tree, &self.outline,
             &self.page_label_tree, &self.destination_profiles,
         &mut self.struct_tree_root);
-        write_fields!(remapper, &mut pdf; &self.pages, &self.struct_elements, &self.page_labels,
+        write_fields!(remapper, &mut pdf; &self.struct_elements, &self.page_labels,
             &self.annotations, &self.fonts, &self.color_spaces, &self.icc_profiles, &self.destinations,
-            &self.ext_g_states, &self.images, &self.masks, &self.x_objects,
-            &self.shading_functions, &self.patterns
+            &self.ext_g_states, &self.masks, &self.x_objects,
+            &self.shading_functions, &self.patterns, &self.pages, &self.images
         );
 
         if self.metadata.as_ref().is_none_or(|m| m.title.is_none()) {

crates/krilla/src/error.rs

@@ -22,4 +22,7 @@ pub enum KrillaError {
     ///
     /// [`SerializeSettings`]: crate::SerializeSettings
     ValidationError(Vec<ValidationError>),
+    /// An image couldn't be processed properly.
+    #[cfg(feature = "raster-images")]
+    ImageError(crate::image::Image),
 }

crates/krilla/src/object/image.rs

@@ -10,6 +10,7 @@
 // TODO: CLean up and update docs
 
 use crate::color::{ICCBasedColorSpace, ICCProfile, ICCProfileWrapper, DEVICE_CMYK, DEVICE_RGB};
+use crate::error::{KrillaError, KrillaResult};
 use crate::object::color::DEVICE_GRAY;
 use crate::resource::RegisterableResource;
 use crate::serialize::SerializerContext;
@@ -128,8 +129,8 @@ impl ImageRepr {
 }
 
 impl Debug for ImageRepr {
-    fn fmt(&self, _: &mut Formatter<'_>) -> std::fmt::Result {
-        todo!()
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        write!(f, "ImageRepr {{..}}")
     }
 }
 
@@ -233,7 +234,11 @@ impl Image {
         self.0.color_space()
     }
 
-    pub(crate) fn serialize(self, sc: &mut SerializerContext, root_ref: Ref) -> Deferred<Chunk> {
+    pub(crate) fn serialize(
+        self,
+        sc: &mut SerializerContext,
+        root_ref: Ref,
+    ) -> Deferred<KrillaResult<Chunk>> {
         let soft_mask_id = sc.new_ref();
         let icc_ref = self.icc().and_then(|ic| {
             if sc
@@ -260,7 +265,12 @@ impl Image {
         Deferred::new(move || {
             let mut chunk = Chunk::new();
 
-            let repr = self.0.inner.wait().as_ref().unwrap();
+            let repr = self
+                .0
+                .inner
+                .wait()
+                .as_ref()
+                .ok_or(KrillaError::ImageError(self.clone()))?;
 
             let alpha_mask = match repr {
                 Repr::Sampled(sampled) => sampled.mask_data.as_ref().map(|mask_data| {
@@ -323,7 +333,7 @@ impl Image {
             }
             image_x_object.finish();
 
-            chunk
+            Ok(chunk)
         })
     }
 }

crates/krilla/src/object/page.rs

@@ -178,10 +178,10 @@ impl InternalPage {
     }
 
     pub(crate) fn serialize(
-        &self,
+        self,
         sc: &mut SerializerContext,
         root_ref: Ref,
-    ) -> KrillaResult<Chunk> {
+    ) -> KrillaResult<Deferred<Chunk>> {
         let mut chunk = Chunk::new();
 
         let mut annotation_refs = vec![];
@@ -238,9 +238,10 @@ impl InternalPage {
 
         page.finish();
 
-        chunk.extend(self.stream_chunk.wait());
-
-        Ok(chunk)
+        Ok(Deferred::new(move || {
+            chunk.extend(self.stream_chunk.wait());
+            chunk
+        }))
     }
 }
 

crates/krilla/src/serialize.rs

@@ -586,17 +586,17 @@ impl SerializerContext {
         }
 
         let pages = std::mem::take(&mut self.pages);
-        for (ref_, page) in &pages {
-            let chunk = page.serialize(&mut self, *ref_)?;
+        for (ref_, page) in pages {
+            let chunk = page.serialize(&mut self, ref_)?;
             self.chunk_container.pages.push(chunk);
         }
 
         if let Some(page_tree_ref) = self.page_tree_ref {
             let mut page_tree_chunk = Chunk::new();
             page_tree_chunk
                 .pages(page_tree_ref)
-                .count(pages.len() as i32)
-                .kids(pages.iter().map(|(r, _)| *r));
+                .count(self.page_infos.len() as i32)
+                .kids(self.page_infos.iter().map(|i| i.ref_));
             self.chunk_container.page_tree = Some((page_tree_ref, page_tree_chunk));
         }
 

crates/krilla/src/stream.rs

@@ -269,13 +269,11 @@ impl<'a> FilterStream<'a> {
     }
 }
 
-#[cfg_attr(feature = "comemo", comemo::memoize)]
 fn deflate_encode(data: &[u8]) -> Vec<u8> {
     const COMPRESSION_LEVEL: u8 = 6;
     miniz_oxide::deflate::compress_to_vec_zlib(data, COMPRESSION_LEVEL)
 }
 
-#[cfg_attr(feature = "comemo", comemo::memoize)]
 fn hex_encode(data: &[u8]) -> Vec<u8> {
     data.iter()
         .enumerate()

crates/krilla/src/util.rs

@@ -293,10 +293,10 @@ pub(crate) fn hash128<T: Hash + ?Sized>(value: &T) -> u128 {
     state.finish128().as_u128()
 }
 
-/// Just a stub, until we re-add the `Deferred` functionality
-/// with rayon.
+#[cfg(not(feature = "rayon"))]
 pub(crate) struct Deferred<T>(T);
 
+#[cfg(not(feature = "rayon"))]
 impl<T: Send + Sync + 'static> Deferred<T> {
     pub fn new<F>(f: F) -> Self
     where
@@ -310,46 +310,48 @@ impl<T: Send + Sync + 'static> Deferred<T> {
     }
 }
 
-// /// A value that is lazily executed on another thread.
-// ///
-// /// Execution will be started in the background and can be waited on.
-// pub(crate) struct Deferred<T>(Arc<OnceCell<T>>);
-//
-// impl<T: Send + Sync + 'static> Deferred<T> {
-//     /// Creates a new deferred value.
-//     ///
-//     /// The closure will be called on a secondary thread such that the value
-//     /// can be initialized in parallel.
-//     pub fn new<F>(f: F) -> Self
-//     where
-//         F: FnOnce() -> T + Send + Sync + 'static,
-//     {
-//         let inner = Arc::new(OnceCell::new());
-//         let cloned = Arc::clone(&inner);
-//         rayon::spawn(move || {
-//             // Initialize the value if it hasn't been initialized yet.
-//             // We do this to avoid panicking in case it was set externally.
-//             cloned.get_or_init(f);
-//         });
-//         Self(inner)
-//     }
-//
-//     /// Waits on the value to be initialized.
-//     ///
-//     /// If the value has already been initialized, this will return
-//     /// immediately. Otherwise, this will block until the value is
-//     /// initialized in another thread.
-//     pub fn wait(&self) -> &T {
-//         // Fast path if the value is already available. We don't want to yield
-//         // to rayon in that case.
-//         if let Some(value) = self.0.get() {
-//             return value;
-//         }
-//
-//         // Ensure that we yield to give the deferred value a chance to compute
-//         // single-threaded platforms (for WASM compatibility).
-//         while let Some(rayon::Yield::Executed) = rayon::yield_now() {}
-//
-//         self.0.wait()
-//     }
-// }
+/// A value that is lazily executed on another thread.
+///
+/// Execution will be started in the background and can be waited on.
+#[cfg(feature = "rayon")]
+pub(crate) struct Deferred<T>(std::sync::Arc<once_cell::sync::OnceCell<T>>);
+
+#[cfg(feature = "rayon")]
+impl<T: Send + Sync + 'static> Deferred<T> {
+    /// Creates a new deferred value.
+    ///
+    /// The closure will be called on a secondary thread such that the value
+    /// can be initialized in parallel.
+    pub fn new<F>(f: F) -> Self
+    where
+        F: FnOnce() -> T + Send + Sync + 'static,
+    {
+        let inner = std::sync::Arc::new(once_cell::sync::OnceCell::new());
+        let cloned = std::sync::Arc::clone(&inner);
+        rayon::spawn(move || {
+            // Initialize the value if it hasn't been initialized yet.
+            // We do this to avoid panicking in case it was set externally.
+            cloned.get_or_init(f);
+        });
+        Self(inner)
+    }
+
+    /// Waits on the value to be initialized.
+    ///
+    /// If the value has already been initialized, this will return
+    /// immediately. Otherwise, this will block until the value is
+    /// initialized in another thread.
+    pub fn wait(&self) -> &T {
+        // Fast path if the value is already available. We don't want to yield
+        // to rayon in that case.
+        if let Some(value) = self.0.get() {
+            return value;
+        }
+
+        // Ensure that we yield to give the deferred value a chance to compute
+        // single-threaded platforms (for WASM compatibility).
+        while let Some(rayon::Yield::Executed) = rayon::yield_now() {}
+
+        self.0.wait()
+    }
+}