DarkRoom/core/dr-pipeline/src/lens.rs

//! Coordinate-domain operations — the geometry half of the pipeline.
//!
//! # Why this is not `Operation`
//!
//! Every [`crate::operation::Operation`] is a function from colour to colour:
//! `wgsl_body` receives `c: vec3<f32>` and produces one. That shape cannot
//! express lens correction, and the reason is worth stating precisely because
//! it is what justifies a second trait rather than an extension of the first.
//!
//! Distortion does not change a pixel's value; it changes **which pixel you
//! read**. Chromatic aberration is worse still: lateral CA is a per-channel
//! radial magnification, so red, green and blue must be fetched from three
//! *different* coordinates. No function of an already-fetched `vec3<f32>` can
//! recover that — by the time a colour reaches an `Operation`, the three
//! channels have been sampled together and the information is gone.
//!
//! So a warp runs **before** the fetch, and composes into the generated
//! shader ahead of it (ARCH §5.2 places lens corrections in the geometry
//! half of the chain).
//!
//! # Inverse mapping
//!
//! A warp declares where an output pixel's colour **came from**, not where an
//! input pixel goes. This is not a stylistic choice:
//!
//! - A forward map is a *scatter* — each input pixel writes somewhere. In a
//!   compute shader that needs atomics, leaves holes where the map expands,
//!   and races where it contracts.
//! - An inverse map is a *gather* — each output pixel reads somewhere. One
//!   dispatch, one write per pixel, no contention, and hole-free by
//!   construction.
//!
//! So `undistort` is expressed as "given this output position, which source
//! position feeds it?". For a barrel-distorting lens that means the warp
//! *magnifies* the radius, which reads backwards until you remember the
//! direction is inverse.
//!
//! # Coordinate space
//!
//! Warps work in **normalised centred** coordinates: the image centre is
//! `(0, 0)`, and the radius is scaled so that `r == 1` at the corner. Both
//! properties matter.
//!
//! Note which variable carries which. The shader's `p` supplies the centring
//! and spans `±0.5 * aspect`; the corner normalisation is applied on top of it
//! by `operation::sample_source`, which publishes the result as `radius`. A
//! warp reading `length(p)` and calling it `r` would be evaluating its
//! polynomial short of where the profile was fitted.
//!
//! Centring is what makes the polynomial meaningful — lens distortion is
//! radially symmetric about the optical axis, so a formula written about any
//! other origin would need cross terms to say the same thing.
//!
//! Corner normalisation is what makes a coefficient **portable across
//! resolutions and aspect ratios**: the same value describes the lens whether
//! applied to a full-resolution export, a 512px thumbnail, or a cropped
//! frame. Normalising to the shorter edge instead — the other obvious choice
//! — would make a coefficient mean different things on a 3:2 and a 16:9 body
//! wearing the same lens, which defeats the point of a lens profile.

use std::fmt::Write as _;
use std::sync::{Arc, LazyLock};

use crate::descriptor::{Attribute, LocalizedKey, OpDescriptor, OpId, ParamDescriptor, ParamId};
use crate::operation::{Helper, Uniform};

/// Lateral chromatic aberration, as a per-channel radial scale.
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct Tca {
    pub red_scale: f32,
    pub blue_scale: f32,
}

/// What a lens profile says about one shot, in the pipeline's own types.
///
/// **Deliberately a mirror of `dr_lens::LensProfile` rather than that type
/// itself.** The dependency would have to run the wrong way: `dr-lens` carries
/// an XML parser and 5.5 MB of Lensfun data, and this crate is organised around
/// having no dependencies so that its codegen is testable without a GPU or a
/// database (ARCH §6.5a). Whatever sits above both does the conversion; it is
/// nine fields and a `match`.
///
/// Every field is independently optional because the database is: it commonly
/// carries distortion for a lens and no vignetting, or covers only part of a
/// zoom. A partial profile is useful and must not be discarded wholesale.
#[derive(Debug, Clone, Copy, PartialEq, Default)]
pub struct LensProfile {
    pub distortion: Option<crate::ops::distortion::PtLens>,
    pub tca: Option<Tca>,
    pub vignetting: Option<crate::ops::vignetting::Pa>,
}

