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

//! The `Operation` trait and WGSL fragment composition.
//!
//! # Composable shaders
//!
//! Each operation contributes a **WGSL fragment**: a function taking a linear
//! RGB colour and returning one. The pipeline concatenates the fragments of
//! the enabled operations into a single generated shader, run as one compute
//! dispatch. This buys the performance of a fused pass without the coupling:
//!
//! - **One texture read and one write per frame**, not one pair per operation.
//!   At 24 MP the difference is the whole frame budget.
//! - **Operations stay independent.** Adding one is a new file implementing
//!   this trait; no central shader to edit and no ordering table to update.
//! - **A disabled operation vanishes from the source** rather than costing a
//!   branch, so an image with two active adjustments compiles to a shader
//!   doing exactly two things.
//! - **Each distinct op-set compiles once** and is cached by the hash of its
//!   generated source (ARCH §5.6).
//!
//! The cost is that WGSL compile errors point at generated source, so the
//! generator emits readable, commented output — see [`compose`].

use std::fmt::Write as _;

use dr_types::{ColourSpace, Transfer};

use crate::descriptor::{OpDescriptor, ParamId, Presentation};
use crate::framing::{Framing, FRAMING_UNIFORM_FIELDS};

/// What an operation's parameters affect, for cache invalidation scoping.
///
/// Adjusting exposure must not invalidate the demosaic result; this is what
/// lets the tile cache reuse everything up to the first changed stage
/// (ARCH §5.3).
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub enum Affects {
    /// Per-pixel colour only. Everything in this milestone.
    Colour,
    /// Pixel positions — crop, rotate. Invalidates geometry-dependent caches.
    Geometry,
}

/// A single scalar a fragment reads from the generated uniform block.
///
/// Operations declare uniforms by name and value; the composer assigns them
/// slots and emits the struct. An operation never knows its own offset, which
/// is what allows fragments to be reordered or omitted freely.
#[derive(Debug, Clone, PartialEq)]
pub struct Uniform {
    /// Field name as it appears in WGSL. Prefixed with the op id by the
    /// composer, so two operations may both declare `amount`.
    pub name: &'static str,
    pub value: f32,
}

/// A develop operation.
///
/// Object-safe: the pipeline holds `Box<dyn Operation>` in graph order, so
/// order is data rather than code (ARCH §3.4).
pub trait Operation: Send + Sync {
    /// Static description, driving UI generation (FR-DEV-3a).
    fn descriptor(&self) -> &'static OpDescriptor;

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

    /// Read a parameter back, for the sidecar and for the UI's initial state.
    fn param(&self, id: ParamId) -> f32;

    /// Whether this operation currently changes the image.
    ///
    /// An operation at its neutral settings returns `false` and is omitted
    /// from the generated shader entirely. This is what makes the common case
    /// — a handful of active adjustments out of many available — cost only
    /// what is actually used.
    fn is_active(&self) -> bool;

    /// The WGSL body of this operation's transform.
    ///
    /// Receives `c` (a `vec3<f32>` of linear RGB) and must produce the
    /// result in `c`. Uniforms are addressed by the names declared in
    /// [`Self::uniforms`], accessed as `u.<prefixed_name>`; the composer
    /// rewrites them, so a fragment writes the bare name.
    ///
    /// The fragment runs inside its own block, so locals need no unique
    /// names.
    fn wgsl_body(&self) -> String;

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

    /// What this operation's parameters affect.
    fn affects(&self) -> Affects {
        Affects::Colour
    }

    /// Any WGSL helper functions the fragment calls.
    ///
    /// Emitted once per *distinct* function name even if several operations
    /// request it, so shared helpers (luminance, soft clipping) are declared
    /// exactly once.
    fn helpers(&self) -> &'static [Helper] {
        &[]
    }

    /// TRACES: FR-DEV-3a | FR-DEV-3b
    /// How this operation would like its parameters presented.
    ///
    /// `None` — the default, and the right answer for nearly every operation
    /// — means one control per parameter, chosen from its
    /// [`crate::ParamKind`]. Returning a [`Presentation`] says that several
    /// parameters form a single conceptual control and names the widget that
    /// draws it.
    ///
    /// Purely a hint. The parameters remain individually addressable
    /// scalars, so a UI that does not implement the named widget falls back
    /// to sliders and stays fully functional.
    fn presentation(&self) -> Option<Presentation> {
        None
    }
}

/// A named WGSL helper function, deduplicated across operations.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Helper {
    pub name: &'static str,
    pub source: &'static str,
}

