onedrop_renderer/

custom_wave.rs

//! Custom-wave (`wavecode_N`) pass.
//!
//! The MD2 `wavecode_N` block defines up to four user-driven waves that
//! run their own `per_frame_init` / `per_frame` / `per_point` equations.
//! Per frame the engine evaluates each enabled wave's per-point loop on
//! the CPU, producing a stream of `(x, y, r, g, b, a)` 6-tuples; this
//! module uploads that stream into a dynamic vertex buffer and dispatches
//! either a line-list or triangle-list draw.
//!
//! Why CPU-emit vs. GPU-derive: each preset's per-point equations can
//! reference arbitrary user state (`q*` channels, the audio buffer,
//! whatever a preset author thought up) — running them on the GPU would
//! require porting the evaluator. The line/dot geometry is tiny enough
//! that uploading expanded vertices each frame is the cheaper path.
//!
//! Pipeline matrix:
//! - **Topology** is chosen at submit time: `LineStrip` for `b_use_dots = 0`
//!   (one strip per wave; segments break at wave boundaries via
//!   per-wave draw calls), `TriangleList` for `b_use_dots = 1`
//!   (each point expands CPU-side into a small 2-tri quad).
//! - **Blend** is alpha or additive, picked by `b_additive`.
//!
//! Dispatch slot: the renderer calls into this pass between the legacy
//! waveform pass and the blur pyramid, so custom waves feed back next
//! frame and contribute to `GetBlur*`. Same slot the historical
//! waveform pass uses.

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

/// One vertex of the custom-wave stream. The shader is pure passthrough
/// — `pos` is already clip-space `[-1, 1]`, `color` is the per-point RGBA
/// after `r/g/b/a` carried through the per-point loop.
#[repr(C)]
#[derive(Debug, Clone, Copy, PartialEq, Pod, Zeroable, Default)]
pub struct CustomWaveVertex {
    pub pos: [f32; 2],
    pub color: [f32; 4],
}

/// One dispatch unit: the slice of the vertex buffer that belongs to a
/// single `wavecode_N` block, plus its blend/topology flags. The renderer
/// reuses a single growable buffer across all waves and walks the batches
/// to issue one draw call each.
#[derive(Debug, Clone, Copy)]
pub struct CustomWaveBatch {
    /// Index into the renderer's vertex buffer where this wave's stream
    /// starts.
    pub start_vertex: u32,
    /// Number of vertices this wave contributes.
    pub vertex_count: u32,
    /// `true` → triangle-list topology. Holds for both `b_use_dots = 1`
    /// (one quad per point) and `b_draw_thick = 1` with lines (one quad
    /// per *segment* between consecutive per-point outputs). `false`
    /// → line-strip from the per-point trail (thin, 1 px wide).
    /// Field name kept for backward compat — historically only `dots`
    /// drove TriangleList.
    pub dots: bool,
    /// `true` → additive blend, otherwise alpha.
    pub additive: bool,
}

/// Maximum total custom-wave vertices we'll buffer in a frame. Four
/// waves × 512 samples × (2 lines = 6 verts/segment, dots = 6
/// verts/point) ≈ 12 288. Round up to 16 384 to leave headroom for
/// the next-sprint custom-shape pass, which will reuse this buffer
/// layout.
pub const MAX_CUSTOM_WAVE_VERTICES: usize = 16_384;

/// Custom-wave GPU pass. Owns one growable vertex buffer + two
/// pipelines (lines and dots topology, alpha and additive blend = 4
/// combinations, indexed by `pipeline_index(dots, additive)`).
pub struct CustomWaveRenderer {
    pipelines: [wgpu::RenderPipeline; 4],
    vertex_buffer: wgpu::Buffer,
    /// Logical vertex count after the most recent upload (0 when no
    /// custom waves are enabled).
    vertex_count: u32,
    /// Per-wave dispatch metadata, in stream order.
    batches: Vec<CustomWaveBatch>,
}

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