/// TRACES: FR-DEV-3
/// The lens profile itself, as one switch a photographer can reach.
///
/// **Not an operation and not a warp** — it corrects nothing on its own. It
/// is the answer to "use the measured profile for this lens, or not", and the
/// three corrections that read the coefficients ([`crate::ops::Distortion`],
/// [`crate::ops::Aberration`] and the vignetting node) are where the
/// correction actually happens. Framing is in the capability list on the same
/// terms: something a photographer sets which is not an `Operation`.
///
/// # Why this exists at all
///
/// The profile arrives from the file's EXIF and a database, and applying it
/// silently was the whole of the interface for it. That reads as the feature
/// being absent: the corrections in the panel are manual sliders, the
/// automatic part is a line of grey text under the camera, and nothing
/// anywhere says "this is on and you may turn it off". `dr_lens`'s honesty
/// rule — an automatic correction the user cannot see is worse than one they
/// can see is unavailable — asks for a control here, not just a caption.
///
/// # Why it is a parameter rather than a flag on the session
///
/// Everything a photographer sets travels by one road: the capability list
/// feeds the generated panel, [`crate::Preset`] captures it, the sidecar
/// stores it and the undo stack replays it (FR-DEV-3c). A `bool` on the
/// session would have needed its own place in each of those four, and would
/// have been forgotten in at least one — which is exactly how the mask stack
/// came to be missing from the history. As a parameter it is in all of them
/// with nothing registered.
pub mod profile_switch {
    use super::*;

    pub const ID: OpId = OpId("lens_profile");
    pub const APPLY: ParamId = ParamId("apply");

    /// On by default: the profile is a measurement of the lens that took the
    /// photograph, so applying it is the neutral state and declining it is
    /// the edit. See [`ParamDescriptor::switch_on`].
    pub(crate) static DESCRIPTOR: LazyLock<Arc<OpDescriptor>> = LazyLock::new(|| {
        Arc::new(OpDescriptor {
            id: ID,
            label: LocalizedKey("op.lens_profile"),
            params: vec![ParamDescriptor::switch_on(
                "apply",
                "param.lens_profile.apply",
            )],
            // Optics, so it sits with the corrections it drives rather than in
            // a group of its own.
            attributes: vec![Attribute::Optics],
        })
    });

    /// This switch's description, for a caller building a control from it.
    pub fn descriptor() -> Arc<OpDescriptor> {
        DESCRIPTOR.clone()
    }
}

/// A coordinate-domain operation, applied before the source is sampled.
///
/// Object-safe for the same reason [`crate::operation::Operation`] is: the
/// graph holds `Box<dyn Warp>` in order, so the geometry chain is data.
pub trait Warp: Send + Sync {
    /// This warp's description, driving UI generation exactly as for an
    /// operation — including being owned rather than `&'static`, for the
    /// reason [`crate::descriptor::OpDescriptor`] gives.
    fn descriptor(&self) -> Arc<OpDescriptor>;

    /// Set a parameter. Values arrive already clamped to the descriptor.
    fn set_param(&mut self, id: ParamId, value: f32);

    /// Read a parameter back.
    fn param(&self, id: ParamId) -> f32;

    /// Whether this warp currently moves any pixel.
    ///
    /// A warp at neutral is omitted from the shader entirely — and if *every*
    /// warp is neutral the generated shader keeps its integer `textureLoad`
    /// path rather than paying for a bilinear sample it does not need.
    fn is_active(&self) -> bool;

    /// The WGSL body of this warp's inverse coordinate transform.
    ///
    /// Receives `p` (a `vec2<f32>`, normalised and centred per the module
    /// docs) and must leave the **source** position in `p`.
    ///
    /// A warp needing per-channel divergence writes `p_r` and `p_b` as well;
    /// they enter the block equal to `p` and are carried out of it. A warp
    /// that ignores them costs nothing — the composer drops the per-channel
    /// path when no active warp declares [`Self::splits_channels`].
    ///
    /// Uniforms are addressed by their bare declared names, as for an
    /// operation; the composer rewrites them to their prefixed fields.
    fn wgsl_body(&self) -> String;

    /// Uniform values this warp's body reads.
    fn uniforms(&self) -> Vec<Uniform>;

