Refactor PDF export (#4154)

Co-authored-by: Laurenz <laurmaedje@gmail.com>

1 files changed, 166 insertions, 130 deletions

diff --git a/crates/typst-pdf/src/gradient.rs b/crates/typst-pdf/src/gradient.rs
index 576c254e..de77df33 100644
--- a/crates/typst-pdf/src/gradient.rs
+++ b/crates/typst-pdf/src/gradient.rs
@@ -1,19 +1,23 @@
+use std::collections::HashMap;
 use std::f32::consts::{PI, TAU};
 use std::sync::Arc;
 
 use ecow::eco_format;
-use pdf_writer::types::{ColorSpaceOperand, FunctionShadingType};
-use pdf_writer::writers::StreamShadingType;
-use pdf_writer::{Filter, Finish, Name, Ref};
+use pdf_writer::{
+    types::{ColorSpaceOperand, FunctionShadingType},
+    writers::StreamShadingType,
+    Filter, Finish, Name, Ref,
+};
+
 use typst::layout::{Abs, Angle, Point, Quadrant, Ratio, Transform};
 use typst::utils::Numeric;
 use typst::visualize::{
     Color, ColorSpace, Gradient, RatioOrAngle, RelativeTo, WeightedColor,
 };
 
-use crate::color::{ColorSpaceExt, PaintEncode, QuantizedColor};
-use crate::page::{PageContext, PageResource, ResourceKind, Transforms};
-use crate::{deflate, transform_to_array, AbsExt, PdfContext};
+use crate::color::{self, ColorSpaceExt, PaintEncode, QuantizedColor};
+use crate::{content, WithGlobalRefs};
+use crate::{deflate, transform_to_array, AbsExt, PdfChunk};
 
 /// A unique-transform-aspect-ratio combination that will be encoded into the
 /// PDF.
@@ -32,122 +36,144 @@ pub struct PdfGradient {
 
 /// Writes the actual gradients (shading patterns) to the PDF.
 /// This is performed once after writing all pages.
-pub(crate) fn write_gradients(ctx: &mut PdfContext) {
-    for PdfGradient { transform, aspect_ratio, gradient, angle } in
-        ctx.gradient_map.items().cloned().collect::<Vec<_>>()
-    {
-        let shading = ctx.alloc.bump();
-        ctx.gradient_refs.push(shading);
-
-        let color_space = if gradient.space().hue_index().is_some() {
-            ColorSpace::Oklab
-        } else {
-            gradient.space()
-        };
-
-        let mut shading_pattern = match &gradient {
-            Gradient::Linear(_) => {
-                let shading_function = shading_function(ctx, &gradient, color_space);
-                let mut shading_pattern = ctx.pdf.shading_pattern(shading);
-                let mut shading = shading_pattern.function_shading();
-                shading.shading_type(FunctionShadingType::Axial);
-
-                ctx.colors.write(color_space, shading.color_space(), &mut ctx.alloc);
-
-                let (mut sin, mut cos) = (angle.sin(), angle.cos());
-
-                // Scale to edges of unit square.
-                let factor = cos.abs() + sin.abs();
-                sin *= factor;
-                cos *= factor;
-
-                let (x1, y1, x2, y2): (f64, f64, f64, f64) = match angle.quadrant() {
-                    Quadrant::First => (0.0, 0.0, cos, sin),
-                    Quadrant::Second => (1.0, 0.0, cos + 1.0, sin),
-                    Quadrant::Third => (1.0, 1.0, cos + 1.0, sin + 1.0),
-                    Quadrant::Fourth => (0.0, 1.0, cos, sin + 1.0),
-                };
-
-                shading
-                    .anti_alias(gradient.anti_alias())
-                    .function(shading_function)
-                    .coords([x1 as f32, y1 as f32, x2 as f32, y2 as f32])
-                    .extend([true; 2]);
-
-                shading.finish();
-
-                shading_pattern
-            }
-            Gradient::Radial(radial) => {
-                let shading_function = shading_function(ctx, &gradient, color_space);
-                let mut shading_pattern = ctx.pdf.shading_pattern(shading);
-                let mut shading = shading_pattern.function_shading();
-                shading.shading_type(FunctionShadingType::Radial);
-
-                ctx.colors.write(color_space, shading.color_space(), &mut ctx.alloc);
-
-                shading
-                    .anti_alias(gradient.anti_alias())
-                    .function(shading_function)
-                    .coords([
-                        radial.focal_center.x.get() as f32,
-                        radial.focal_center.y.get() as f32,
-                        radial.focal_radius.get() as f32,
-                        radial.center.x.get() as f32,
-                        radial.center.y.get() as f32,
-                        radial.radius.get() as f32,
-                    ])
-                    .extend([true; 2]);
-
-                shading.finish();
-
-                shading_pattern
+pub fn write_gradients(
+    context: &WithGlobalRefs,
+) -> (PdfChunk, HashMap<PdfGradient, Ref>) {
+    let mut chunk = PdfChunk::new();
+    let mut out = HashMap::new();
+    context.resources.traverse(&mut |resources| {
+        for pdf_gradient in resources.gradients.items() {
+            if out.contains_key(pdf_gradient) {
+                continue;
             }
-            Gradient::Conic(_) => {
-                let vertices = compute_vertex_stream(&gradient, aspect_ratio);
-
-                let stream_shading_id = ctx.alloc.bump();
-                let mut stream_shading =
-                    ctx.pdf.stream_shading(stream_shading_id, &vertices);
-
-                ctx.colors.write(
-                    color_space,
-                    stream_shading.color_space(),
-                    &mut ctx.alloc,
-                );
-
-                let range = color_space.range();
-                stream_shading
-                    .bits_per_coordinate(16)
-                    .bits_per_component(16)
-                    .bits_per_flag(8)
-                    .shading_type(StreamShadingType::CoonsPatch)
-                    .decode([
-                        0.0, 1.0, 0.0, 1.0, range[0], range[1], range[2], range[3],
-                        range[4], range[5],
-                    ])
-                    .anti_alias(gradient.anti_alias())
-                    .filter(Filter::FlateDecode);
-
-                stream_shading.finish();
-
-                let mut shading_pattern = ctx.pdf.shading_pattern(shading);
-                shading_pattern.shading_ref(stream_shading_id);
-                shading_pattern
-            }
-        };
 
-        shading_pattern.matrix(transform_to_array(transform));
-    }
+            let shading = chunk.alloc();
+            out.insert(pdf_gradient.clone(), shading);
+
+            let PdfGradient { transform, aspect_ratio, gradient, angle } = pdf_gradient;
+
+            let color_space = if gradient.space().hue_index().is_some() {
+                ColorSpace::Oklab
+            } else {
+                gradient.space()
+            };
+
+            let mut shading_pattern = match &gradient {
+                Gradient::Linear(_) => {
+                    let shading_function =
+                        shading_function(gradient, &mut chunk, color_space);
+                    let mut shading_pattern = chunk.chunk.shading_pattern(shading);
+                    let mut shading = shading_pattern.function_shading();
+                    shading.shading_type(FunctionShadingType::Axial);
+
+                    color::write(
+                        color_space,
+                        shading.color_space(),
+                        &context.globals.color_functions,
+                    );
+
+                    let (mut sin, mut cos) = (angle.sin(), angle.cos());
+
+                    // Scale to edges of unit square.
+                    let factor = cos.abs() + sin.abs();
+                    sin *= factor;
+                    cos *= factor;
+
+                    let (x1, y1, x2, y2): (f64, f64, f64, f64) = match angle.quadrant() {
+                        Quadrant::First => (0.0, 0.0, cos, sin),
+                        Quadrant::Second => (1.0, 0.0, cos + 1.0, sin),
+                        Quadrant::Third => (1.0, 1.0, cos + 1.0, sin + 1.0),
+                        Quadrant::Fourth => (0.0, 1.0, cos, sin + 1.0),
+                    };
+
+                    shading
+                        .anti_alias(gradient.anti_alias())
+                        .function(shading_function)
+                        .coords([x1 as f32, y1 as f32, x2 as f32, y2 as f32])
+                        .extend([true; 2]);
+
+                    shading.finish();
+
+                    shading_pattern
+                }
+                Gradient::Radial(radial) => {
+                    let shading_function =
+                        shading_function(gradient, &mut chunk, color_space_of(gradient));
+                    let mut shading_pattern = chunk.chunk.shading_pattern(shading);
+                    let mut shading = shading_pattern.function_shading();
+                    shading.shading_type(FunctionShadingType::Radial);
+
+                    color::write(
+                        color_space,
+                        shading.color_space(),
+                        &context.globals.color_functions,
+                    );
+
+                    shading
+                        .anti_alias(gradient.anti_alias())
+                        .function(shading_function)
+                        .coords([
+                            radial.focal_center.x.get() as f32,
+                            radial.focal_center.y.get() as f32,
+                            radial.focal_radius.get() as f32,
+                            radial.center.x.get() as f32,
+                            radial.center.y.get() as f32,
+                            radial.radius.get() as f32,
+                        ])
+                        .extend([true; 2]);
+
+                    shading.finish();
+
+                    shading_pattern
+                }
+                Gradient::Conic(_) => {
+                    let vertices = compute_vertex_stream(gradient, *aspect_ratio);
+
+                    let stream_shading_id = chunk.alloc();
+                    let mut stream_shading =
+                        chunk.chunk.stream_shading(stream_shading_id, &vertices);
+
+                    color::write(
+                        color_space,
+                        stream_shading.color_space(),
+                        &context.globals.color_functions,
+                    );
+
+                    let range = color_space.range();
+                    stream_shading
+                        .bits_per_coordinate(16)
+                        .bits_per_component(16)
+                        .bits_per_flag(8)
+                        .shading_type(StreamShadingType::CoonsPatch)
+                        .decode([
+                            0.0, 1.0, 0.0, 1.0, range[0], range[1], range[2], range[3],
+                            range[4], range[5],
+                        ])
+                        .anti_alias(gradient.anti_alias())
+                        .filter(Filter::FlateDecode);
+
+                    stream_shading.finish();
+
+                    let mut shading_pattern = chunk.shading_pattern(shading);
+                    shading_pattern.shading_ref(stream_shading_id);
+                    shading_pattern
+                }
+            };
+
+            shading_pattern.matrix(transform_to_array(*transform));
+        }
+    });
+
+    (chunk, out)
 }
 
 /// Writes an expotential or stitched function that expresses the gradient.
 fn shading_function(
-    ctx: &mut PdfContext,
     gradient: &Gradient,
+    chunk: &mut PdfChunk,
     color_space: ColorSpace,
 ) -> Ref {
-    let function = ctx.alloc.bump();
+    let function = chunk.alloc();
     let mut functions = vec![];
     let mut bounds = vec![];
     let mut encode = vec![];
@@ -166,7 +192,7 @@ fn shading_function(
                 let real_t = first.1.get() * (1.0 - t) + second.1.get() * t;
 
                 let c = gradient.sample(RatioOrAngle::Ratio(Ratio::new(real_t)));
-                functions.push(single_gradient(ctx, last_c, c, color_space));
+                functions.push(single_gradient(chunk, last_c, c, color_space));
                 bounds.push(real_t as f32);
                 encode.extend([0.0, 1.0]);
                 last_c = c;
@@ -174,7 +200,7 @@ fn shading_function(
         }
 
         bounds.push(second.1.get() as f32);
-        functions.push(single_gradient(ctx, first.0, second.0, color_space));
+        functions.push(single_gradient(chunk, first.0, second.0, color_space));
         encode.extend([0.0, 1.0]);
     }
 
@@ -187,7 +213,7 @@ fn shading_function(
     bounds.pop();
 
     // Create the stitching function.
-    ctx.pdf
+    chunk
         .stitching_function(function)
         .domain([0.0, 1.0])
         .range(color_space.range())
@@ -201,14 +227,13 @@ fn shading_function(
 /// Writes an expontential function that expresses a single segment (between two
 /// stops) of a gradient.
 fn single_gradient(
-    ctx: &mut PdfContext,
+    chunk: &mut PdfChunk,
     first_color: Color,
     second_color: Color,
     color_space: ColorSpace,
 ) -> Ref {
-    let reference = ctx.alloc.bump();
-
-    ctx.pdf
+    let reference = chunk.alloc();
+    chunk
         .exponential_function(reference)
         .range(color_space.range())
         .c0(color_space.convert(first_color))
@@ -220,7 +245,12 @@ fn single_gradient(
 }
 
 impl PaintEncode for Gradient {
-    fn set_as_fill(&self, ctx: &mut PageContext, on_text: bool, transforms: Transforms) {
+    fn set_as_fill(
+        &self,
+        ctx: &mut content::Builder,
+        on_text: bool,
+        transforms: content::Transforms,
+    ) {
         ctx.reset_fill_color_space();
 
         let index = register_gradient(ctx, self, on_text, transforms);
@@ -229,15 +259,13 @@ impl PaintEncode for Gradient {
 
         ctx.content.set_fill_color_space(ColorSpaceOperand::Pattern);
         ctx.content.set_fill_pattern(None, name);
-        ctx.resources
-            .insert(PageResource::new(ResourceKind::Gradient, id), index);
     }
 
     fn set_as_stroke(
         &self,
-        ctx: &mut PageContext,
+        ctx: &mut content::Builder,
         on_text: bool,
-        transforms: Transforms,
+        transforms: content::Transforms,
     ) {
         ctx.reset_stroke_color_space();
 
@@ -247,17 +275,15 @@ impl PaintEncode for Gradient {
 
         ctx.content.set_stroke_color_space(ColorSpaceOperand::Pattern);
         ctx.content.set_stroke_pattern(None, name);
-        ctx.resources
-            .insert(PageResource::new(ResourceKind::Gradient, id), index);
     }
 }
 
 /// Deduplicates a gradient to a named PDF resource.
 fn register_gradient(
-    ctx: &mut PageContext,
+    ctx: &mut content::Builder,
     gradient: &Gradient,
     on_text: bool,
-    mut transforms: Transforms,
+    mut transforms: content::Transforms,
 ) -> usize {
     // Edge cases for strokes.
     if transforms.size.x.is_zero() {
@@ -307,7 +333,9 @@ fn register_gradient(
         angle: Gradient::correct_aspect_ratio(rotation, size.aspect_ratio()),
     };
 
-    ctx.parent.gradient_map.insert(pdf_gradient)
+    ctx.resources.colors.mark_as_used(color_space_of(gradient));
+
+    ctx.resources.gradients.insert(pdf_gradient)
 }
 
 /// Writes a single Coons Patch as defined in the PDF specification
@@ -466,3 +494,11 @@ fn compute_vertex_stream(gradient: &Gradient, aspect_ratio: Ratio) -> Arc<Vec<u8
 
     Arc::new(deflate(&vertices))
 }
+
+fn color_space_of(gradient: &Gradient) -> ColorSpace {
+    if gradient.space().hue_index().is_some() {
+        ColorSpace::Oklab
+    } else {
+        gradient.space()
+    }
+}