onedrop_renderer/

waveform.rs

//! Waveform pass — MD2 modes 0..7 wired into the renderer.
//!
//! `WaveformRenderer` owns four pipelines (alpha vs additive × lines vs
//! dots) that share a uniform + storage layout. Per frame the host:
//!
//! 1. Smooths the audio samples (CPU IIR) and uploads them to
//!    [`WaveformRenderer::update_wave_samples`].
//! 2. Updates the [`WaveUniforms`] payload from `WaveParams`.
//! 3. Calls [`WaveformRenderer::render`] into the warp output texture so
//!    the wave participates in the feedback loop (matches MD2's order).
//!
//! Mode 6 (`double_line`) is drawn in two passes: the second call sets
//! `flip_y = 1` so the wave mirrors around `wave_y`. Modes 0..5 and 7
//! ignore `flip_y`.

use bytemuck::{Pod, Zeroable};
use wgpu::util::DeviceExt;

use crate::config::WaveParams;

/// MilkDrop 2 waveform mode (`nWaveMode`, 0..7).
///
/// Numbering matches the MD2 preset format: presets read
/// `nWaveMode = 4` as `DerivativeLine`, etc.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[repr(u32)]
pub enum WaveformMode {
    Circle = 0,
    XyOscopeOctagonal = 1,
    XyOscopeQuadrilateral = 2,
    Blob = 3,
    DerivativeLine = 4,
    ExplosiveHash = 5,
    DoubleLine = 6,
    DoubleHorizontal = 7,
}

impl WaveformMode {
    /// Convert a raw `nWaveMode` int to a mode, clamping anything outside
    /// `[0, 7]` to `Circle`. MD2 silently does the same.
    pub fn from_i32(v: i32) -> Self {
        match v {
            0 => Self::Circle,
            1 => Self::XyOscopeOctagonal,
            2 => Self::XyOscopeQuadrilateral,
            3 => Self::Blob,
            4 => Self::DerivativeLine,
            5 => Self::ExplosiveHash,
            6 => Self::DoubleLine,
            7 => Self::DoubleHorizontal,
            _ => Self::Circle,
        }
    }

    /// True when the wave is drawn as a closed loop (samples wrap).
    pub fn is_closed(self) -> bool {
        matches!(self, Self::Circle | Self::Blob)
    }

    /// True when MD2 draws the mode in two passes (mirrored). Currently
    /// only `DoubleLine`.
    pub fn is_double_pass(self) -> bool {
        matches!(self, Self::DoubleLine)
    }
}

/// Standard MD2 waveform sample length — 512 samples per channel before
/// downmix. The shader treats `wave_samples` as mono; the host is
/// responsible for the L/R mix.
pub const NUM_WAVE_SAMPLES: usize = 512;

/// Default IIR smoothing factor when `wave_smoothing` is 0. MD2 uses a
/// 5-tap moving average; a single-pole IIR with α≈0.5 is a cheap visual
/// match.
const DEFAULT_SMOOTHING_ALPHA: f32 = 0.0;

/// Per-frame uniforms consumed by the waveform shader. Field order
/// MUST match `Uniforms` in `waveform_advanced.wgsl`.
#[repr(C)]
#[derive(Debug, Clone, Copy, Pod, Zeroable)]
pub struct WaveUniforms {
    pub resolution: [f32; 2],
    pub time: f32,
    pub mode: u32,

    pub wave_x: f32,
    pub wave_y: f32,
    pub wave_scale: f32,
    pub wave_alpha: f32,

    pub wave_param: f32,
    pub aspect: f32,
    pub thickness: f32,
    pub smoothing: f32,

    pub wave_color: [f32; 4],

    pub num_samples: u32,
    pub is_thick: u32,
    pub is_dots: u32,
    /// Index offset into the (now 2N-wide) sample buffer. `0` picks the
    /// left-channel half (and the mono-fallback when no right channel
    /// is uploaded); `NUM_WAVE_SAMPLES` picks the right-channel half.
    /// The split-channel path picks left vs. right by changing this
    /// offset between two consecutive draws at different `wave_y`
    /// offsets when the top-vs-bottom L/R split is enabled.
    pub sample_offset: u32,
}