impl CustomWaveRenderer {
    pub fn new(device: &wgpu::Device, format: wgpu::TextureFormat) -> Self {
        let shader = crate::pipeline_helpers::load_wgsl(
            device,
            "Custom Wave Shader",
            include_str!("../shaders/custom_wave.wgsl"),
        );

        let initial: Vec<CustomWaveVertex> =
            vec![CustomWaveVertex::default(); MAX_CUSTOM_WAVE_VERTICES];
        let vertex_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
            label: Some("Custom Wave Vertices"),
            contents: bytemuck::cast_slice(&initial),
            usage: wgpu::BufferUsages::VERTEX | wgpu::BufferUsages::COPY_DST,
        });

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

        let vertex_attributes = [
            wgpu::VertexAttribute {
                offset: 0,
                shader_location: 0,
                format: wgpu::VertexFormat::Float32x2,
            },
            wgpu::VertexAttribute {
                offset: std::mem::size_of::<[f32; 2]>() as u64,
                shader_location: 1,
                format: wgpu::VertexFormat::Float32x4,
            },
        ];
        let vertex_layout = wgpu::VertexBufferLayout {
            array_stride: std::mem::size_of::<CustomWaveVertex>() as u64,
            step_mode: wgpu::VertexStepMode::Vertex,
            attributes: &vertex_attributes,
        };

        let make_pipeline = |label: &str,
                             topology: wgpu::PrimitiveTopology,
                             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: std::slice::from_ref(&vertex_layout),
                    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,
                    ..Default::default()
                },
                depth_stencil: None,
                multisample: wgpu::MultisampleState::default(),
                multiview_mask: None,
                cache: None,
            })
        };

        let pipelines = [
            make_pipeline(
                "CustomWave Lines Alpha",
                wgpu::PrimitiveTopology::LineStrip,
                false,
            ),
            make_pipeline(
                "CustomWave Lines Additive",
                wgpu::PrimitiveTopology::LineStrip,
                true,
            ),
            make_pipeline(
                "CustomWave Dots Alpha",
                wgpu::PrimitiveTopology::TriangleList,
                false,
            ),
            make_pipeline(
                "CustomWave Dots Additive",
                wgpu::PrimitiveTopology::TriangleList,
                true,
            ),
        ];

        Self {
            pipelines,
            vertex_buffer,
            vertex_count: 0,
            batches: Vec::new(),
        }
    }

    /// Replace the per-frame vertex stream and batch list. The combined
    /// vertex count is clamped to [`MAX_CUSTOM_WAVE_VERTICES`]; any tail
    /// past that is dropped (with a `log::warn` so it's visible in
    /// long-running sessions).
    pub fn update(
        &mut self,
        queue: &wgpu::Queue,
        vertices: &[CustomWaveVertex],
        batches: &[CustomWaveBatch],
    ) {
        let n = vertices.len().min(MAX_CUSTOM_WAVE_VERTICES);
        if vertices.len() > MAX_CUSTOM_WAVE_VERTICES {
            log::warn!(
                "custom-wave vertex stream truncated: {} > cap {}",
                vertices.len(),
                MAX_CUSTOM_WAVE_VERTICES
            );
        }
        if n > 0 {
            queue.write_buffer(&self.vertex_buffer, 0, bytemuck::cast_slice(&vertices[..n]));
        }
        self.vertex_count = n as u32;
        self.batches.clear();
        // Filter out any batch that points past the truncation boundary.
        for b in batches {
            if b.start_vertex + b.vertex_count <= self.vertex_count {
                self.batches.push(*b);
            }
        }
    }

    /// Issue one draw call per stored batch into `view`. Loads existing
    /// contents (no clear) — meant to draw on top of the warp + legacy
    /// waveform output.
    pub fn render(&self, encoder: &mut wgpu::CommandEncoder, view: &wgpu::TextureView) {
        if self.batches.is_empty() || self.vertex_count == 0 {
            return;
        }
        let mut rp = encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
            label: Some("Custom Wave 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_vertex_buffer(0, self.vertex_buffer.slice(..));
        for b in &self.batches {
            if b.vertex_count == 0 {
                continue;
            }
            let pidx = pipeline_index(b.dots, b.additive);
            rp.set_pipeline(&self.pipelines[pidx]);
            let end = b.start_vertex + b.vertex_count;
            rp.draw(b.start_vertex..end, 0..1);
        }
    }

    /// Number of vertices held by the most recent [`Self::update`] call.
    /// Exposed for tests asserting the upload size.
    pub fn vertex_count(&self) -> u32 {
        self.vertex_count
    }

    /// Number of dispatch batches the next [`Self::render`] will issue.
    pub fn batch_count(&self) -> usize {
        self.batches.len()
    }
}