/// The result of composing a set of operations into one shader.
#[derive(Debug, Clone, PartialEq)]
pub struct ComposedShader {
    /// Complete, compilable WGSL.
    pub source: String,
    /// Uniform values in the order the generated struct declares them.
    pub uniforms: Vec<f32>,
    /// Identifies this shader's *structure* — the op-set and their order,
    /// not their values. Two edits differing only in slider positions share
    /// a compiled pipeline and differ only in the uniform upload.
    pub structure_hash: u64,
}

/// Fields the generated uniform struct always carries, before op uniforms.
///
/// WGSL requires a uniform struct to be non-empty and 16-byte aligned; these
/// are needed by every generated shader in any case.
const BASE_UNIFORM_FIELDS: usize = 16;

/// Where an operation's own uniforms begin in the generated block.
///
/// The base fields, then framing's. Exported because `dr-gpu` writes the
/// camera matrix into the leading slots by index and would otherwise carry
/// its own copy of this arithmetic — a duplicate that silently corrupts every
/// operation's uniforms the moment either block changes size.
pub const RESERVED_UNIFORM_FIELDS: usize = BASE_UNIFORM_FIELDS + FRAMING_UNIFORM_FIELDS;

/// Compose enabled operations into a single compute shader, for the display.
///
/// Inactive operations are skipped entirely — they contribute no code, no
/// uniforms, and nothing to the structure hash.
///
/// Equivalent to [`compose_with_framing`] with neutral framing and an sRGB
/// output.
pub fn compose(ops: &[Box<dyn Operation>]) -> ComposedShader {
    compose_with_framing(ops, &Framing::new(), ColourSpace::Srgb)
}

