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.
//!
//! 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;

use crate::descriptor::{OpDescriptor, ParamId};
use crate::operation::{Helper, Uniform};

/// 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] {
        &[]
    }
}

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