onedrop_codegen/

prelude.rs

//! Standard WGSL prelude injected ahead of every translated user shader.
//!
//! The prelude exposes the MilkDrop 2 uniform set (audio, time, aspect,
//! `q1..q32`, etc.) under a single `uniforms` binding so warp and comp
//! HLSL shaders can be translated to WGSL with a known reference target.
//!
//! Two artefacts must stay in lockstep:
//! - [`ShaderUniforms`] — the Rust mirror, repr-C, Pod/Zeroable. The
//!   renderer fills it each frame from `RenderState`/`MilkContext` and
//!   uploads it via `wgpu::Queue::write_buffer`.
//! - [`SHADER_UNIFORMS_WGSL`] — the WGSL `struct` declaration plus the
//!   `@group(0) @binding(0) var<uniform> uniforms: ShaderUniforms;` line.
//!   Concatenated to the front of every translated user shader before
//!   naga validates and the renderer compiles it.
//!
//! Layout discipline (matches WGSL/std140-ish uniform alignment):
//! - All scalar `f32` fields packed at the top in groups of four (each
//!   group is one 16-byte slot).
//! - All `vec4<f32>` fields after, contiguous (each is exactly 16 bytes,
//!   so no further padding is needed).
//! - Trailing `array<vec4<f32>, N>` for `q1..q32` — eight vec4 slots = 32
//!   floats.
//!
//! The size in bytes is asserted by the
//! `tests::layout_is_byte_for_byte_compatible` test in this module
//! against a naga-parsed WGSL module: any drift between the Rust struct
//! and the WGSL declaration trips the test.

use bytemuck::{Pod, Zeroable};

/// Number of `q*` channels exposed to user shaders. MilkDrop 2 ships
/// `q1..q32`; the engine tracks 64 internally but only the first 32 cross
/// into shader-land.
pub const SHADER_Q_CHANNELS: usize = 32;

/// Rust mirror of the WGSL `ShaderUniforms` struct.
///
/// Field order, alignment, and padding must exactly match
/// [`SHADER_UNIFORMS_WGSL`]; the `tests` module in this file enforces
/// this with a byte-for-byte naga round-trip.
#[repr(C)]
#[derive(Clone, Copy, Debug, Pod, Zeroable)]
pub struct ShaderUniforms {
    // ---- scalar slot 0 (16 B) ----
    /// Wall-clock time in seconds since preset load.
    pub time: f32,
    /// Smoothed FPS estimate.
    pub fps: f32,
    /// Frame counter (cast from u32).
    pub frame: f32,
    /// Position within the preset's display window, `[0, 1]`. 0 at load,
    /// 1 just before the next hard cut.
    pub progress: f32,

    // ---- scalar slot 1 (16 B) ----
    /// Bass band level (raw, ~0 baseline, can spike >1 on beats).
    pub bass: f32,
    /// Mid band level.
    pub mid: f32,
    /// Treble band level.
    pub treb: f32,
    /// `(bass + mid + treb) / 3`.
    pub vol: f32,

    // ---- scalar slot 2 (16 B) ----
    /// Smoothed (attenuated) bass level — exponential moving average.
    pub bass_att: f32,
    pub mid_att: f32,
    pub treb_att: f32,
    pub vol_att: f32,

    // ---- scalar slot 3 (16 B): blur level 1 / 2 ----
    /// `bN_min` / `bN_max` are remap bounds — the comp shader applies
    /// `lerp(min, max, sample)` to recover usable contrast from the
    /// downsampled blur pyramid.
    pub blur1_min: f32,
    pub blur1_max: f32,
    pub blur2_min: f32,
    pub blur2_max: f32,

    // ---- scalar slot 4 (16 B): blur level 3 + display gamma ----
    pub blur3_min: f32,
    pub blur3_max: f32,
    /// MilkDrop's `f_gamma_adj`. Display-only multiplier (not a true gamma);
    /// applied in the comp pass after any user shader work.
    pub gamma_adj: f32,
    /// Reserved for a future scalar; keeps the slot count a multiple of 4.
    pub _reserved1: f32,