    /// Whether this warp moves the channels independently.
    ///
    /// True only for chromatic aberration. When no active warp declares it,
    /// the composer emits a single sample instead of three — a 3× saving in
    /// texture bandwidth for the common case of distortion alone, which at
    /// 24 MP is the difference the tile budget is measured in.
    fn splits_channels(&self) -> bool {
        false
    }

    /// Any WGSL helper functions the body calls.
    fn helpers(&self) -> &'static [Helper] {
        &[]
    }

    /// Take whatever this warp needs from a lens profile.
    ///
    /// Handed the *whole* profile rather than its own slice of it, so that the
    /// graph fanning one out does not have to know which correction wants
    /// which coefficients — the same reason an operation is handed a
    /// [`ParamId`] rather than a field. `None` clears any profile in place,
    /// which is what opening a photograph from an unrecognised lens must do:
    /// leaving the previous one standing would correct this frame for the
    /// optics of the last one.
    ///
    /// Defaulted, because a warp need not be profile-driven. Nothing about the
    /// coordinate stage requires a database behind it.
    fn set_profile(&mut self, _profile: Option<&LensProfile>) {}
}

/// The composed geometry stage: WGSL, uniforms, and what it needs from the
/// sampler.
#[derive(Debug, Clone, PartialEq, Default)]
pub struct ComposedWarp {
    /// The WGSL block computing source coordinates, or empty when no warp is
    /// active.
    pub body: String,
    /// Helper functions the body calls.
    pub helpers: Vec<Helper>,
    /// Uniform declarations, to be appended to the generated struct.
    pub uniform_fields: String,
    /// Uniform values, in declaration order.
    pub uniforms: Vec<f32>,
    /// Whether any active warp samples the channels separately.
    pub splits_channels: bool,
}

impl ComposedWarp {
    /// Whether any warp is active. When false the shader samples with an
    /// integer `textureLoad` and no interpolation at all.
    pub fn is_active(&self) -> bool {
        !self.body.is_empty()
    }
}

/// Compose the active warps into one coordinate transform.
///
/// Warps chain in order: each receives the position the previous one produced,
/// so correcting distortion and then CA composes as a single expression with
/// no intermediate buffer.
pub fn compose_warps(warps: &[Box<dyn Warp>]) -> ComposedWarp {
    let active: Vec<&dyn Warp> = warps
        .iter()
        .map(|w| w.as_ref())
        .filter(|w| w.is_active())
        .collect();

    if active.is_empty() {
        return ComposedWarp::default();
    }

    let mut out = ComposedWarp {
        splits_channels: active.iter().any(|w| w.splits_channels()),
        ..Default::default()
    };

    for warp in &active {
        let id = warp.descriptor().id.0;
        let prefix = sanitise(id);

        let warp_uniforms = warp.uniforms();
        if !warp_uniforms.is_empty() {
            let _ = writeln!(out.uniform_fields, "    // {id}");
        }
        for u in &warp_uniforms {
            let _ = writeln!(out.uniform_fields, "    {prefix}_{}: f32,", u.name);
            out.uniforms.push(u.value);
        }

        for h in warp.helpers() {
            if !out.helpers.iter().any(|e| e.name == h.name) {
                out.helpers.push(*h);
            }
        }

        let mut fragment = warp.wgsl_body();
        for u in &warp_uniforms {
            fragment = crate::operation::rewrite_uniform(
                &fragment,
                u.name,
                &format!("u.{prefix}_{}", u.name),
            );
        }

        let _ = writeln!(out.body, "\n    // ---- warp: {id} ----");
        let _ = writeln!(out.body, "    {{");
        for line in fragment.lines() {
            let _ = writeln!(out.body, "        {line}");
        }
        let _ = writeln!(out.body, "    }}");
    }

    out
}

fn sanitise(id: &str) -> String {
    id.chars()
        .map(|c| if c.is_ascii_alphanumeric() { c } else { '_' })
        .collect()
}

#[cfg(test)]
mod tests {
    use std::sync::LazyLock;

    use super::*;
    use crate::descriptor::Attribute;
    use crate::descriptor::{LocalizedKey, OpDescriptor, OpId, ParamDescriptor};