impl Default for WaveUniforms {
    fn default() -> Self {
        Self {
            resolution: [1.0, 1.0],
            time: 0.0,
            mode: 0,
            wave_x: 0.5,
            wave_y: 0.5,
            wave_scale: 0.5,
            wave_alpha: 1.0,
            wave_param: 0.0,
            aspect: 1.0,
            thickness: 0.004,
            smoothing: 0.0,
            wave_color: [1.0, 1.0, 1.0, 1.0],
            num_samples: NUM_WAVE_SAMPLES as u32,
            is_thick: 0,
            is_dots: 0,
            sample_offset: 0,
        }
    }
}

/// Legacy storage element kept for backwards-compatibility with older
/// tests that referenced `WavePoint`. The new pipeline streams a flat
/// `array<f32>` and does NOT consume `WavePoint`; the type is retained
/// only so external code that still uses it keeps compiling.
#[repr(C)]
#[derive(Debug, Clone, Copy, Pod, Zeroable)]
pub struct WavePoint {
    pub position: [f32; 2],
    pub value: f32,
    pub _padding: f32,
}

/// GPU waveform pass — owns its pipelines, uniforms, and a sample
/// storage buffer of fixed length `NUM_WAVE_SAMPLES`.
pub struct WaveformRenderer {
    /// 4 pipelines: (alpha lines, additive lines, alpha dots, additive
    /// dots). Indexed by `pipeline_index(is_dots, is_additive)`.
    pipelines: [wgpu::RenderPipeline; 4],

    bind_group: wgpu::BindGroup,
    uniform_buffer: wgpu::Buffer,
    sample_buffer: wgpu::Buffer,

    /// Last uniforms uploaded, kept for `flip_y` double-pass dispatch.
    last_uniforms: WaveUniforms,
}

#[inline]
fn pipeline_index(is_dots: bool, is_additive: bool) -> usize {
    (is_additive as usize) | ((is_dots as usize) << 1)
}

