onedrop_renderer/

custom_shape.rs

//! Custom-shape (`shapecode_N`) pass.
//!
//! Up to four user-defined shapes per preset. Each shape spawns
//! `num_inst` instances (1..1024) drawn as triangle fans of `sides`
//! vertices around `(x, y)` at radius `rad` rotated by `ang`. The
//! engine runs the shape's per-frame equations once per instance,
//! collects the resulting state into a flat storage buffer, and
//! issues one fill draw call per shape; an optional outline pass
//! follows for `b_thick_outline = 1`.
//!
//! Why a storage buffer + instanced draw (vs. CPU-expanded vertex
//! stream like custom waves): a single shape with `num_inst = 1024` and
//! `sides = 64` would emit `1024 × 64 × 3 ≈ 200K vertices` per frame —
//! the upload bandwidth dominates the work. With instancing the upload
//! is one struct (~96 bytes) per instance, and the GPU expands the
//! geometry from `(@builtin(vertex_index), @builtin(instance_index))`.
//!
//! Textured mode: instead of a solid colour, the fragment shader
//! samples `prev_texture` (the post-warp / post-wave output of the
//! previous frame) at UV warped by `tex_zoom` / `tex_ang`. The output
//! is multiplied by the centre→edge gradient so the user's RGB still
//! tints the result.

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

/// GPU-side per-instance record. `repr(C)` + `Pod` for direct
/// `bytemuck::cast_slice` upload to a storage buffer. Each field is a
/// `vec4`-aligned 16-byte block so the WGSL `array<ShapeInstance>` layout
/// matches without an explicit `align(16)` shim.
#[repr(C)]
#[derive(Debug, Clone, Copy, PartialEq, Pod, Zeroable, Default)]
pub struct CustomShapeInstance {
    /// `[center_x, center_y, radius_x, radius_y]` — center is
    /// clip-space (-1..1); radius is in clip units before the shader's
    /// aspect correction.
    pub center_radius: [f32; 4],
    /// `[angle (radians), tex_zoom, tex_ang (radians), sides (f32 → u32)]`.
    pub angle_zoom_ang_sides: [f32; 4],
    pub color_center: [f32; 4],
    pub color_edge: [f32; 4],
    pub border_color: [f32; 4],
    /// `[border_size, flag_bits, pad, pad]` where `flag_bits` is:
    /// - bit 0: `textured`
    /// - bit 1: `additive`
    /// - bit 2: `thick_outline`
    pub border_flags: [f32; 4],
}

/// Per-instance flag bits packed into `CustomShapeInstance::border_flags[1]`.
/// Reading side lives in `custom_shape.wgsl` (`flag_bits & 1u`, etc.).
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct ShapeFlags(u32);

impl ShapeFlags {
    pub const TEXTURED: Self = Self(0b001);
    pub const ADDITIVE: Self = Self(0b010);
    pub const THICK_OUTLINE: Self = Self(0b100);

    pub const fn empty() -> Self {
        Self(0)
    }
    pub const fn bits(self) -> u32 {
        self.0
    }
    pub fn contains(self, other: Self) -> bool {
        (self.0 & other.0) == other.0
    }
    pub fn pack(self) -> f32 {
        // Round-trip through f32 since the storage buffer layout uses
        // f32 throughout. `u32 → f32` is exact for values up to 2^24.
        self.0 as f32
    }
}

impl std::ops::BitOr for ShapeFlags {
    type Output = Self;
    fn bitor(self, rhs: Self) -> Self {
        Self(self.0 | rhs.0)
    }
}

impl std::ops::BitOrAssign for ShapeFlags {
    fn bitor_assign(&mut self, rhs: Self) {
        self.0 |= rhs.0;
    }
}

/// One dispatch unit: a contiguous run of instances in the storage
/// buffer that belong to the same shape and share `sides` + blend mode.
#[derive(Debug, Clone, Copy)]
pub struct CustomShapeBatch {
    pub start_instance: u32,
    pub instance_count: u32,
    /// Effective sides for the draw call's vertex count: `sides * 3`
    /// triangle-list vertices per instance. Clamped to `[3, MAX_SIDES]`.
    pub sides: u32,
    /// `true` → additive blend pipeline, `false` → alpha.
    pub additive: bool,
    /// `true` → also dispatch the LineStrip outline pass for this batch
    /// using the same instance range.
    pub thick_outline: bool,
}