/// Convert MD2 preset-space `(x, y)` (origin top-left, range `[0, 1]`)
/// into clip-space `(x', y')` (origin centre, range `[-1, 1]`, Y-up).
#[inline]
pub fn preset_xy_to_clip(x: f64, y: f64) -> [f32; 2] {
    [(x * 2.0 - 1.0) as f32, (1.0 - y * 2.0) as f32]
}

/// Expand one MD2 per-point output `(x, y, r, g, b, a)` into a small
/// screen-space quad (two triangles, six vertices). Used by the dots
/// pipeline. `radius_clip` is the half-side in clip space — pick
/// `2.0 / height` for a single-pixel-equivalent dot at default sizing.
pub fn point_to_dot_quad(
    x: f64,
    y: f64,
    color: [f32; 4],
    radius_clip: f32,
) -> [CustomWaveVertex; 6] {
    let [cx, cy] = preset_xy_to_clip(x, y);
    let lo_x = cx - radius_clip;
    let hi_x = cx + radius_clip;
    let lo_y = cy - radius_clip;
    let hi_y = cy + radius_clip;
    let tl = CustomWaveVertex {
        pos: [lo_x, hi_y],
        color,
    };
    let tr = CustomWaveVertex {
        pos: [hi_x, hi_y],
        color,
    };
    let bl = CustomWaveVertex {
        pos: [lo_x, lo_y],
        color,
    };
    let br = CustomWaveVertex {
        pos: [hi_x, lo_y],
        color,
    };
    [tl, tr, bl, tr, br, bl]
}