/// TRACES: FR-EXP-2 | FR-DSP-6
/// Compose operations and framing into a single compute shader.
///
/// Framing generates the shader's **prologue** — the map from an output pixel
/// back to a source position — where [`compose`] would emit a fixed identity
/// scale. The fused-dispatch property is unaffected: a cropped, straightened
/// edit with three adjustments is still one dispatch, one read, one write.
///
/// # The output space is a parameter, not a constant
///
/// `output` decides the primaries and transfer function the last two lines of
/// the shader encode into. It is passed per composition rather than held
/// anywhere because it is a property of *this render*: the same edit goes to
/// the screen in the display's space and to a file in whatever the export asks
/// for, and neither is more authoritative than the other.
///
/// It also enters the structure hash, so the two do not collide in the
/// pipeline cache — a screen render and a Display P3 export are different
/// shaders, however identical their sliders.
pub fn compose_with_framing(
    ops: &[Box<dyn Operation>],
    framing: &Framing,
    output: ColourSpace,
) -> ComposedShader {
    let active: Vec<&dyn Operation> = ops
        .iter()
        .map(|o| o.as_ref())
        .filter(|o| o.is_active())
        .collect();

    let mut uniform_fields = String::new();
    let mut uniform_values: Vec<f32> = Vec::new();
    let mut body = String::new();
    let mut helpers: Vec<Helper> = Vec::new();

    // The base block: the camera matrix and output settings every generated
    // shader needs. Declared first so their slots are fixed regardless of
    // which operations are present.
    uniform_fields.push_str(
        "    // Camera RGB -> linear sRGB. Rows padded to vec4 for std140\n\
         \x20   // alignment; a bare mat3x3 is laid out as three vec4 anyway.\n\
         \x20   cam_to_srgb_0: vec4<f32>,\n\
         \x20   cam_to_srgb_1: vec4<f32>,\n\
         \x20   cam_to_srgb_2: vec4<f32>,\n\
         \x20   // As-shot white balance, the neutral point for the WB control.\n\
         \x20   // `.w` is not padding: it flags a non-linear source (1.0 for a\n\
         \x20   // gamma-encoded JPEG, 0.0 for demosaiced sensor data), which the\n\
         \x20   // prologue reads to decide whether to linearise.\n\
         \x20   as_shot_wb: vec4<f32>,\n",
    );
    uniform_values.resize(BASE_UNIFORM_FIELDS, 0.0);

    // Framing's block follows the base one at a fixed offset, for the same
    // reason: the prologue is emitted whether or not any operation is active,
    // so these slots cannot be positioned by the op loop below.
    uniform_fields.push_str(
        "    // Framing: the crop rect (origin, extent) and the straightening\n\
         \x20   // angle as sin/cos — a trig call per pixel would recompute a\n\
         \x20   // value that is constant across the dispatch.\n\
         \x20   crop_rect: vec4<f32>,\n\
         \x20   framing_angle: vec4<f32>,\n",
    );
    uniform_values.extend_from_slice(&framing.uniforms());

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

        // Each op's uniforms are prefixed, so two operations may both declare
        // a field called `amount` without colliding.
        let op_uniforms = op.uniforms();
        if !op_uniforms.is_empty() {
            let _ = writeln!(uniform_fields, "    // {id}");
        }
        for u in &op_uniforms {
            let _ = writeln!(uniform_fields, "    {prefix}_{}: f32,", u.name);
            uniform_values.push(u.value);
        }

        for h in op.helpers() {
            if !helpers.iter().any(|existing| existing.name == h.name) {
                helpers.push(*h);
            }
        }

        // Rewrite bare uniform names to their prefixed struct fields, so a
        // fragment is written without knowing about any other operation.
        let mut fragment = op.wgsl_body();
        for u in &op_uniforms {
            fragment = rewrite_uniform(&fragment, u.name, &format!("u.{prefix}_{}", u.name));
        }

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

    // Pad the uniform block to a 16-byte boundary. A struct whose size is not
    // a multiple of 16 is rejected by the WGSL uniform address space rules.
    let pad = (4 - (uniform_values.len() % 4)) % 4;
    for i in 0..pad {
        let _ = writeln!(uniform_fields, "    _pad{i}: f32,");
        uniform_values.push(0.0);
    }

    let mut helper_src = String::new();
    for h in &helpers {
        let _ = writeln!(helper_src, "{}\n", h.source.trim_end());
    }

    // The coordinate stage: output pixel -> source position -> colour. Emitted
    // ahead of the operation fragments, which receive the sampled `c`.
    let prologue = format!(
        "{}\n{}",
        framing.wgsl_prologue(),
        sample_source(framing.needs_interpolation())
    );
    let sampler_helper = if framing.needs_interpolation() {
        BILINEAR_HELPER
    } else {
        ""
    };

    let to_output = primaries_conversion(output);
    let encode_output = encode_output_fn(output);

    let source = format!(
        "// GENERATED — do not edit.
//
// Composed by dr-pipeline from {} active operation(s). Each block below is
// one operation's fragment, run in graph order over a linear scene-referred
// colour. Operations at neutral settings are omitted rather than branched
// over, so this shader does exactly the work the current edit requires.

struct Params {{
{uniform_fields}}}

@group(0) @binding(0) var source: texture_2d<f32>;
@group(0) @binding(1) var<uniform> u: Params;
@group(0) @binding(2) var output: texture_storage_2d<rgba8unorm, write>;

{sampler_helper}{helper_src}{encode_output}
// Display-encoded sRGB back to linear, for sources that arrive that way.
//
// A JPEG is uploaded with its bytes untouched, so its values are gamma-encoded
// where the demosaicer's are linear. Every operation below assumes linear
// scene-referred colour — exposure is a multiply, and doubling a gamma-encoded
// value is not a stop — so the encoding is undone here, once, at the only
// point where the two source kinds still differ.
fn decode_srgb(c: vec3<f32>) -> vec3<f32> {{
    let lo = c / 12.92;
    let hi = pow((max(c, vec3<f32>(0.04045)) + 0.055) / 1.055, vec3<f32>(2.4));
    return select(hi, lo, c <= vec3<f32>(0.04045));
}}

@compute @workgroup_size(8, 8, 1)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {{
    let dims = textureDimensions(output);
    if (gid.x >= dims.x || gid.y >= dims.y) {{
        return;
    }}

{prologue}
    // A non-linear source is already display-encoded; undo that so the
    // operations below see linear colour whatever the source was.
    let non_linear = u.as_shot_wb.w > 0.5;
    if (non_linear) {{
        c = decode_srgb(c);
    }}

    // As-shot white balance. Applied unconditionally, before any operation,
    // because it is part of *interpreting* the sensor rather than an edit: a
    // Bayer sensor's green photosites collect far more signal than its red
    // and blue, so raw camera-space values are strongly green and no amount
    // of later correction recovers a neutral image from them. The white
    // balance operation, when active, applies its own offset on top of this.
    //
    // A non-linear source has already had this applied in-camera; the uniform
    // is neutral there, so this is a multiply by one rather than a branch.
    c = c * u.as_shot_wb.rgb;
{body}
    // Camera space -> linear sRGB. Applied after the adjustments so white
    // balance and exposure act on sensor-native values, which is where they
    // are physically meaningful.
    //
    // Identity for a non-linear source, which is already in sRGB primaries.
    c = vec3<f32>(
        dot(u.cam_to_srgb_0.rgb, c),
        dot(u.cam_to_srgb_1.rgb, c),
        dot(u.cam_to_srgb_2.rgb, c),
    );
{to_output}
    // Clip to the output gamut and encode. The clip is last for the reason the
    // matrix above is: a colour outside sRGB is still inside a wider space, and
    // clipping before the conversion would throw it away for no one's benefit.
    c = clamp(c, vec3<f32>(0.0), vec3<f32>(1.0));
    textureStore(output, vec2<i32>(gid.xy), vec4<f32>(encode_output(c), 1.0));
}}
",
        active.len()
    );

    // Framing enters the hash by structure only — which branches its prologue
    // generated, never how far a slider moved. Dragging the crop handles must
    // reuse the compiled pipeline and re-upload uniforms.
    //
    // The output space enters it too, and must: it changes the source, so two
    // spaces sharing a hash would have the second silently rendered with the
    // first's shader — a Display P3 export that came out sRGB and said
    // otherwise.
    let structure_hash = mix(
        mix(hash_structure(&active), framing.structure_key()),
        output as u64,
    );

    ComposedShader {
        source,
        uniforms: uniform_values,
        structure_hash,
    }
}

