typst/crates/typst-pdf/src/gradient.rs

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 typst::layout::{Abs, Angle, Point, Quadrant, Ratio, Transform};
use typst::util::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};

/// A unique-transform-aspect-ratio combination that will be encoded into the
/// PDF.
#[derive(Debug, Clone, Eq, PartialEq, Hash)]
pub struct PdfGradient {
    /// The transform to apply to the gradient.
    pub transform: Transform,
    /// The aspect ratio of the gradient.
    /// Required for aspect ratio correction.
    pub aspect_ratio: Ratio,
    /// The gradient.
    pub gradient: Gradient,
    /// The corrected angle of the gradient.
    pub angle: Angle,
}

/// 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
            }
            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));
    }
}

/// Writes an expotential or stitched function that expresses the gradient.
fn shading_function(
    ctx: &mut PdfContext,
    gradient: &Gradient,
    color_space: ColorSpace,
) -> Ref {
    let function = ctx.alloc.bump();
    let mut functions = vec![];
    let mut bounds = vec![];
    let mut encode = vec![];

    // Create the individual gradient functions for each pair of stops.
    for window in gradient.stops_ref().windows(2) {
        let (first, second) = (window[0], window[1]);

        // If we have a hue index, we will create several stops in-between
        // to make the gradient smoother without interpolation issues with
        // native color spaces.
        let mut last_c = first.0;
        if gradient.space().hue_index().is_some() {
            for i in 0..=32 {
                let t = i as f64 / 32.0;
                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));
                bounds.push(real_t as f32);
                encode.extend([0.0, 1.0]);
                last_c = c;
            }
        }

        bounds.push(second.1.get() as f32);
        functions.push(single_gradient(ctx, first.0, second.0, color_space));
        encode.extend([0.0, 1.0]);
    }

    // Special case for gradients with only two stops.
    if functions.len() == 1 {
        return functions[0];
    }

    // Remove the last bound, since it's not needed for the stitching function.
    bounds.pop();

    // Create the stitching function.
    ctx.pdf
        .stitching_function(function)
        .domain([0.0, 1.0])
        .range(color_space.range())
        .functions(functions)
        .bounds(bounds)
        .encode(encode);

    function
}

/// Writes an expontential function that expresses a single segment (between two
/// stops) of a gradient.
fn single_gradient(
    ctx: &mut PdfContext,
    first_color: Color,
    second_color: Color,
    color_space: ColorSpace,
) -> Ref {
    let reference = ctx.alloc.bump();

    ctx.pdf
        .exponential_function(reference)
        .range(color_space.range())
        .c0(color_space.convert(first_color))
        .c1(color_space.convert(second_color))
        .domain([0.0, 1.0])
        .n(1.0);

    reference
}