    static DESC_A: LazyLock<Arc<OpDescriptor>> = LazyLock::new(|| {
        Arc::new(OpDescriptor {
            id: OpId("warp_a"),
            label: LocalizedKey("a"),
            params: vec![ParamDescriptor::amount("amount", "a.amount")],
            attributes: vec![Attribute::Tone],
        })
    });
    static DESC_B: LazyLock<Arc<OpDescriptor>> = LazyLock::new(|| {
        Arc::new(OpDescriptor {
            id: OpId("warp_b"),
            label: LocalizedKey("b"),
            params: vec![ParamDescriptor::amount("amount", "b.amount")],
            attributes: vec![Attribute::Tone],
        })
    });

    struct Fake {
        desc: Arc<OpDescriptor>,
        amount: f32,
        splits: bool,
    }

    impl Warp for Fake {
        fn descriptor(&self) -> Arc<OpDescriptor> {
            self.desc.clone()
        }
        fn set_param(&mut self, _id: ParamId, value: f32) {
            self.amount = value;
        }
        fn param(&self, _id: ParamId) -> f32 {
            self.amount
        }
        fn is_active(&self) -> bool {
            self.amount != 0.0
        }
        fn wgsl_body(&self) -> String {
            "p = p * amount;".into()
        }
        fn uniforms(&self) -> Vec<Uniform> {
            vec![Uniform {
                name: "amount",
                value: self.amount,
            }]
        }
        fn splits_channels(&self) -> bool {
            self.splits
        }
    }

    fn fake(desc: Arc<OpDescriptor>, amount: f32, splits: bool) -> Box<dyn Warp> {
        Box::new(Fake {
            desc,
            amount,
            splits,
        })
    }

    #[test]
    fn no_active_warp_composes_to_nothing() {
        // The property that keeps the common case free: an image with no lens
        // correction must not pay for a bilinear sample.
        let composed = compose_warps(&[fake(DESC_A.clone(), 0.0, false)]);
        assert!(!composed.is_active());
        assert!(composed.uniforms.is_empty());
        assert!(!composed.splits_channels);
    }

    #[test]
    fn an_active_warp_appears_once() {
        let composed = compose_warps(&[fake(DESC_A.clone(), 2.0, false)]);
        assert!(composed.is_active());
        assert!(composed.body.contains("---- warp: warp_a ----"));
        assert!(composed.body.contains("u.warp_a_amount"));
    }

    #[test]
    fn uniforms_are_prefixed_so_warps_cannot_collide() {
        // Both fakes declare `amount`; without prefixing the generated struct
        // would carry a duplicate field and fail to compile.
        let composed = compose_warps(&[
            fake(DESC_A.clone(), 1.0, false),
            fake(DESC_B.clone(), 2.0, false),
        ]);
        assert!(composed.uniform_fields.contains("warp_a_amount: f32"));
        assert!(composed.uniform_fields.contains("warp_b_amount: f32"));
        assert_eq!(composed.uniforms, vec![1.0, 2.0]);
    }

    #[test]
    fn channel_splitting_is_requested_by_any_active_warp() {
        // One CA warp among several must switch the whole stage to the
        // three-sample path.
        let composed = compose_warps(&[
            fake(DESC_A.clone(), 1.0, false),
            fake(DESC_B.clone(), 1.0, true),
        ]);
        assert!(composed.splits_channels);
    }

    #[test]
    fn an_inactive_splitting_warp_does_not_force_three_samples() {
        // CA present but at neutral must cost nothing — otherwise every image
        // with the panel visible pays triple bandwidth.
        let composed = compose_warps(&[
            fake(DESC_A.clone(), 1.0, false),
            fake(DESC_B.clone(), 0.0, true),
        ]);
        assert!(composed.is_active());
        assert!(!composed.splits_channels);
    }

    #[test]
    fn warps_compose_in_order() {
        let composed = compose_warps(&[
            fake(DESC_A.clone(), 1.0, false),
            fake(DESC_B.clone(), 1.0, false),
        ]);
        let a = composed.body.find("warp_a").expect("a present");
        let b = composed.body.find("warp_b").expect("b present");
        assert!(a < b, "warps must chain in graph order");
    }
}