impl WaveformRenderer {
    /// Create a new renderer targeting `format`. `format` must match the
    /// texture the host will dispatch into (typically `render_texture`).
    pub fn new(device: &wgpu::Device, format: wgpu::TextureFormat) -> Self {
        let shader = crate::pipeline_helpers::load_wgsl(
            device,
            "Waveform Shader",
            include_str!("../shaders/waveform_advanced.wgsl"),
        );

        let uniform_buffer = device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("Waveform Uniforms"),
            size: std::mem::size_of::<WaveUniforms>() as u64,
            usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });

        // Initial sample buffer is zero-filled silence. Allocate 2 ×
        // `NUM_WAVE_SAMPLES` to hold both channels back-to-back:
        // `[0 .. N)` = left, `[N .. 2N)` = right. The mono path
        // duplicates the left buffer into the right half so a
        // shader that reads either offset sees coherent data.
        let zeros = [0.0f32; NUM_WAVE_SAMPLES * 2];
        let sample_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
            label: Some("Waveform Samples"),
            contents: bytemuck::cast_slice(&zeros),
            usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
        });

        let bgl = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
            label: Some("Waveform BGL"),
            entries: &[
                wgpu::BindGroupLayoutEntry {
                    binding: 0,
                    visibility: wgpu::ShaderStages::VERTEX_FRAGMENT,
                    ty: wgpu::BindingType::Buffer {
                        ty: wgpu::BufferBindingType::Uniform,
                        has_dynamic_offset: false,
                        min_binding_size: None,
                    },
                    count: None,
                },
                wgpu::BindGroupLayoutEntry {
                    binding: 1,
                    visibility: wgpu::ShaderStages::VERTEX,
                    ty: wgpu::BindingType::Buffer {
                        ty: wgpu::BufferBindingType::Storage { read_only: true },
                        has_dynamic_offset: false,
                        min_binding_size: None,
                    },
                    count: None,
                },
            ],
        });

        let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some("Waveform BG"),
            layout: &bgl,
            entries: &[
                wgpu::BindGroupEntry {
                    binding: 0,
                    resource: uniform_buffer.as_entire_binding(),
                },
                wgpu::BindGroupEntry {
                    binding: 1,
                    resource: sample_buffer.as_entire_binding(),
                },
            ],
        });

        let layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
            label: Some("Waveform Layout"),
            bind_group_layouts: &[Some(&bgl)],
            immediate_size: 0,
        });

        let make_pipeline = |label: &str, additive: bool| -> wgpu::RenderPipeline {
            let blend = crate::pipeline_helpers::blend_state_for(additive);
            device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
                label: Some(label),
                layout: Some(&layout),
                vertex: wgpu::VertexState {
                    module: &shader,
                    entry_point: Some("vs_main"),
                    buffers: &[],
                    compilation_options: Default::default(),
                },
                fragment: Some(wgpu::FragmentState {
                    module: &shader,
                    entry_point: Some("fs_main"),
                    targets: &[Some(wgpu::ColorTargetState {
                        format,
                        blend: Some(blend),
                        write_mask: wgpu::ColorWrites::ALL,
                    })],
                    compilation_options: Default::default(),
                }),
                primitive: wgpu::PrimitiveState {
                    topology: wgpu::PrimitiveTopology::TriangleList,
                    ..Default::default()
                },
                depth_stencil: None,
                multisample: wgpu::MultisampleState::default(),
                multiview_mask: None,
                cache: None,
            })
        };

        let pipelines = [
            make_pipeline("Wave Lines Alpha", false),
            make_pipeline("Wave Lines Additive", true),
            make_pipeline("Wave Dots Alpha", false),
            make_pipeline("Wave Dots Additive", true),
        ];

        Self {
            pipelines,
            bind_group,
            uniform_buffer,
            sample_buffer,
            last_uniforms: WaveUniforms::default(),
        }
    }

    /// Push a fresh sample buffer (mono, length capped at
    /// `NUM_WAVE_SAMPLES`; shorter buffers are zero-padded). Apply CPU
    /// IIR smoothing first with [`apply_smoothing`] if you want
    /// `wave_smoothing` to actually attenuate spikes.
    ///
    /// Mono path: the same data is written into both halves of the
    /// 2N-wide buffer so a shader pass at either `sample_offset` sees
    /// coherent samples. Stereo callers use
    /// [`Self::update_wave_samples_lr`] instead.
    pub fn update_wave_samples(&self, queue: &wgpu::Queue, samples: &[f32]) {
        let mut buf = [0.0f32; NUM_WAVE_SAMPLES * 2];
        let n = samples.len().min(NUM_WAVE_SAMPLES);
        buf[..n].copy_from_slice(&samples[..n]);
        buf[NUM_WAVE_SAMPLES..NUM_WAVE_SAMPLES + n].copy_from_slice(&samples[..n]);
        queue.write_buffer(&self.sample_buffer, 0, bytemuck::cast_slice(&buf));
    }

    /// Push split L/R sample buffers. Left lands at offset `0` (the
    /// default `sample_offset` of `0` reads it); right lands at offset
    /// `NUM_WAVE_SAMPLES`. Used by the top-vs-bottom split path —
    /// `RenderChain::record_passes` dispatches two waveform passes with
    /// distinct `sample_offset` + `wave_y` so the listener sees L on
    /// the upper half and R on the lower half.
    pub fn update_wave_samples_lr(&self, queue: &wgpu::Queue, left: &[f32], right: &[f32]) {
        let mut buf = [0.0f32; NUM_WAVE_SAMPLES * 2];
        let nl = left.len().min(NUM_WAVE_SAMPLES);
        let nr = right.len().min(NUM_WAVE_SAMPLES);
        buf[..nl].copy_from_slice(&left[..nl]);
        buf[NUM_WAVE_SAMPLES..NUM_WAVE_SAMPLES + nr].copy_from_slice(&right[..nr]);
        queue.write_buffer(&self.sample_buffer, 0, bytemuck::cast_slice(&buf));
    }

    /// Push fresh uniforms. Keeps a CPU-side copy for double-pass mode 6.
    pub fn update_uniforms(&mut self, queue: &wgpu::Queue, u: &WaveUniforms) {
        self.last_uniforms = *u;
        queue.write_buffer(&self.uniform_buffer, 0, bytemuck::bytes_of(u));
    }

    /// Render the waveform into `view` using the previously uploaded
    /// uniforms + samples. Loads existing contents (no clear) — meant
    /// to draw on top of the warp pass output.
    pub fn render(&self, encoder: &mut wgpu::CommandEncoder, view: &wgpu::TextureView) {
        let mode = WaveformMode::from_i32(self.last_uniforms.mode as i32);
        let is_dots = self.last_uniforms.is_dots != 0 || mode == WaveformMode::ExplosiveHash;
        // Additive blending follows `b_additive_waves`. The
        // uniform doesn't carry it directly — the host sets blend by
        // calling `set_additive` before `render`. Default to alpha.
        let is_additive = false; // overridden by `render_with_blend`
        self.render_with_blend(encoder, view, is_dots, is_additive);
    }

    /// Same as [`Self::render`] but with explicit blend selection. Used
    /// by the host to honour `b_additive_waves` without round-tripping a
    /// flag through the uniforms.
    pub fn render_with_blend(
        &self,
        encoder: &mut wgpu::CommandEncoder,
        view: &wgpu::TextureView,
        is_dots: bool,
        is_additive: bool,
    ) {
        let mode = WaveformMode::from_i32(self.last_uniforms.mode as i32);
        let num_samples = self.last_uniforms.num_samples;
        if num_samples == 0 {
            return;
        }
        // For closed shapes we draw the same number of segments as
        // samples (wrap), open shapes use samples-1.
        let segments = if mode.is_closed() {
            num_samples
        } else {
            num_samples.saturating_sub(1)
        };
        let vertex_count = segments * 6;

        let pidx = pipeline_index(is_dots, is_additive);
        let pipeline = &self.pipelines[pidx];

        let mut rp = encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
            label: Some("Waveform Pass"),
            color_attachments: &[Some(wgpu::RenderPassColorAttachment {
                view,
                depth_slice: None,
                resolve_target: None,
                ops: wgpu::Operations {
                    load: wgpu::LoadOp::Load,
                    store: wgpu::StoreOp::Store,
                },
            })],
            depth_stencil_attachment: None,
            timestamp_writes: None,
            occlusion_query_set: None,
            multiview_mask: None,
        });
        rp.set_pipeline(pipeline);
        rp.set_bind_group(0, &self.bind_group, &[]);
        rp.draw(0..vertex_count, 0..1);
    }
}

