onedrop_renderer/

motion_vector.rs

//! Motion-vector grid pass.
//!
//! MD2 draws a grid of `n_x × n_y` short line segments into the warp
//! output every frame. The renderer is stateless — segments live for
//! exactly one frame, painted at the same screen-space grid each time
//! — but because the segments end up inside the feedback loop (we
//! dispatch between the shape overlays and the blur pyramid), the
//! warp pulls them into trails. Those trails are the "motion vectors"
//! the player sees.
//!
//! The grid is clamped to MD2's hard cap of `64 × 48`; combined with
//! the 2 vertices per segment that gives `MAX_MOTION_VECTOR_SEGMENTS *
//! 2` vertices per frame upload — small enough to keep the buffer at
//! a static size.
//!
//! Segment geometry per cell `(i, j)`:
//!
//! - Anchor: `(x, y) = ((i + 0.5) / n_x + dx, (j + 0.5) / n_y + dy)` in
//!   MD2 preset space `[0, 1]` (origin top-left), then mapped to clip
//!   space with `preset_xy_to_clip`.
//! - Tail: anchor + `(length / n_x, 0)` — a small horizontal segment.
//!   `length` is `mv_l` (MD2 default 0.9) expressed in cells so a value
//!   of `1.0` would span an entire cell, with the visible trail length
//!   emerging from how aggressively the feedback warp pulls it.

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

use crate::config::MotionVectorParams;
use crate::warp_pipeline::WarpVertex;

/// Maximum cells along each axis. MD2 caps motion-vector grids at this
/// size in the preset format; the renderer enforces the cap regardless
/// of what the preset declares.
pub const MAX_MOTION_VECTOR_GRID_X: u32 = 64;
pub const MAX_MOTION_VECTOR_GRID_Y: u32 = 48;

/// Maximum number of segments stored in the GPU buffer.
pub const MAX_MOTION_VECTOR_SEGMENTS: usize =
    (MAX_MOTION_VECTOR_GRID_X as usize) * (MAX_MOTION_VECTOR_GRID_Y as usize);

/// One vertex of the motion-vector stream. Two per segment (head +
/// tail). Pure passthrough — the shader doesn't transform or shade,
/// just emits the per-segment colour.
#[repr(C)]
#[derive(Debug, Clone, Copy, PartialEq, Pod, Zeroable, Default)]
pub struct MotionVectorVertex {
    pub pos: [f32; 2],
    pub color: [f32; 4],
}