    // ---- scalar slot 5 (16 B): echo + stereo ----
    /// MilkDrop's `f_video_echo_zoom`. Comp pass samples the previous
    /// frame's display at `(uv-0.5)*echo_zoom + 0.5` and blends it with
    /// the current warp output. `0.0` disables the echo branch entirely
    /// (treated as `echo_alpha = 0`).
    pub echo_zoom: f32,
    /// MilkDrop's `f_video_echo_alpha`. Mix weight: `0.0` = no echo, full
    /// current frame; `1.0` = pure echo from the previous frame.
    pub echo_alpha: f32,
    /// MilkDrop's `n_video_echo_orientation`. 0 = no flip, 1 = flip X,
    /// 2 = flip Y, 3 = flip both. Stored as `f32` to avoid mixed
    /// scalar types in the uniform layout; the comp shader casts to
    /// an integer at use site.
    pub echo_orient: f32,
    /// MilkDrop's `b_red_blue_stereo`. `1.0` = anaglyph mode active,
    /// `0.0` = off. Wired through to `comp.wgsl`.
    pub red_blue_stereo: f32,

    // ---- vec4 slot 0: aspect (16 B) ----
    /// `(aspect_x, aspect_y, 1/aspect_x, 1/aspect_y)`. MD2 historically
    /// supplies all four to spare the shader an extra division.
    pub aspect: [f32; 4],

    // ---- vec4 slot 1: texsize_main (16 B) ----
    /// `(width, height, 1/width, 1/height)` of the warp/comp render
    /// target — the texture user shaders sample as `sampler_main`.
    pub texsize: [f32; 4],

    // ---- vec4 slot 2: per-preset random vector (16 B) ----
    /// Four `[0, 1)` floats sampled once at preset load. MD2's `rand_preset`.
    pub rand_preset: [f32; 4],

    // ---- vec4 slot 3: per-frame random vector (16 B) ----
    /// Four `[0, 1)` floats resampled every frame.
    pub rand_frame: [f32; 4],

    // ---- vec4 slots 4-7: roaming animations (4 × 16 B) ----
    /// Pre-computed `cos` of the four MD2 "slow roam" oscillators.
    pub slow_roam_cos: [f32; 4],
    /// Pre-computed `sin` of the same.
    pub slow_roam_sin: [f32; 4],
    /// Faster-cycling roam oscillators.
    pub roam_cos: [f32; 4],
    pub roam_sin: [f32; 4],

    // ---- vec4 slot 8: hue_shader tint (16 B) ----
    /// MD2's per-frame RGB tint (`.xyz`). Many comp shaders multiply
    /// `ret * hue_shader` for a slow rainbow rotation across frames;
    /// without this populated, the multiplication degrades to identity
    /// (`vec3(1, 1, 1)`) and the rotation disappears. The engine
    /// computes a time-driven tri-phase oscillator each frame; the
    /// `.w` lane is reserved for a future per-corner mix factor.
    pub hue_shader: [f32; 4],

    // ---- vec4 array: q channels (8 × 16 B = 128 B) ----
    /// `q[0]` packs `q1..q4`, `q[1]` packs `q5..q8`, etc. Filled from
    /// `MilkEngine::q_snapshot()`.
    pub q: [[f32; 4]; 8],
}

impl ShaderUniforms {
    /// Size of the uniform buffer in bytes. Useful when allocating GPU
    /// buffers without taking `size_of::<Self>()` at the call site.
    pub const SIZE: usize = std::mem::size_of::<Self>();

    /// Pack `q1..q32` into the eight-vec4 layout the WGSL struct expects.
    pub fn set_q_channels(&mut self, q: &[f32; SHADER_Q_CHANNELS]) {
        for (i, slot) in self.q.iter_mut().enumerate() {
            slot[0] = q[i * 4];
            slot[1] = q[i * 4 + 1];
            slot[2] = q[i * 4 + 2];
            slot[3] = q[i * 4 + 3];
        }
    }
}