/// The WGSL converting linear sRGB into the output space's primaries.
///
/// A constant matrix rather than a uniform: the space is chosen when the
/// shader is composed, so the numbers are known at generation time and the
/// driver can fold them into the surrounding arithmetic.
///
/// Empty for sRGB, which is the space the pipeline already works in — the
/// camera matrix converts into it, which is what `cam_to_srgb` is named for.
/// Emitting an identity there would put nine constants and three dot products
/// into the display path's shader, the one compiled most often, to compute the
/// value it already had. The identity is detected rather than special-cased by
/// name, so a space that happens to share sRGB's primaries would be spared
/// too.
fn primaries_conversion(output: ColourSpace) -> String {
    let m = output.from_linear_srgb();
    const IDENTITY: [f32; 9] = [1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0];
    // A tolerance rather than equality: the matrix is an inverse multiplied by
    // a product, so sRGB's own comes back a few ULP off the identity. A
    // millionth of a channel is four decimal places below an 8-bit step.
    if m.iter().zip(IDENTITY).all(|(a, b)| (a - b).abs() < 1e-6) {
        return String::new();
    }

    let mut out = format!(
        "\n    // Linear sRGB -> linear {}. The last colour transform before the\n\
         \x20   // encode, and the reason a colour sRGB could not hold survives\n\
         \x20   // this far: it is still inside this gamut.\n\
         \x20   c = vec3<f32>(\n",
        output.label()
    );
    // Entries below the printed precision are zero as far as the shader is
    // concerned, and a shared primary produces one every time. Snapping them
    // avoids emitting `-0.000000`, which reads as a sign error to whoever is
    // debugging a shader at the time.
    let show = |v: f32| if v.abs() < 5e-7 { 0.0 } else { v };
    for row in 0..3 {
        let _ = writeln!(
            out,
            "        dot(vec3<f32>({:.6}, {:.6}, {:.6}), c),",
            show(m[row * 3]),
            show(m[row * 3 + 1]),
            show(m[row * 3 + 2])
        );
    }
    out.push_str("    );\n");
    out
}

/// The WGSL of the output space's transfer function.
///
/// Named `encode_output` whatever the space, so the call site at the end of
/// `main` does not have to know which one it got.
fn encode_output_fn(output: ColourSpace) -> String {
    let body = match output.transfer() {
        Transfer::Srgb => "    let lo = c * 12.92;
    let hi = 1.055 * pow(max(c, vec3<f32>(0.0031308)), vec3<f32>(1.0 / 2.4)) - 0.055;
    return select(hi, lo, c <= vec3<f32>(0.0031308));"
            .to_string(),
        // No linear segment at all, so no `select`: Adobe RGB (1998) is a
        // pure power curve, and inventing a toe for it would be a different
        // space wearing its name.
        Transfer::Gamma(g) => format!("    return pow(c, vec3<f32>(1.0 / {g:.8}));"),
        Transfer::Prophoto => "    let lo = c * 16.0;
    let hi = pow(max(c, vec3<f32>(0.001953125)), vec3<f32>(1.0 / 1.8));
    return select(hi, lo, c < vec3<f32>(0.001953125));"
            .to_string(),
    };

    format!(
        "// Linear {} to its transfer function.
//
// The one place quantisation happens: everything above runs in linear f16,
// and this is the final encode (ARCH §5.2).
fn encode_output(c: vec3<f32>) -> vec3<f32> {{
{body}
}}
",
        output.label()
    )
}