/// Expand two consecutive per-point outputs into a thick-line segment
/// quad — six clip-space vertices forming a screen-space rectangle of
/// `thickness_clip` perpendicular to the segment direction. Mirrors
/// the `line_vertex` shape used by `waveform_advanced.wgsl` for
/// `b_wave_thick`, so MD2's "thick custom-wave trail" looks identical
/// to the thick-static-wave path.
///
/// `color_a` / `color_b` come from the two source points so the segment
/// gracefully interpolates colour along a fading per-point trail.
/// Caller supplies `(x, y)` in MD2 preset space; we convert internally.
pub fn segment_to_thick_quad(
    ax: f64,
    ay: f64,
    color_a: [f32; 4],
    bx: f64,
    by: f64,
    color_b: [f32; 4],
    thickness_clip: f32,
) -> [CustomWaveVertex; 6] {
    let [ax_c, ay_c] = preset_xy_to_clip(ax, ay);
    let [bx_c, by_c] = preset_xy_to_clip(bx, by);
    let dx = bx_c - ax_c;
    let dy = by_c - ay_c;
    let len = (dx * dx + dy * dy).sqrt().max(1e-6);
    // Perpendicular `(–dy, dx) / len` × half-thickness.
    let half = thickness_clip * 0.5;
    let nx = -dy / len * half;
    let ny = dx / len * half;
    let a_minus = CustomWaveVertex {
        pos: [ax_c - nx, ay_c - ny],
        color: color_a,
    };
    let a_plus = CustomWaveVertex {
        pos: [ax_c + nx, ay_c + ny],
        color: color_a,
    };
    let b_minus = CustomWaveVertex {
        pos: [bx_c - nx, by_c - ny],
        color: color_b,
    };
    let b_plus = CustomWaveVertex {
        pos: [bx_c + nx, by_c + ny],
        color: color_b,
    };
    [a_minus, b_minus, b_plus, a_minus, b_plus, a_plus]
}

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

    #[test]
    fn preset_xy_to_clip_center_is_origin() {
        let p = preset_xy_to_clip(0.5, 0.5);
        assert!((p[0] - 0.0).abs() < 1e-5);
        assert!((p[1] - 0.0).abs() < 1e-5);
    }

    #[test]
    fn preset_xy_to_clip_topleft_is_neg1_pos1() {
        let p = preset_xy_to_clip(0.0, 0.0);
        assert!((p[0] - -1.0).abs() < 1e-5);
        assert!((p[1] - 1.0).abs() < 1e-5);
    }

    #[test]
    fn preset_xy_to_clip_bottomright_is_pos1_neg1() {
        let p = preset_xy_to_clip(1.0, 1.0);
        assert!((p[0] - 1.0).abs() < 1e-5);
        assert!((p[1] - -1.0).abs() < 1e-5);
    }

    /// A horizontal segment `(0, 0) → (1, 0)` in preset space becomes
    /// a screen-space rectangle whose perpendicular offset is
    /// straight up/down (positive Y is the screen *up* direction in
    /// clip-space; preset space flips Y so MD2 y=0 maps to clip y=1).
    /// We only care that:
    /// - the 6 vertices form a closed quad (no degenerate triangles),
    /// - the half-width offset matches the requested thickness,
    /// - the two endpoints' colours carry through to their respective
    ///   pairs.
    #[test]
    fn segment_to_thick_quad_horizontal_yields_perpendicular_offset() {
        let q = segment_to_thick_quad(
            0.0,
            0.5,
            [1.0, 0.0, 0.0, 1.0],
            1.0,
            0.5,
            [0.0, 1.0, 0.0, 1.0],
            0.01,
        );
        // Endpoint A (preset (0, 0.5)) → clip (-1, 0).
        // Endpoint B (preset (1, 0.5)) → clip ( 1, 0).
        // Perpendicular to a horizontal segment is vertical, so x ≈
        // endpoint x and y deviates by ±0.005 (half of 0.01).
        let half = 0.005f32;
        for v in q.iter() {
            // Either near x=-1 or near x=1 (no horizontal drift).
            let near_a = (v.pos[0] + 1.0).abs() < 1e-4;
            let near_b = (v.pos[0] - 1.0).abs() < 1e-4;
            assert!(near_a || near_b, "x off-axis: {v:?}");
            assert!(
                (v.pos[1].abs() - half).abs() < 1e-4,
                "y not at ±half-thickness: {v:?}"
            );
            if near_a {
                assert_eq!(v.color, [1.0, 0.0, 0.0, 1.0]);
            }
            if near_b {
                assert_eq!(v.color, [0.0, 1.0, 0.0, 1.0]);
            }
        }
    }

    #[test]
    fn point_to_dot_quad_centered_at_xy() {
        let q = point_to_dot_quad(0.5, 0.5, [1.0, 0.0, 0.0, 1.0], 0.01);
        // Two triangles share two vertices; all six should be near origin.
        for v in q.iter() {
            assert!(v.pos[0].abs() <= 0.011);
            assert!(v.pos[1].abs() <= 0.011);
            assert_eq!(v.color, [1.0, 0.0, 0.0, 1.0]);
        }
    }

    #[test]
    fn pipeline_index_table() {
        assert_eq!(pipeline_index(false, false), 0); // lines alpha
        assert_eq!(pipeline_index(false, true), 1); // lines additive
        assert_eq!(pipeline_index(true, false), 2); // dots alpha
        assert_eq!(pipeline_index(true, true), 3); // dots additive
    }
}