impl Default for ShaderUniforms {
    fn default() -> Self {
        Self {
            time: 0.0,
            fps: 60.0,
            frame: 0.0,
            progress: 0.0,
            bass: 0.0,
            mid: 0.0,
            treb: 0.0,
            vol: 0.0,
            bass_att: 0.0,
            mid_att: 0.0,
            treb_att: 0.0,
            vol_att: 0.0,
            blur1_min: 0.0,
            blur1_max: 1.0,
            blur2_min: 0.0,
            blur2_max: 1.0,
            blur3_min: 0.0,
            blur3_max: 1.0,
            gamma_adj: 2.0,
            _reserved1: 0.0,
            echo_zoom: 1.0,
            echo_alpha: 0.0,
            echo_orient: 0.0,
            red_blue_stereo: 0.0,
            aspect: [1.0, 1.0, 1.0, 1.0],
            texsize: [1.0, 1.0, 1.0, 1.0],
            rand_preset: [0.0; 4],
            rand_frame: [0.0; 4],
            slow_roam_cos: [0.0; 4],
            slow_roam_sin: [0.0; 4],
            roam_cos: [0.0; 4],
            roam_sin: [0.0; 4],
            hue_shader: [1.0, 1.0, 1.0, 1.0],
            q: [[0.0; 4]; 8],
        }
    }
}

/// WGSL `struct` declaration mirroring [`ShaderUniforms`] plus the standard
/// `@group(0) @binding(0)` uniform binding.
///
/// Field order/types are identical; this string is concatenated to the front
/// of every translated user shader.
pub const SHADER_UNIFORMS_WGSL: &str = r#"
struct ShaderUniforms {
    // scalar slot 0
    time: f32,
    fps: f32,
    frame: f32,
    progress: f32,
    // scalar slot 1
    bass: f32,
    mid: f32,
    treb: f32,
    vol: f32,
    // scalar slot 2
    bass_att: f32,
    mid_att: f32,
    treb_att: f32,
    vol_att: f32,
    // scalar slot 3
    blur1_min: f32,
    blur1_max: f32,
    blur2_min: f32,
    blur2_max: f32,
    // scalar slot 4
    blur3_min: f32,
    blur3_max: f32,
    gamma_adj: f32,
    _reserved1: f32,
    // scalar slot 5 — echo + anaglyph stereo
    echo_zoom: f32,
    echo_alpha: f32,
    echo_orient: f32,
    red_blue_stereo: f32,
    // vec4 slots
    aspect: vec4<f32>,
    texsize: vec4<f32>,
    rand_preset: vec4<f32>,
    rand_frame: vec4<f32>,
    slow_roam_cos: vec4<f32>,
    slow_roam_sin: vec4<f32>,
    roam_cos: vec4<f32>,
    roam_sin: vec4<f32>,
    // Per-frame RGB tint (`.xyz`); `.w` reserved.
    hue_shader: vec4<f32>,
    // q1..q32 packed as 8 × vec4
    q: array<vec4<f32>, 8>,
}

@group(0) @binding(0)
var<uniform> uniforms: ShaderUniforms;
"#;

/// Build a full prelude string by appending the user's shader body to the
/// standard `ShaderUniforms` declaration. `body` is expected to already be
/// translated to WGSL.
pub fn with_prelude(body: &str) -> String {
    let mut out = String::with_capacity(SHADER_UNIFORMS_WGSL.len() + body.len() + 1);
    out.push_str(SHADER_UNIFORMS_WGSL);
    out.push('\n');
    out.push_str(body);
    out
}

#[cfg(test)]
mod tests {
    use super::*;

    /// 6 scalar slots × 16 B + 9 vec4 fields × 16 B + 8 × vec4 q-array × 16 B
    /// = (6 + 9 + 8) × 16 = 368 bytes. Slot 5 holds the echo +
    /// red_blue_stereo block; vec4 slot 8 is `hue_shader`.
    const EXPECTED_SIZE: usize = (6 + 9 + 8) * 16;

    #[test]
    fn rust_struct_matches_expected_layout() {
        assert_eq!(std::mem::size_of::<ShaderUniforms>(), EXPECTED_SIZE);
    }