/// The WGSL turning the framed source position `p` into the colour `c`.
///
/// Split out because it is the join between the coordinate stage and the
/// colour stage, and because the choice it makes — an exact integer load, or
/// a filtered sample — is the one thing the free-angle case changes.
fn sample_source(interpolate: bool) -> &'static str {
    if interpolate {
        "    // Back to texture coordinates.
    let uv_src = p / aspect + vec2<f32>(0.5);

    // Outside the source there is no pixel. A straightened frame exposes its
    // corners; render them black rather than clamping, which would smear an
    // edge pixel across them.
    if (any(uv_src < vec2<f32>(0.0)) || any(uv_src >= vec2<f32>(1.0))) {
        textureStore(output, vec2<i32>(gid.xy), vec4<f32>(0.0, 0.0, 0.0, 1.0));
        return;
    }

    // A free angle puts output pixels between source pixels. Nearest-neighbour
    // here is what makes a straightened horizon stair-step, so interpolate.
    var c = sample_bilinear(uv_src, src_dims);
"
    } else {
        "    // Back to texture coordinates.
    let uv_src = p / aspect + vec2<f32>(0.5);

    // Outside the source there is no pixel — possible once the frame has been
    // transformed at all. Render it black rather than clamping, which would
    // smear an edge pixel across the gap.
    if (any(uv_src < vec2<f32>(0.0)) || any(uv_src >= vec2<f32>(1.0))) {
        textureStore(output, vec2<i32>(gid.xy), vec4<f32>(0.0, 0.0, 0.0, 1.0));
        return;
    }

    // Every output pixel lands on a source pixel, so load it directly: exact,
    // and with no interpolation to soften detail.
    let coord = min(vec2<i32>(uv_src * vec2<f32>(src_dims)), vec2<i32>(src_dims) - vec2<i32>(1));
    var c = textureLoad(source, coord, 0).rgb;
"
    }
}

/// Bilinear sampling against an unfiltered `texture_2d`.
///
/// Hand-rolled rather than done with a sampler: the source is bound as a plain
/// texture, and adding a sampler for the straightening case alone would change
/// a bind group layout that every pass shares.
const BILINEAR_HELPER: &str = "fn sample_bilinear(uv: vec2<f32>, dims: vec2<u32>) -> vec3<f32> {
    let last = vec2<i32>(dims) - vec2<i32>(1);

    // Half-texel offset: sample positions are texel *centres*. Without it the
    // image shifts by half a pixel and every rotation comes out slightly soft.
    let q = uv * vec2<f32>(dims) - vec2<f32>(0.5);
    let base = floor(q);
    let f = q - base;
    let i0 = clamp(vec2<i32>(base), vec2<i32>(0), last);
    let i1 = min(i0 + vec2<i32>(1), last);

    let c00 = textureLoad(source, vec2<i32>(i0.x, i0.y), 0).rgb;
    let c10 = textureLoad(source, vec2<i32>(i1.x, i0.y), 0).rgb;
    let c01 = textureLoad(source, vec2<i32>(i0.x, i1.y), 0).rgb;
    let c11 = textureLoad(source, vec2<i32>(i1.x, i1.y), 0).rgb;

    return mix(mix(c00, c10, f.x), mix(c01, c11, f.x), f.y);
}

";

/// Fold a value into a hash. FNV-1a's mixing step, over eight bytes.
fn mix(mut h: u64, value: u64) -> u64 {
    for byte in value.to_le_bytes() {
        h ^= u64::from(byte);
        h = h.wrapping_mul(0x100_0000_01b3);
    }
    h
}

/// Hash the op-set and order — the structure, not the values.
///
/// Two edits with the same operations at different slider positions produce
/// the same hash and reuse one compiled pipeline (ARCH §6.13: the hash is
/// over integer state only, so it is exactly deterministic).
fn hash_structure(active: &[&dyn Operation]) -> u64 {
    // FNV-1a: no dependency, stable across runs and platforms, which the
    // shader cache key requires.
    let mut h: u64 = 0xcbf2_9ce4_8422_2325;
    for op in active {
        for byte in op.descriptor().id.0.as_bytes() {
            h ^= u64::from(*byte);
            h = h.wrapping_mul(0x100_0000_01b3);
        }
        // A separator, so ["ab", "c"] and ["a", "bc"] differ.
        h ^= 0xff;
        h = h.wrapping_mul(0x100_0000_01b3);
    }
    h
}

/// Replace whole-word occurrences of `name` with `replacement`.
///
/// Whole-word matching matters: an operation with uniforms `amount` and
/// `amount_hi` must not have the first rewrite corrupt the second.
///
/// Shared with [`crate::lens`], which prefixes its uniforms by the same rule
/// and must not diverge from it.
/// Comments are skipped. A fragment explaining what `factor` does should not
/// have its prose rewritten to `u.saturation_factor` — the generated source
/// is meant to be read when a shader fails to compile, and mangled comments
/// make that harder rather than easier.
pub(crate) fn rewrite_uniform(src: &str, name: &str, replacement: &str) -> String {
    let mut out = String::with_capacity(src.len());
    let bytes = src.as_bytes();
    let mut i = 0;
    // Tracks whether the cursor sits inside a `//` comment. WGSL fragments
    // use line comments only, so this needs no block-comment handling.
    let mut in_comment = false;

    while i < src.len() {
        if bytes[i] == b'\n' {
            in_comment = false;
        } else if !in_comment && src[i..].starts_with("//") {
            in_comment = true;
        }

        if !in_comment && src[i..].starts_with(name) {
            let before_ok = i == 0 || !is_ident_byte(bytes[i - 1]);
            let after = i + name.len();
            let after_ok = after >= src.len() || !is_ident_byte(bytes[after]);
            if before_ok && after_ok {
                out.push_str(replacement);
                i = after;
                continue;
            }
        }
        // Push one full character, not one byte, so non-ASCII in a comment
        // does not split a UTF-8 sequence.
        let ch = src[i..].chars().next().expect("in bounds");
        out.push(ch);
        i += ch.len_utf8();
    }
    out
}