/// GPU pass for the MD2 motion-vector grid.
pub struct MotionVectorRenderer {
    pipeline: wgpu::RenderPipeline,
    vertex_buffer: wgpu::Buffer,
    /// Logical vertex count after the most recent [`Self::update`].
    /// Always `segments * 2`.
    vertex_count: u32,
}

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

        // Allocate the buffer at max size up front so we never need a
        // reallocation in the steady state.
        let initial: Vec<MotionVectorVertex> =
            vec![MotionVectorVertex::default(); MAX_MOTION_VECTOR_SEGMENTS * 2];
        let vertex_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
            label: Some("Motion Vector Vertices"),
            contents: bytemuck::cast_slice(&initial),
            usage: wgpu::BufferUsages::VERTEX | wgpu::BufferUsages::COPY_DST,
        });

        let layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
            label: Some("Motion Vector 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::<MotionVectorVertex>() as u64,
            step_mode: wgpu::VertexStepMode::Vertex,
            attributes: &vertex_attributes,
        };

        let pipeline = device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
            label: Some("Motion Vector Pipeline"),
            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(wgpu::BlendState::ALPHA_BLENDING),
                    write_mask: wgpu::ColorWrites::ALL,
                })],
                compilation_options: Default::default(),
            }),
            primitive: wgpu::PrimitiveState {
                topology: wgpu::PrimitiveTopology::LineList,
                ..Default::default()
            },
            depth_stencil: None,
            multisample: wgpu::MultisampleState::default(),
            multiview_mask: None,
            cache: None,
        });

        Self {
            pipeline,
            vertex_buffer,
            vertex_count: 0,
        }
    }

    /// Rebuild the per-frame segment stream from a [`MotionVectorParams`]
    /// snapshot and upload it. Returns the number of segments produced
    /// — `0` when the pass is a no-op (zero alpha, zero grid, etc).
    ///
    /// `warp_field` (when `Some`) carries the per-vertex `uv_warp` array
    /// from the warp mesh; the segment at grid cell `(i, j)` then points
    /// along the local warp displacement (`uv_warp - uv_orig`) instead
    /// of a flat horizontal stub. Matches MD2's "streaks" feel: the
    /// segments anticipate the trail the feedback loop will pull them
    /// into. Pass `None` to fall back to the horizontal-stub default.
    pub fn update(
        &mut self,
        queue: &wgpu::Queue,
        params: MotionVectorParams,
        warp_field: Option<WarpField<'_>>,
    ) -> u32 {
        let nx = params.grid_x.min(MAX_MOTION_VECTOR_GRID_X);
        let ny = params.grid_y.min(MAX_MOTION_VECTOR_GRID_Y);
        let alpha = params.color[3];
        if nx == 0 || ny == 0 || alpha <= 0.0 {
            self.vertex_count = 0;
            return 0;
        }

        let mut vertices: Vec<MotionVectorVertex> =
            Vec::with_capacity((nx as usize) * (ny as usize) * 2);
        build_segments(nx, ny, params, warp_field, &mut vertices);

        if !vertices.is_empty() {
            queue.write_buffer(&self.vertex_buffer, 0, bytemuck::cast_slice(&vertices));
        }
        self.vertex_count = vertices.len() as u32;
        (vertices.len() / 2) as u32
    }

    pub fn render(&self, encoder: &mut wgpu::CommandEncoder, view: &wgpu::TextureView) {
        if self.vertex_count == 0 {
            return;
        }
        let mut rp = encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
            label: Some("Motion Vector 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(&self.pipeline);
        rp.set_vertex_buffer(0, self.vertex_buffer.slice(..));
        rp.draw(0..self.vertex_count, 0..1);
    }

    /// Number of segments the next [`Self::render`] will issue (two
    /// vertices each). Exposed for tests.
    pub fn segment_count(&self) -> u32 {
        self.vertex_count / 2
    }
}

/// Per-vertex warp-displacement field, borrowed from the renderer's
/// most recent `update_warp_vertices` call. `vertices[r * cols + c]`
/// matches the warp mesh's row-major layout; the segment builder
/// recovers `uv_orig` from `pos_clip` (inverse of the same mapping
/// `WarpMesh::new` uses) and reports the displacement
/// `uv_warp - uv_orig` at the grid cell's nearest mesh vertex.
pub struct WarpField<'a> {
    pub cols: u32,
    pub rows: u32,
    pub vertices: &'a [WarpVertex],
}

/// Build the per-frame segment list. Pulled out as a free function so
/// it can be unit-tested without a GPU device. When `warp_field` is
/// `Some`, segments point along the local warp displacement; without
/// it the function falls back to a horizontal stub.
fn build_segments(
    nx: u32,
    ny: u32,
    params: MotionVectorParams,
    warp_field: Option<WarpField<'_>>,
    vertices: &mut Vec<MotionVectorVertex>,
) {
    let inv_nx = 1.0 / nx as f32;
    let inv_ny = 1.0 / ny as f32;
    // Segment length expressed as a fraction of cell spacing. MD2's
    // `mv_l` is already in those units; we just multiply by the cell
    // width in preset space.
    let seg_len_preset = params.length * inv_nx;

    for j in 0..ny {
        for i in 0..nx {
            let cx_preset = (i as f32 + 0.5) * inv_nx + params.dx;
            let cy_preset = (j as f32 + 0.5) * inv_ny + params.dy;
            let (dir_x, dir_y) = if let Some(field) = warp_field.as_ref() {
                sample_warp_displacement(field, cx_preset, cy_preset)
            } else {
                // Default horizontal stub.
                (seg_len_preset, 0.0)
            };
            // When the warp displacement is supplied in *preset-space
            // UV units*, scale by `length` so `mv_l` keeps its MD2
            // meaning ("fraction of a cell wide"). The default stub
            // already incorporated `length` so we skip the scaling
            // there.
            let (sx, sy) = if warp_field.is_some() {
                (dir_x * params.length * 8.0, dir_y * params.length * 8.0)
            } else {
                (dir_x, dir_y)
            };
            let head = preset_xy_to_clip(cx_preset, cy_preset);
            let tail = preset_xy_to_clip(cx_preset + sx, cy_preset + sy);
            vertices.push(MotionVectorVertex {
                pos: head,
                color: params.color,
            });
            vertices.push(MotionVectorVertex {
                pos: tail,
                color: params.color,
            });
        }
    }
}