/// MD2 caps `sides` at 100 in the preset format. We clamp to 64 here
/// to keep the per-instance vertex count manageable (64 × 3 = 192
/// vertices/instance × 1024 instances = ~200K verts max per shape).
pub const MAX_SHAPE_SIDES: u32 = 64;

/// Hard cap on total instances buffered per frame. 4 shapes × 1024
/// MD2 max = 4 096; round up to 4 800 for headroom.
pub const MAX_CUSTOM_SHAPE_INSTANCES: usize = 4_800;

#[repr(C)]
#[derive(Debug, Clone, Copy, Pod, Zeroable)]
struct ShapeUniforms {
    aspect_pad: [f32; 4],
}

/// GPU pass. Owns the per-frame instance storage buffer + the four
/// pipelines (fill × alpha/additive, outline × alpha/additive). The
/// fragment shader picks textured vs. solid via a per-instance flag bit,
/// so we don't need a separate textured pipeline.
pub struct CustomShapeRenderer {
    fill_pipelines: [wgpu::RenderPipeline; 2],
    outline_pipelines: [wgpu::RenderPipeline; 2],
    bind_group: wgpu::BindGroup,
    bgl: wgpu::BindGroupLayout,
    instance_buffer: wgpu::Buffer,
    uniform_buffer: wgpu::Buffer,
    sampler: wgpu::Sampler,
    instance_count: u32,
    batches: Vec<CustomShapeBatch>,
}