    #[test]
    fn rust_struct_size_matches_wgsl_layout() {
        // Compile the prelude through naga, look up `ShaderUniforms` in the
        // type arena, and confirm naga's computed byte size matches the
        // Rust mirror's `size_of`. Any drift between the two declarations
        // (a missing field, a swapped vec4/scalar, an alignment surprise)
        // makes the two sides disagree and trips this test.
        let trivial_shader = format!(
            "{SHADER_UNIFORMS_WGSL}\n@vertex fn vs() -> @builtin(position) vec4<f32> {{ return vec4<f32>(0.0); }}"
        );
        let module = naga::front::wgsl::parse_str(&trivial_shader).unwrap();
        let mut info = naga::valid::Validator::new(
            naga::valid::ValidationFlags::all(),
            naga::valid::Capabilities::all(),
        );
        let _ = info.validate(&module).expect("module must validate");

        let (_, ty) = module
            .types
            .iter()
            .find(|(_, t)| t.name.as_deref() == Some("ShaderUniforms"))
            .expect("ShaderUniforms struct missing from prelude");
        // naga exposes the layout-computed size on the inner Struct variant.
        let wgsl_size = match &ty.inner {
            naga::TypeInner::Struct { span, .. } => *span as usize,
            _ => panic!("ShaderUniforms is not a struct in WGSL"),
        };
        assert_eq!(
            wgsl_size,
            std::mem::size_of::<ShaderUniforms>(),
            "Rust struct ({} B) and WGSL struct ({} B) disagree on size — \
             field order or padding has drifted",
            std::mem::size_of::<ShaderUniforms>(),
            wgsl_size,
        );
    }

    #[test]
    fn wgsl_declaration_parses() {
        // The prelude alone is just a struct + uniform binding — naga
        // accepts it as a valid module fragment when wrapped in a no-op
        // shader.
        let trivial_shader = format!(
            "{SHADER_UNIFORMS_WGSL}\n@vertex fn vs() -> @builtin(position) vec4<f32> {{ return vec4<f32>(0.0); }}"
        );
        let module = naga::front::wgsl::parse_str(&trivial_shader)
            .expect("WGSL prelude must parse with naga");
        // Sanity: the ShaderUniforms struct should be present in the
        // type arena.
        let has_struct = module.types.iter().any(|(_, ty)| {
            matches!(&ty.inner, naga::TypeInner::Struct { .. })
                && ty.name.as_deref() == Some("ShaderUniforms")
        });
        assert!(has_struct, "naga did not see ShaderUniforms in the prelude");
    }

    #[test]
    fn with_prelude_concatenates_in_order() {
        let body = "// user body";
        let combined = with_prelude(body);
        let prelude_end = combined.find(body).expect("body not found");
        assert!(
            combined[..prelude_end].contains("struct ShaderUniforms"),
            "prelude must come before the user body"
        );
    }

    #[test]
    fn set_q_channels_packs_correctly() {
        let mut u = ShaderUniforms::default();
        let mut q = [0.0f32; SHADER_Q_CHANNELS];
        for (i, slot) in q.iter_mut().enumerate() {
            *slot = i as f32 + 1.0; // q1=1, q2=2, ...
        }
        u.set_q_channels(&q);
        // q1..q4 land in u.q[0]
        assert_eq!(u.q[0], [1.0, 2.0, 3.0, 4.0]);
        // q29..q32 land in u.q[7]
        assert_eq!(u.q[7], [29.0, 30.0, 31.0, 32.0]);
    }

    #[test]
    fn default_is_pod_safe_zero_initialisable() {
        // Pod requires "all bit patterns valid"; the default is a sane
        // starting point but we also want the all-zeroes pattern (used by
        // `Zeroable`) to be a valid uniform — exercised here.
        let zeroed: ShaderUniforms = bytemuck::Zeroable::zeroed();
        assert_eq!(zeroed.time, 0.0);
        assert_eq!(zeroed.q, [[0.0; 4]; 8]);
    }
}