/// Look up the warp displacement at preset-space position `(x, y)`
/// from `field`. Bilinear-interpolates across the four nearest mesh
/// cells so the segment direction stays smooth across the grid
/// boundaries. Returns `(dx, dy)` in preset-space UV units
/// (`uv_warp - uv_orig`).
fn sample_warp_displacement(field: &WarpField<'_>, x: f32, y: f32) -> (f32, f32) {
    if field.cols < 2 || field.rows < 2 || field.vertices.is_empty() {
        return (0.0, 0.0);
    }
    let cols_f = (field.cols - 1) as f32;
    let rows_f = (field.rows - 1) as f32;
    let mc = (x.clamp(0.0, 1.0)) * cols_f;
    let mr = (y.clamp(0.0, 1.0)) * rows_f;
    let c0 = mc.floor().clamp(0.0, cols_f - 1.0) as u32;
    let r0 = mr.floor().clamp(0.0, rows_f - 1.0) as u32;
    let c1 = (c0 + 1).min(field.cols - 1);
    let r1 = (r0 + 1).min(field.rows - 1);
    let tu = mc - c0 as f32;
    let tv = mr - r0 as f32;

    let idx = |c: u32, r: u32| -> usize { (r * field.cols + c) as usize };
    let disp = |v: &WarpVertex| -> (f32, f32) {
        // `pos_clip` is `(uv_orig.x * 2 - 1, 1 - uv_orig.y * 2)`; invert to
        // recover the original UV the mesh was sampled at.
        let ox = (v.pos_clip[0] + 1.0) * 0.5;
        let oy = (1.0 - v.pos_clip[1]) * 0.5;
        (v.uv_warp[0] - ox, v.uv_warp[1] - oy)
    };
    let d00 = disp(&field.vertices[idx(c0, r0)]);
    let d10 = disp(&field.vertices[idx(c1, r0)]);
    let d01 = disp(&field.vertices[idx(c0, r1)]);
    let d11 = disp(&field.vertices[idx(c1, r1)]);

    let mix = |a: f32, b: f32, t: f32| a + (b - a) * t;
    let dx = mix(mix(d00.0, d10.0, tu), mix(d01.0, d11.0, tu), tv);
    let dy = mix(mix(d00.1, d10.1, tu), mix(d01.1, d11.1, tu), tv);
    (dx, dy)
}