#[inline]
fn additive_index(additive: bool) -> usize {
    additive as usize
}

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

        // Zero-initialised storage buffer. We always upload the
        // exact span the current frame uses, so any tail past
        // `instance_count` reads stale data — but no batch ever
        // points there.
        let initial = vec![CustomShapeInstance::default(); MAX_CUSTOM_SHAPE_INSTANCES];
        let instance_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
            label: Some("Custom Shape Instances"),
            contents: bytemuck::cast_slice(&initial),
            usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
        });

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

        let sampler = device.create_sampler(&wgpu::SamplerDescriptor {
            label: Some("Custom Shape Sampler"),
            address_mode_u: wgpu::AddressMode::Repeat,
            address_mode_v: wgpu::AddressMode::Repeat,
            mag_filter: wgpu::FilterMode::Linear,
            min_filter: wgpu::FilterMode::Linear,
            ..Default::default()
        });

        let bgl = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
            label: Some("Custom Shape BGL"),
            entries: &[
                wgpu::BindGroupLayoutEntry {
                    binding: 0,
                    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,
                },
                wgpu::BindGroupLayoutEntry {
                    binding: 1,
                    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: 2,
                    visibility: wgpu::ShaderStages::FRAGMENT,
                    ty: wgpu::BindingType::Texture {
                        sample_type: wgpu::TextureSampleType::Float { filterable: true },
                        view_dimension: wgpu::TextureViewDimension::D2,
                        multisampled: false,
                    },
                    count: None,
                },
                wgpu::BindGroupLayoutEntry {
                    binding: 3,
                    visibility: wgpu::ShaderStages::FRAGMENT,
                    ty: wgpu::BindingType::Sampler(wgpu::SamplerBindingType::Filtering),
                    count: None,
                },
            ],
        });

        let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some("Custom Shape BG"),
            layout: &bgl,
            entries: &[
                wgpu::BindGroupEntry {
                    binding: 0,
                    resource: instance_buffer.as_entire_binding(),
                },
                wgpu::BindGroupEntry {
                    binding: 1,
                    resource: uniform_buffer.as_entire_binding(),
                },
                wgpu::BindGroupEntry {
                    binding: 2,
                    resource: wgpu::BindingResource::TextureView(prev_texture_view),
                },
                wgpu::BindGroupEntry {
                    binding: 3,
                    resource: wgpu::BindingResource::Sampler(&sampler),
                },
            ],
        });

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

        let make_pipeline = |label: &str,
                             topology: wgpu::PrimitiveTopology,
                             vs_entry: &str,
                             fs_entry: &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_entry),
                    buffers: &[],
                    compilation_options: Default::default(),
                },
                fragment: Some(wgpu::FragmentState {
                    module: &shader,
                    entry_point: Some(fs_entry),
                    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 fill_pipelines = [
            make_pipeline(
                "Shape Fill Alpha",
                wgpu::PrimitiveTopology::TriangleList,
                "vs_main",
                "fs_main",
                false,
            ),
            make_pipeline(
                "Shape Fill Additive",
                wgpu::PrimitiveTopology::TriangleList,
                "vs_main",
                "fs_main",
                true,
            ),
        ];
        let outline_pipelines = [
            make_pipeline(
                "Shape Outline Alpha",
                wgpu::PrimitiveTopology::TriangleList,
                "vs_outline",
                "fs_outline",
                false,
            ),
            make_pipeline(
                "Shape Outline Additive",
                wgpu::PrimitiveTopology::TriangleList,
                "vs_outline",
                "fs_outline",
                true,
            ),
        ];

        Self {
            fill_pipelines,
            outline_pipelines,
            bind_group,
            bgl,
            instance_buffer,
            uniform_buffer,
            sampler,
            instance_count: 0,
            batches: Vec::new(),
        }
    }

    /// Re-bind `prev_texture_view` after a resize / texture realloc. The
    /// instance + uniform buffers persist, only the texture view
    /// changes.
    pub fn rebind_prev_texture(
        &mut self,
        device: &wgpu::Device,
        prev_texture_view: &wgpu::TextureView,
    ) {
        self.bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some("Custom Shape BG (rebind)"),
            layout: &self.bgl,
            entries: &[
                wgpu::BindGroupEntry {
                    binding: 0,
                    resource: self.instance_buffer.as_entire_binding(),
                },
                wgpu::BindGroupEntry {
                    binding: 1,
                    resource: self.uniform_buffer.as_entire_binding(),
                },
                wgpu::BindGroupEntry {
                    binding: 2,
                    resource: wgpu::BindingResource::TextureView(prev_texture_view),
                },
                wgpu::BindGroupEntry {
                    binding: 3,
                    resource: wgpu::BindingResource::Sampler(&self.sampler),
                },
            ],
        });
    }

    /// Replace the per-frame instance stream + batch list. Aspect ratio
    /// goes into the uniform buffer so the shader can correct the X
    /// radius without each instance carrying a copy.
    pub fn update(
        &mut self,
        queue: &wgpu::Queue,
        instances: &[CustomShapeInstance],
        batches: &[CustomShapeBatch],
        aspect: f32,
    ) {
        let u = ShapeUniforms {
            aspect_pad: [aspect.max(1e-3), 0.0, 0.0, 0.0],
        };
        queue.write_buffer(&self.uniform_buffer, 0, bytemuck::bytes_of(&u));

        let n = instances.len().min(MAX_CUSTOM_SHAPE_INSTANCES);
        if instances.len() > MAX_CUSTOM_SHAPE_INSTANCES {
            log::warn!(
                "custom-shape instance stream truncated: {} > cap {}",
                instances.len(),
                MAX_CUSTOM_SHAPE_INSTANCES
            );
        }
        if n > 0 {
            queue.write_buffer(
                &self.instance_buffer,
                0,
                bytemuck::cast_slice(&instances[..n]),
            );
        }
        self.instance_count = n as u32;
        self.batches.clear();
        for b in batches {
            if b.start_instance + b.instance_count <= self.instance_count {
                self.batches.push(*b);
            }
        }
    }

    /// Issue fill + (optional) outline draw calls for each batch. Loads
    /// existing contents (no clear); meant to overlay on the warp +
    /// wave output.
    pub fn render(&self, encoder: &mut wgpu::CommandEncoder, view: &wgpu::TextureView) {
        if self.batches.is_empty() || self.instance_count == 0 {
            return;
        }
        let mut rp = encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
            label: Some("Custom Shape 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_bind_group(0, &self.bind_group, &[]);

        for b in &self.batches {
            if b.instance_count == 0 {
                continue;
            }
            let sides = b.sides.clamp(3, MAX_SHAPE_SIDES);
            let fill_pipe = &self.fill_pipelines[additive_index(b.additive)];
            rp.set_pipeline(fill_pipe);
            let end_instance = b.start_instance + b.instance_count;
            rp.draw(0..(sides * 3), b.start_instance..end_instance);

            if b.thick_outline {
                let outline_pipe = &self.outline_pipelines[additive_index(b.additive)];
                rp.set_pipeline(outline_pipe);
                // TriangleList: 6 vertices per perimeter segment build a
                // screen-space quad of `OUTLINE_THICKNESS` perpendicular
                // to the segment. `sides * 6` total — `sides` segments
                // closing the perimeter (segment `sides - 1` wraps
                // vertex `sides - 1 → 0` via the shader's modulo).
                rp.draw(0..(sides * 6), b.start_instance..end_instance);
            }
        }
    }

    pub fn instance_count(&self) -> u32 {
        self.instance_count
    }

    pub fn batch_count(&self) -> usize {
        self.batches.len()
    }
}