impl PaintEncode for Gradient {
    fn set_as_fill(&self, ctx: &mut PageContext, on_text: bool, transforms: Transforms) {
        ctx.reset_fill_color_space();

        let index = register_gradient(ctx, self, on_text, transforms);
        let id = eco_format!("Gr{index}");
        let name = Name(id.as_bytes());

        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,
        on_text: bool,
        transforms: Transforms,
    ) {
        ctx.reset_stroke_color_space();

        let index = register_gradient(ctx, self, on_text, transforms);
        let id = eco_format!("Gr{index}");
        let name = Name(id.as_bytes());

        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,
    gradient: &Gradient,
    on_text: bool,
    mut transforms: Transforms,
) -> usize {
    // Edge cases for strokes.
    if transforms.size.x.is_zero() {
        transforms.size.x = Abs::pt(1.0);
    }

    if transforms.size.y.is_zero() {
        transforms.size.y = Abs::pt(1.0);
    }
    let size = match gradient.unwrap_relative(on_text) {
        RelativeTo::Self_ => transforms.size,
        RelativeTo::Parent => transforms.container_size,
    };

    let (offset_x, offset_y) = match gradient {
        Gradient::Conic(conic) => (
            -size.x * (1.0 - conic.center.x.get() / 2.0) / 2.0,
            -size.y * (1.0 - conic.center.y.get() / 2.0) / 2.0,
        ),
        _ => (Abs::zero(), Abs::zero()),
    };

    let rotation = gradient.angle().unwrap_or_else(Angle::zero);

    let transform = match gradient.unwrap_relative(on_text) {
        RelativeTo::Self_ => transforms.transform,
        RelativeTo::Parent => transforms.container_transform,
    };

    let scale_offset = match gradient {
        Gradient::Conic(_) => 4.0_f64,
        _ => 1.0,
    };

    let pdf_gradient = PdfGradient {
        aspect_ratio: size.aspect_ratio(),
        transform: transform
            .pre_concat(Transform::translate(
                offset_x * scale_offset,
                offset_y * scale_offset,
            ))
            .pre_concat(Transform::scale(
                Ratio::new(size.x.to_pt() * scale_offset),
                Ratio::new(size.y.to_pt() * scale_offset),
            )),
        gradient: gradient.clone(),
        angle: Gradient::correct_aspect_ratio(rotation, size.aspect_ratio()),
    };

    ctx.parent.gradient_map.insert(pdf_gradient)
}

/// Writes a single Coons Patch as defined in the PDF specification
/// to a binary vec.
///
/// Structure:
///  - flag: `u8`
///  - points: `[u16; 24]`
///  - colors: `[u16; 12]`
fn write_patch(
    target: &mut Vec<u8>,
    t: f32,
    t1: f32,
    c0: [u16; 3],
    c1: [u16; 3],
    angle: Angle,
) {
    let theta = -TAU * t + angle.to_rad() as f32 + PI;
    let theta1 = -TAU * t1 + angle.to_rad() as f32 + PI;

    let (cp1, cp2) =
        control_point(Point::new(Abs::pt(0.5), Abs::pt(0.5)), 0.5, theta, theta1);

    // Push the flag
    target.push(0);

    let p1 =
        [u16::quantize(0.5, [0.0, 1.0]).to_be(), u16::quantize(0.5, [0.0, 1.0]).to_be()];

    let p2 = [
        u16::quantize(theta.cos(), [-1.0, 1.0]).to_be(),
        u16::quantize(theta.sin(), [-1.0, 1.0]).to_be(),
    ];

    let p3 = [
        u16::quantize(theta1.cos(), [-1.0, 1.0]).to_be(),
        u16::quantize(theta1.sin(), [-1.0, 1.0]).to_be(),
    ];

    let cp1 = [
        u16::quantize(cp1.x.to_f32(), [0.0, 1.0]).to_be(),
        u16::quantize(cp1.y.to_f32(), [0.0, 1.0]).to_be(),
    ];

    let cp2 = [
        u16::quantize(cp2.x.to_f32(), [0.0, 1.0]).to_be(),
        u16::quantize(cp2.y.to_f32(), [0.0, 1.0]).to_be(),
    ];

    // Push the points
    target.extend_from_slice(bytemuck::cast_slice(&[
        p1, p1, p2, p2, cp1, cp2, p3, p3, p1, p1, p1, p1,
    ]));

    let colors =
        [c0.map(u16::to_be), c0.map(u16::to_be), c1.map(u16::to_be), c1.map(u16::to_be)];

    // Push the colors.
    target.extend_from_slice(bytemuck::cast_slice(&colors));
}

fn control_point(c: Point, r: f32, angle_start: f32, angle_end: f32) -> (Point, Point) {
    let n = (TAU / (angle_end - angle_start)).abs();
    let f = ((angle_end - angle_start) / n).tan() * 4.0 / 3.0;

    let p1 = c + Point::new(
        Abs::pt((r * angle_start.cos() - f * r * angle_start.sin()) as f64),
        Abs::pt((r * angle_start.sin() + f * r * angle_start.cos()) as f64),
    );

    let p2 = c + Point::new(
        Abs::pt((r * angle_end.cos() + f * r * angle_end.sin()) as f64),
        Abs::pt((r * angle_end.sin() - f * r * angle_end.cos()) as f64),
    );

    (p1, p2)
}

#[comemo::memoize]
fn compute_vertex_stream(gradient: &Gradient, aspect_ratio: Ratio) -> Arc<Vec<u8>> {
    let Gradient::Conic(conic) = gradient else { unreachable!() };

    // Generated vertices for the Coons patches
    let mut vertices = Vec::new();

    // Correct the gradient's angle
    let angle = Gradient::correct_aspect_ratio(conic.angle, aspect_ratio);

    for window in conic.stops.windows(2) {
        let ((c0, t0), (c1, t1)) = (window[0], window[1]);

        // Precision:
        // - On an even color, insert a stop every 90deg
        // - For a hue-based color space, insert 200 stops minimum
        // - On any other, insert 20 stops minimum
        let max_dt = if c0 == c1 {
            0.25
        } else if conic.space.hue_index().is_some() {
            0.005
        } else {
            0.05
        };
        let encode_space = conic
            .space
            .hue_index()
            .map(|_| ColorSpace::Oklab)
            .unwrap_or(conic.space);
        let mut t_x = t0.get();
        let dt = (t1.get() - t0.get()).min(max_dt);

        // Special casing for sharp gradients.
        if t0 == t1 {
            write_patch(
                &mut vertices,
                t0.get() as f32,
                t1.get() as f32,
                encode_space.convert(c0),
                encode_space.convert(c1),
                angle,
            );
            continue;
        }

        while t_x < t1.get() {
            let t_next = (t_x + dt).min(t1.get());

            // The current progress in the current window.
            let t = |t| (t - t0.get()) / (t1.get() - t0.get());
            let c = Color::mix_iter(
                [WeightedColor::new(c0, 1.0 - t(t_x)), WeightedColor::new(c1, t(t_x))],
                conic.space,
            )
            .unwrap();

            let c_next = Color::mix_iter(
                [
                    WeightedColor::new(c0, 1.0 - t(t_next)),
                    WeightedColor::new(c1, t(t_next)),
                ],
                conic.space,
            )
            .unwrap();

            write_patch(
                &mut vertices,
                t_x as f32,
                t_next as f32,
                encode_space.convert(c),
                encode_space.convert(c_next),
                angle,
            );

            t_x = t_next;
        }
    }

    Arc::new(deflate(&vertices))
}