/// MD2 preset-space → clip-space, identical to the helpers in
/// `custom_wave` / `custom_shape`. Origin moves from top-left at
/// `(0, 0)` to centre at `(0, 0)`, range `[0, 1]` becomes `[-1, 1]`,
/// and y is flipped.
#[inline]
fn preset_xy_to_clip(x: f32, y: f32) -> [f32; 2] {
    [x * 2.0 - 1.0, 1.0 - y * 2.0]
}

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

    fn enabled_params() -> MotionVectorParams {
        MotionVectorParams {
            grid_x: 4,
            grid_y: 3,
            dx: 0.0,
            dy: 0.0,
            length: 0.5,
            color: [1.0, 1.0, 1.0, 1.0],
        }
    }

    #[test]
    fn segment_count_matches_grid_size() {
        let mut v = Vec::new();
        build_segments(4, 3, enabled_params(), None, &mut v);
        assert_eq!(v.len(), 4 * 3 * 2);
    }

    #[test]
    fn segment_endpoints_are_horizontal() {
        let mut v = Vec::new();
        build_segments(4, 3, enabled_params(), None, &mut v);
        for pair in v.chunks(2) {
            // Each pair has the same y, different x.
            assert!((pair[0].pos[1] - pair[1].pos[1]).abs() < 1e-6);
            assert!(pair[0].pos[0] < pair[1].pos[0]);
        }
    }

    #[test]
    fn dx_dy_shift_anchor() {
        let mut a = Vec::new();
        let mut b = Vec::new();
        let p0 = enabled_params();
        let mut p1 = p0;
        p1.dx = 0.25;
        p1.dy = -0.1;
        build_segments(2, 2, p0, None, &mut a);
        build_segments(2, 2, p1, None, &mut b);
        // Both should produce the same count.
        assert_eq!(a.len(), b.len());
        // The shift `dx = 0.25` in preset space is `+0.5` in clip-x;
        // `dy = -0.1` is `+0.2` in clip-y (y is flipped).
        let dx_clip = 0.25 * 2.0;
        let dy_clip = -(-0.1) * 2.0;
        for (av, bv) in a.iter().zip(b.iter()) {
            assert!((bv.pos[0] - av.pos[0] - dx_clip).abs() < 1e-5);
            assert!((bv.pos[1] - av.pos[1] - dy_clip).abs() < 1e-5);
        }
    }

    #[test]
    fn segment_carries_color() {
        let mut v = Vec::new();
        let mut p = enabled_params();
        p.color = [0.2, 0.4, 0.6, 0.5];
        build_segments(2, 2, p, None, &mut v);
        for vert in &v {
            assert_eq!(vert.color, p.color);
        }
    }

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

    /// `sample_warp_displacement` should bilinear-interpolate a
    /// constant offset across a 2×2 mesh and recover that offset
    /// (within float epsilon) at every sampled position. Locks the
    /// inverse `pos_clip → uv_orig` math against the warp mesh's
    /// forward mapping.
    #[test]
    fn warp_field_sampler_returns_constant_offset() {
        // 2 × 2 warp vertices laid out exactly like `WarpMesh::new`:
        // pos_clip in {(-1, -1), (1, -1), (-1, 1), (1, 1)}, uv_orig
        // in {(0, 1), (1, 1), (0, 0), (1, 0)}. Each gets a constant
        // displacement `(+0.1, -0.05)`.
        let off = (0.1f32, -0.05f32);
        let mk = |pc: [f32; 2], uv_orig: [f32; 2]| WarpVertex {
            pos_clip: pc,
            uv_warp: [uv_orig[0] + off.0, uv_orig[1] + off.1],
        };
        let verts = vec![
            mk([-1.0, -1.0], [0.0, 1.0]), // row 0, col 0
            mk([1.0, -1.0], [1.0, 1.0]),  // row 0, col 1
            mk([-1.0, 1.0], [0.0, 0.0]),  // row 1, col 0
            mk([1.0, 1.0], [1.0, 0.0]),   // row 1, col 1
        ];
        let field = WarpField {
            cols: 2,
            rows: 2,
            vertices: &verts,
        };
        for (x, y) in [(0.0, 0.0), (0.5, 0.5), (1.0, 1.0), (0.25, 0.75)] {
            let (dx, dy) = sample_warp_displacement(&field, x, y);
            assert!(
                (dx - off.0).abs() < 1e-5 && (dy - off.1).abs() < 1e-5,
                "displacement at ({x}, {y}) was ({dx}, {dy})"
            );
        }
    }

    /// When `warp_field` is supplied with a non-zero per-vertex offset,
    /// segment tails should depart from the horizontal default. The
    /// test seeds a uniform vertical-only displacement `(0, +0.2)` and
    /// checks that every segment ends up with a vertical-only delta.
    #[test]
    fn segments_pick_up_warp_field_direction() {
        // 2 × 2 mesh, uniform offset `(0, +0.2)` in preset-space UV.
        let off_y = 0.2f32;
        let mk = |pc: [f32; 2], uv_orig: [f32; 2]| WarpVertex {
            pos_clip: pc,
            uv_warp: [uv_orig[0], uv_orig[1] + off_y],
        };
        let verts = vec![
            mk([-1.0, -1.0], [0.0, 1.0]),
            mk([1.0, -1.0], [1.0, 1.0]),
            mk([-1.0, 1.0], [0.0, 0.0]),
            mk([1.0, 1.0], [1.0, 0.0]),
        ];
        let field = WarpField {
            cols: 2,
            rows: 2,
            vertices: &verts,
        };
        let mut v = Vec::new();
        build_segments(2, 2, enabled_params(), Some(field), &mut v);
        for pair in v.chunks(2) {
            let dx = pair[1].pos[0] - pair[0].pos[0];
            let dy = pair[1].pos[1] - pair[0].pos[1];
            // Vertical-only offset in preset space → preset y stays
            // unchanged in x, shifts down in clip-y (y flipped). So
            // dx ≈ 0 and dy < 0.
            assert!(dx.abs() < 1e-5, "x drift: {dx}");
            assert!(dy < -1e-3, "expected downward tail, got {dy}");
        }
    }
}