/// Convert MD2 preset-space `(x, y)` (origin top-left, range `[0, 1]`)
/// to clip-space `(x', y')` (origin centre, range `[-1, 1]`, Y-up).
/// Same convention as `custom_wave::preset_xy_to_clip`.
#[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]
}

/// Build a `CustomShapeInstance` from MD2-space eval output. `rad` is
/// converted from preset units (roughly [0, 1] for full-screen) to
/// clip-space half-extent (radius 0.5 = quarter-screen-width before
/// aspect correction).
#[allow(clippy::too_many_arguments)]
pub fn instance_from_md2_state(
    x: f64,
    y: f64,
    rad: f64,
    ang: f64,
    tex_zoom: f64,
    tex_ang: f64,
    sides: u32,
    color_center: [f32; 4],
    color_edge: [f32; 4],
    border_color: [f32; 4],
    border_size: f32,
    flags: ShapeFlags,
) -> CustomShapeInstance {
    let [cx, cy] = preset_xy_to_clip(x, y);
    // MD2's `rad` is a screen-relative half-extent already. Map 0..1 →
    // 0..2 clip units (full screen) without further scaling.
    let r = (rad as f32) * 2.0;
    CustomShapeInstance {
        center_radius: [cx, cy, r, r],
        angle_zoom_ang_sides: [ang as f32, tex_zoom as f32, tex_ang as f32, sides as f32],
        color_center,
        color_edge,
        border_color,
        border_flags: [border_size, flags.pack(), 0.0, 0.0],
    }
}

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

    #[test]
    fn preset_xy_to_clip_matches_wave_convention() {
        // Same mapping as custom_wave::preset_xy_to_clip — both modules
        // own a copy so the renderer can be sliced apart later, but
        // they must agree on the convention.
        assert_eq!(
            preset_xy_to_clip(0.5, 0.5),
            crate::custom_wave::preset_xy_to_clip(0.5, 0.5)
        );
        assert_eq!(
            preset_xy_to_clip(0.0, 1.0),
            crate::custom_wave::preset_xy_to_clip(0.0, 1.0)
        );
    }

    #[test]
    fn instance_from_md2_state_packs_flags() {
        let inst = instance_from_md2_state(
            0.5,
            0.5,
            0.1,
            0.0,
            1.0,
            0.0,
            6,
            [1.0, 0.0, 0.0, 1.0],
            [0.0, 1.0, 0.0, 1.0],
            [1.0, 1.0, 1.0, 1.0],
            0.0,
            ShapeFlags::TEXTURED | ShapeFlags::ADDITIVE,
        );
        let flag_bits = inst.border_flags[1] as u32;
        assert!(flag_bits & ShapeFlags::TEXTURED.bits() != 0);
        assert!(flag_bits & ShapeFlags::ADDITIVE.bits() != 0);
        assert!(flag_bits & ShapeFlags::THICK_OUTLINE.bits() == 0);
        assert_eq!(inst.angle_zoom_ang_sides[3], 6.0);
    }

    #[test]
    fn shape_flags_round_trip_through_f32() {
        let f = ShapeFlags::THICK_OUTLINE.pack();
        let bits = f as u32;
        assert_eq!(bits, ShapeFlags::THICK_OUTLINE.bits());
    }

    #[test]
    fn additive_index_table() {
        assert_eq!(additive_index(false), 0);
        assert_eq!(additive_index(true), 1);
    }
}