fn is_ident_byte(b: u8) -> bool {
    b.is_ascii_alphanumeric() || b == b'_'
}

/// Make an operation id safe to embed in a WGSL identifier.
fn sanitise(id: &str) -> String {
    id.chars()
        .map(|c| if c.is_ascii_alphanumeric() { c } else { '_' })
        .collect()
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::descriptor::{LocalizedKey, OpId, ParamDescriptor};

    static DESC_A: OpDescriptor = OpDescriptor {
        id: OpId("op_a"),
        label: LocalizedKey("a"),
        params: &[ParamDescriptor::amount("amount", "a.amount")],
    };
    static DESC_B: OpDescriptor = OpDescriptor {
        id: OpId("op_b"),
        label: LocalizedKey("b"),
        params: &[ParamDescriptor::amount("amount", "b.amount")],
    };

    struct Fake {
        desc: &'static OpDescriptor,
        amount: f32,
        helper: Option<Helper>,
    }

    impl Operation for Fake {
        fn descriptor(&self) -> &'static OpDescriptor {
            self.desc
        }
        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 {
            "c = c * amount;".into()
        }
        fn uniforms(&self) -> Vec<Uniform> {
            vec![Uniform {
                name: "amount",
                value: self.amount,
            }]
        }
        fn helpers(&self) -> &'static [Helper] {
            match self.helper {
                Some(_) => SHARED,
                None => &[],
            }
        }
    }

    static SHARED: &[Helper] = &[Helper {
        name: "luma",
        source: "fn luma(c: vec3<f32>) -> f32 { return c.g; }",
    }];

    fn fake(desc: &'static OpDescriptor, amount: f32, helper: bool) -> Box<dyn Operation> {
        Box::new(Fake {
            desc,
            amount,
            helper: helper.then_some(SHARED[0]),
        })
    }

    #[test]
    fn an_inactive_operation_contributes_nothing() {
        // The point of composing rather than branching: an op at neutral
        // must not appear in the source at all.
        let ops = vec![fake(&DESC_A, 0.0, false)];
        let shader = compose(&ops);
        assert!(
            !shader.source.contains("op_a"),
            "a neutral operation must not reach the generated shader"
        );
        assert_eq!(
            shader.uniforms.len(),
            PREAMBLE_FIELDS,
            "it must contribute no uniforms either"
        );
    }

    /// Uniform slots reserved before any operation's own: the camera matrix
    /// and as-shot white balance, plus framing. The same constant `dr-gpu`
    /// writes against, so these offsets cannot agree with each other while
    /// disagreeing with the shader.
    const PREAMBLE_FIELDS: usize = RESERVED_UNIFORM_FIELDS;

    #[test]
    fn an_active_operation_appears_once() {
        let ops = vec![fake(&DESC_A, 2.0, false)];
        let shader = compose(&ops);
        assert!(shader.source.contains("---- op_a ----"));
        assert!(shader.source.contains("u.op_a_amount"));
    }

    #[test]
    fn uniforms_are_prefixed_so_operations_cannot_collide() {
        // Both fakes declare a uniform called `amount`. Without prefixing,
        // the generated struct would have a duplicate field and fail to
        // compile — the failure mode that makes naive concatenation fragile.
        let ops = vec![fake(&DESC_A, 1.0, false), fake(&DESC_B, 2.0, false)];
        let shader = compose(&ops);
        assert!(shader.source.contains("op_a_amount: f32"));
        assert!(shader.source.contains("op_b_amount: f32"));
        assert!(shader.source.contains("c = c * u.op_a_amount;"));
        assert!(shader.source.contains("c = c * u.op_b_amount;"));
    }

    #[test]
    fn uniform_values_follow_declaration_order() {
        let ops = vec![fake(&DESC_A, 1.5, false), fake(&DESC_B, 2.5, false)];
        let shader = compose(&ops);
        assert_eq!(shader.uniforms[PREAMBLE_FIELDS], 1.5);
        assert_eq!(shader.uniforms[PREAMBLE_FIELDS + 1], 2.5);
    }

    #[test]
    fn a_shared_helper_is_emitted_once() {
        // Two operations wanting the same helper must not produce a
        // duplicate function definition.
        let ops = vec![fake(&DESC_A, 1.0, true), fake(&DESC_B, 1.0, true)];
        let shader = compose(&ops);
        assert_eq!(
            shader.source.matches("fn luma(").count(),
            1,
            "a helper requested twice must be declared once"
        );
    }

    #[test]
    fn the_uniform_block_is_16_byte_aligned() {
        // WGSL rejects a uniform struct whose size is not a multiple of 16.
        for n in 0..6 {
            let ops: Vec<Box<dyn Operation>> = (0..n)
                .map(|i| fake(if i % 2 == 0 { &DESC_A } else { &DESC_B }, 1.0, false))
                .collect();
            let shader = compose(&ops);
            assert_eq!(
                shader.uniforms.len() % 4,
                0,
                "{n} operations produced {} floats, not a multiple of 4",
                shader.uniforms.len()
            );
        }
    }

    #[test]
    fn structure_hash_ignores_values_but_tracks_the_op_set() {
        // The property the shader cache depends on: moving a slider must not
        // trigger a recompile, but enabling an operation must.
        let a1 = compose(&[fake(&DESC_A, 1.0, false)]).structure_hash;
        let a2 = compose(&[fake(&DESC_A, 9.0, false)]).structure_hash;
        assert_eq!(a1, a2, "a value change must reuse the compiled pipeline");

        let both = compose(&[fake(&DESC_A, 1.0, false), fake(&DESC_B, 1.0, false)]);
        assert_ne!(a1, both.structure_hash, "a different op-set must recompile");
    }

    #[test]
    fn structure_hash_is_order_sensitive() {
        // Operation order is data (ARCH §3.4); two orders are different
        // shaders and must not share a cache entry.
        let ab = compose(&[fake(&DESC_A, 1.0, false), fake(&DESC_B, 1.0, false)]);
        let ba = compose(&[fake(&DESC_B, 1.0, false), fake(&DESC_A, 1.0, false)]);
        assert_ne!(ab.structure_hash, ba.structure_hash);
    }

    #[test]
    fn rewriting_respects_word_boundaries() {
        // `amount` must not corrupt `amount_hi` — the bug a naive
        // string replace would introduce.
        let got = rewrite_uniform("x = amount + amount_hi;", "amount", "u.p_amount");
        assert_eq!(got, "x = u.p_amount + amount_hi;");
    }

    #[test]
    fn rewriting_leaves_comments_alone() {
        // Found in generated source: a comment reading "A factor of 0 is
        // monochrome" came out as "A u.saturation_factor of 0 is monochrome".
        // The generated source is what gets read when a shader fails to
        // compile, so mangling it works against the one time it matters.
        let got = rewrite_uniform(
            "// A factor of 0 is monochrome\nc = c * factor;",
            "factor",
            "u.op_factor",
        );
        assert_eq!(got, "// A factor of 0 is monochrome\nc = c * u.op_factor;");
    }

    #[test]
    fn rewriting_resumes_after_a_comment_ends() {
        let got = rewrite_uniform(
            "// factor here is prose\nlet x = factor;\n// factor again\n",
            "factor",
            "u.p",
        );
        assert_eq!(
            got,
            "// factor here is prose\nlet x = u.p;\n// factor again\n"
        );
    }

    #[test]
    fn rewriting_leaves_substrings_alone() {
        let got = rewrite_uniform("total_amount = 1.0;", "amount", "u.a");
        assert_eq!(got, "total_amount = 1.0;");
    }

    #[test]
    fn as_shot_white_balance_is_applied_even_with_no_operations() {
        // The bug this catches, seen on a real CR2: a Bayer sensor's green
        // photosites collect roughly twice the signal of its red and blue,
        // so an image rendered without the as-shot multipliers comes out
        // violently green. It must not depend on the white balance operation
        // being active — that one carries only the user's offset.
        let shader = compose(&[]);
        assert!(
            shader.source.contains("u.as_shot_wb"),
            "a neutral edit must still apply as-shot white balance"
        );
    }

    #[test]
    fn white_balance_is_applied_before_the_operations() {
        // Exposure and the tonal controls act on white-balanced values; if
        // the multiply came afterwards, every operation would be reasoning
        // about a green-cast image.
        let ops = vec![fake(&DESC_A, 2.0, false)];
        let source = compose(&ops).source;
        let wb = source.find("u.as_shot_wb").expect("wb applied");
        let op = source.find("---- op_a ----").expect("op present");
        assert!(wb < op, "as-shot white balance must precede the operations");
    }

    #[test]
    fn the_camera_matrix_is_applied_after_the_operations() {
        // Adjustments are meaningful in sensor-native space, where highlight
        // headroom still exists; converting first would clip it away.
        let ops = vec![fake(&DESC_A, 2.0, false)];
        let source = compose(&ops).source;
        let op = source.find("---- op_a ----").expect("op present");
        let matrix = source.find("u.cam_to_srgb_0").expect("matrix applied");
        assert!(op < matrix, "the camera matrix must come after operations");
    }

    /// Compose with neutral framing into a chosen output space.
    fn compose_to(ops: &[Box<dyn Operation>], output: ColourSpace) -> ComposedShader {
        compose_with_framing(ops, &Framing::new(), output)
    }

    #[test]
    fn an_srgb_render_is_byte_for_byte_what_it_was_before_output_spaces_existed() {
        // The display path is the shader compiled on nearly every frame, and
        // it must not pick up an identity matrix multiply for the sake of
        // generality. Asserted against the source rather than against timing,
        // which would not fail reliably.
        let ops = vec![fake(&DESC_A, 1.0, false)];
        let srgb = compose_to(&ops, ColourSpace::Srgb).source;
        assert!(
            !srgb.contains("Linear sRGB -> linear sRGB"),
            "sRGB in, sRGB out must emit no conversion:\n{srgb}"
        );
        assert_eq!(srgb, compose(&ops).source);
    }

    #[test]
    fn a_wide_output_space_converts_after_the_camera_matrix_and_before_the_clip() {
        // The whole point of the ordering. The camera matrix lands the colour
        // in linear sRGB, the primaries conversion carries it into the wider
        // space, and only then is it clipped — clipping first would discard
        // exactly the colours the wider space was chosen to keep.
        let source = compose_to(&[fake(&DESC_A, 1.0, false)], ColourSpace::DisplayP3).source;
        let camera = source.find("u.cam_to_srgb_0").expect("camera matrix");
        let convert = source
            .find("Linear sRGB -> linear Display P3")
            .expect("primaries conversion");
        let clip = source.find("c = clamp(c,").expect("clip");
        assert!(camera < convert, "the camera matrix must come first");
        assert!(convert < clip, "the clip must come after the conversion");
    }

    #[test]
    fn the_generated_matrix_is_the_one_the_profile_writer_will_use() {
        // The shader encodes the pixels and `dr-export` describes them, from
        // the same table in `dr-types`. If the composer ever grew its own copy
        // of these numbers the file would be labelled with primaries it does
        // not contain, which is the failure the whole feature exists to avoid.
        let m = ColourSpace::DisplayP3.from_linear_srgb();
        let first_row = format!("dot(vec3<f32>({:.6}, {:.6}, {:.6}), c)", m[0], m[1], m[2]);
        let source = compose_to(&[], ColourSpace::DisplayP3).source;
        assert!(
            source.contains(&first_row),
            "expected {first_row} in:\n{source}"
        );
    }

    #[test]
    fn every_output_space_encodes_with_its_own_transfer_function() {
        // Adobe RGB's pure 2.199 gamma and ProPhoto's 1.8-with-a-toe are not
        // the sRGB curve, and a file encoded with the wrong one is wrong in a
        // way no amount of correct primaries repairs.
        let marks = [
            (ColourSpace::Srgb, "1.0 / 2.4"),
            (ColourSpace::DisplayP3, "1.0 / 2.4"),
            (ColourSpace::AdobeRgb, "1.0 / 2.19921875"),
            (ColourSpace::ProPhoto, "1.0 / 1.8"),
        ];
        for (space, mark) in marks {
            let source = compose_to(&[], space).source;
            assert!(
                source.contains(mark),
                "{space:?} should encode with {mark}:\n{source}"
            );
        }
    }

    #[test]
    fn the_output_space_changes_the_structure_hash() {
        // The pipeline cache is keyed on this hash. Two spaces sharing one
        // would have the second rendered with the first's compiled shader —
        // an export that came out sRGB and claimed to be Display P3.
        let mut seen: Vec<u64> = Vec::new();
        for space in ColourSpace::ALL {
            let h = compose_to(&[fake(&DESC_A, 1.0, false)], space).structure_hash;
            assert!(!seen.contains(&h), "{space:?} collides with another space");
            seen.push(h);
        }
    }

    #[test]
    fn generated_source_carries_a_do_not_edit_banner() {
        // Someone will eventually find this in a debugger and try to fix it
        // in place.
        let shader = compose(&[fake(&DESC_A, 1.0, false)]);
        assert!(shader.source.starts_with("// GENERATED"));
    }
}