/// Apply a single-pole IIR low-pass to a buffer of waveform samples,
/// matching MD2's `f_wave_smoothing`. `smoothing` ∈ `[0, 1]`:
/// 0 = no smoothing, 0.9 ≈ heavy attenuation. Mutates in place.
///
/// The pole α is mapped from `smoothing` so that the perceptual range
/// matches MilkDrop's slider: a `smoothing` of 0.75 → α ≈ 0.625, which
/// is the MD2 sweet spot for visible spike rejection without making the
/// waveform look like a flat line.
pub fn apply_smoothing(samples: &mut [f32], smoothing: f32) {
    let s = smoothing.clamp(0.0, 1.0);
    if s <= DEFAULT_SMOOTHING_ALPHA {
        return;
    }
    // Map smoothing slider to IIR pole. f_wave_smoothing=0.75 → α≈0.625.
    let alpha = (s * 0.833).clamp(0.0, 0.95);
    let mut prev = samples.first().copied().unwrap_or(0.0);
    for s in samples.iter_mut() {
        let next = prev * alpha + *s * (1.0 - alpha);
        *s = next;
        prev = next;
    }
}

/// Build the `WaveUniforms` payload from a `WaveParams` snapshot, the
/// current frame's instantaneous volume (used for
/// `b_mod_wave_alpha_by_volume`), and the current frame resolution +
/// aspect.
pub fn build_uniforms(
    params: WaveParams,
    instantaneous_vol: f32,
    width: u32,
    height: u32,
    time: f32,
) -> WaveUniforms {
    let aspect = width as f32 / height.max(1) as f32;
    let alpha_mod = if params.mod_alpha_by_volume {
        let v = instantaneous_vol.clamp(0.0, 1.0);
        let span = (params.mod_alpha_end - params.mod_alpha_start)
            .abs()
            .max(1e-3);
        ((v - params.mod_alpha_start) / span).clamp(0.0, 1.0)
    } else {
        1.0
    };
    let mut color = [params.r, params.g, params.b, params.a * alpha_mod];
    if params.maximize_color {
        // MD2 b_maximize_wave_color — push the colour to full saturation
        // by rescaling so max(rgb) = 1.
        let m = color[0].max(color[1]).max(color[2]).max(1e-3);
        color[0] /= m;
        color[1] /= m;
        color[2] /= m;
    }
    let thickness = if params.thick { 0.0075 } else { 0.0035 };
    WaveUniforms {
        resolution: [width as f32, height as f32],
        time,
        mode: WaveformMode::from_i32(params.mode) as u32,
        wave_x: params.x,
        wave_y: params.y,
        wave_scale: params.scale,
        wave_alpha: params.a,
        wave_param: params.param,
        aspect,
        thickness,
        smoothing: params.smoothing,
        wave_color: color,
        num_samples: NUM_WAVE_SAMPLES as u32,
        is_thick: params.thick as u32,
        is_dots: params.dots as u32,
        sample_offset: 0,
    }
}

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

    #[test]
    fn waveform_mode_from_i32_known() {
        assert_eq!(WaveformMode::from_i32(0), WaveformMode::Circle);
        assert_eq!(WaveformMode::from_i32(1), WaveformMode::XyOscopeOctagonal);
        assert_eq!(WaveformMode::from_i32(7), WaveformMode::DoubleHorizontal);
    }

    #[test]
    fn waveform_mode_from_i32_clamps_out_of_range() {
        assert_eq!(WaveformMode::from_i32(-1), WaveformMode::Circle);
        assert_eq!(WaveformMode::from_i32(99), WaveformMode::Circle);
        assert_eq!(WaveformMode::from_i32(8), WaveformMode::Circle);
    }

    #[test]
    fn waveform_mode_is_closed_for_circle_and_blob() {
        assert!(WaveformMode::Circle.is_closed());
        assert!(WaveformMode::Blob.is_closed());
        assert!(!WaveformMode::DoubleLine.is_closed());
        assert!(!WaveformMode::DerivativeLine.is_closed());
    }

    #[test]
    fn smoothing_zero_is_passthrough() {
        let mut s = vec![1.0_f32, -1.0, 1.0, -1.0];
        let orig = s.clone();
        apply_smoothing(&mut s, 0.0);
        assert_eq!(s, orig);
    }

    #[test]
    fn smoothing_attenuates_oscillation() {
        let mut s = vec![1.0_f32, -1.0, 1.0, -1.0, 1.0, -1.0, 1.0, -1.0];
        apply_smoothing(&mut s, 0.9);
        // After smoothing, peak-to-peak amplitude should drop significantly
        // (one-pole IIR with α≈0.75).
        let max = s.iter().cloned().fold(f32::MIN, f32::max);
        let min = s.iter().cloned().fold(f32::MAX, f32::min);
        assert!(
            (max - min) < 1.5,
            "expected smoothed wave to attenuate (got peak-to-peak {})",
            max - min
        );
    }

    #[test]
    fn build_uniforms_maps_fields() {
        let params = WaveParams {
            r: 1.0,
            g: 0.5,
            b: 0.25,
            a: 0.8,
            x: 0.4,
            y: 0.6,
            mode: 4,
            scale: 0.7,
            param: 0.3,
            smoothing: 0.5,
            thick: true,
            dots: false,
            additive: true,
            maximize_color: false,
            mod_alpha_by_volume: false,
            mod_alpha_start: 0.0,
            mod_alpha_end: 1.0,
            split_lr: false,
        };
        let u = build_uniforms(params, 0.0, 1280, 720, 1.5);
        assert_eq!(u.mode, WaveformMode::DerivativeLine as u32);
        assert!((u.wave_x - 0.4).abs() < 1e-5);
        assert!((u.aspect - (1280.0 / 720.0)).abs() < 1e-3);
        assert_eq!(u.is_thick, 1);
        assert_eq!(u.is_dots, 0);
    }

    #[test]
    fn build_uniforms_maximize_color_normalizes_rgb() {
        let params = WaveParams {
            r: 0.5,
            g: 0.25,
            b: 0.25,
            a: 1.0,
            x: 0.5,
            y: 0.5,
            mode: 0,
            scale: 0.5,
            param: 0.0,
            smoothing: 0.0,
            thick: false,
            dots: false,
            additive: false,
            maximize_color: true,
            mod_alpha_by_volume: false,
            mod_alpha_start: 0.0,
            mod_alpha_end: 1.0,
            split_lr: false,
        };
        let u = build_uniforms(params, 0.0, 100, 100, 0.0);
        assert!((u.wave_color[0] - 1.0).abs() < 1e-5);
        assert!((u.wave_color[1] - 0.5).abs() < 1e-5);
        assert!((u.wave_color[2] - 0.5).abs() < 1e-5);
    }

    #[test]
    fn build_uniforms_mod_alpha_by_volume_scales_alpha() {
        let params = WaveParams {
            r: 1.0,
            g: 1.0,
            b: 1.0,
            a: 1.0,
            x: 0.5,
            y: 0.5,
            mode: 0,
            scale: 0.5,
            param: 0.0,
            smoothing: 0.0,
            thick: false,
            dots: false,
            additive: false,
            maximize_color: false,
            mod_alpha_by_volume: true,
            mod_alpha_start: 0.0,
            mod_alpha_end: 1.0,
            split_lr: false,
        };
        let low = build_uniforms(params, 0.1, 100, 100, 0.0);
        let high = build_uniforms(params, 0.9, 100, 100, 0.0);
        assert!(low.wave_color[3] < high.wave_color[3]);
    }
}