onedrop_renderer/

sprite_pipeline.rs

//! Sprite (§5) GPU pipeline.
//!
//! Renders one textured quad per active sprite into `render_texture`,
//! between the custom-shape pass and the borders pass. Two pipelines
//! cover the two blend modes the engine emits ([`SpriteBlendKind`] —
//! alpha or additive). Each sprite draw owns its own bind group
//! recreated per frame so we can bind a different texture per draw
//! without an atlas (sprites in the wild are few, ~handful per
//! frame, so per-draw bind groups are cheap).
//!
//! The renderer drives this through [`SpritePipeline::record`], which
//! takes a slice of [`SpriteDrawCmd`] produced by
//! [`crate::sprite_pipeline::SpritePool::build_draw_commands`]. The
//! pool resolves the engine-side `texture_index` to a real
//! [`wgpu::TextureView`], applies aspect ratio correction, and packs
//! the per-sprite uniform.

use std::path::PathBuf;
use std::sync::Arc;

use crate::pipeline_helpers::{blend_state_for, load_wgsl};

/// Hard cap on simultaneously active sprites. The active list is
/// pruned every tick so this only matters during a key-mash spawn
/// burst; 64 is well above MD2's practical limit (~5-10).
pub const MAX_ACTIVE_SPRITES: usize = 64;

/// Which of the two GPU pipelines to use for a given sprite. Mirrors
/// `onedrop_engine::SpriteBlendMode`; kept independent so the
/// renderer compiles without an engine-side dep cycle.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SpriteBlendKind {
    Alpha,
    Additive,
}

/// Per-sprite uniform pushed to the GPU each draw. 48 bytes; the
/// pipeline allocates `MAX_ACTIVE_SPRITES * 256 B` total uniform
/// storage so each draw can dynamic-offset into its own slot.
#[repr(C)]
#[derive(Debug, Clone, Copy, Default)]
pub struct SpriteUniform {
    /// `(center_x, center_y, size_x, size_y)` in mixed units —
    /// centre is `[0, 1]` screen coords (top-left origin), size is
    /// half-width / half-height in clip-space units (the shader
    /// adds and subtracts them around the centre).
    center_size: [f32; 4],
    /// `(rot, _, _, _)` — three slots reserved for `flipx/flipy/_`.
    rot_pad: [f32; 4],
    /// Tint (`r, g, b, a`), premultiplied by `a` host-side.
    rgba: [f32; 4],
}

// SAFETY: `SpriteUniform` is `#[repr(C)]` with only POD members.
unsafe impl bytemuck::Pod for SpriteUniform {}
unsafe impl bytemuck::Zeroable for SpriteUniform {}

/// One per-frame sprite the engine asks the renderer to draw. POD —
/// no GPU resources, just numbers. The renderer resolves the
/// `texture_index` through its [`SpritePool`] and builds the
/// [`SpriteDrawCmd`] internally each frame.
#[derive(Debug, Clone, Copy)]
pub struct SpriteFrame {
    pub texture_index: u32,
    pub x: f32,
    pub y: f32,
    pub sx: f32,
    pub sy: f32,
    pub rot: f32,
    pub rgba: [f32; 4],
    pub blend: SpriteBlendKind,
    pub burn: bool,
}

/// One sprite ready to draw. The renderer assembles these from the
/// engine's [`SpriteFrame`] list + the sprite [`SpritePool`].
pub struct SpriteDrawCmd<'a> {
    pub uniform: SpriteUniform,
    pub texture_view: &'a wgpu::TextureView,
    pub blend: SpriteBlendKind,
    /// Burn-pending: caller treats it identically (one draw); the
    /// flag is forwarded for future fence/persistence hooks (none
    /// today — the quad already lands in `render_texture` which is
    /// the feedback target, so "burning into the background" is
    /// implicit).
    pub burn: bool,
}

/// One GPU-resident sprite texture.
pub struct SpriteTexture {
    pub texture: wgpu::Texture,
    pub view: wgpu::TextureView,
    pub width: u32,
    pub height: u32,
    /// Name the engine resolves against — typically the canonical
    /// stem of the source filename.
    pub name: String,
}

/// Sprite texture pool: one texture per sprite def, in load order.
/// Engine-side `SpriteRenderInstance::texture_index` indexes into
/// `textures` directly. Missing / undecodable files get a 1×1
/// transparent fallback so the bind group always populates.
pub struct SpritePool {
    textures: Vec<Arc<SpriteTexture>>,
    fallback: Arc<SpriteTexture>,
    sampler: wgpu::Sampler,
}

impl SpritePool {
    /// Build an empty pool with just the transparent fallback. Used
    /// by tests and as the renderer default when no `MILK_IMG.INI`
    /// has been loaded yet.
    pub fn new(device: &wgpu::Device, queue: &wgpu::Queue) -> Self {
        let fallback = Arc::new(make_fallback(device, queue));
        let sampler = device.create_sampler(&wgpu::SamplerDescriptor {
            label: Some("Sprite Sampler"),
            address_mode_u: wgpu::AddressMode::ClampToEdge,
            address_mode_v: wgpu::AddressMode::ClampToEdge,
            mag_filter: wgpu::FilterMode::Linear,
            min_filter: wgpu::FilterMode::Linear,
            ..Default::default()
        });
        Self {
            textures: Vec::new(),
            fallback,
            sampler,
        }
    }

    /// Borrow the GPU sampler for this pool. The sprite pipeline binds
    /// it identically across every draw.
    pub fn sampler(&self) -> &wgpu::Sampler {
        &self.sampler
    }

    /// Number of loaded sprite textures.
    pub fn len(&self) -> usize {
        self.textures.len()
    }

    pub fn is_empty(&self) -> bool {
        self.textures.is_empty()
    }

    /// Resolve a sprite by index. Returns the fallback for any
    /// out-of-range or missing entry — keeps the per-frame draw path
    /// branchless on the renderer side.
    pub fn get_or_fallback(&self, idx: u32) -> &Arc<SpriteTexture> {
        self.textures.get(idx as usize).unwrap_or(&self.fallback)
    }

    /// Reload the pool from a directory of image files, indexed by
    /// the engine's `SpriteDef` list. Missing files / decode failures
    /// emit a single warn log and substitute the fallback so the
    /// engine's `texture_index` mapping stays dense.
    ///
    /// `def_imgs` is the parallel list of `img=` filenames from the
    /// `MILK_IMG.INI`; `dir` is the directory those filenames resolve
    /// against (e.g. `~/.local/share/onedrop/sprites`).
    pub fn load_from_defs(
        &mut self,
        device: &wgpu::Device,
        queue: &wgpu::Queue,
        dir: &std::path::Path,
        def_imgs: &[String],
    ) {
        self.textures.clear();
        for img in def_imgs {
            let path = dir.join(img);
            match load_sprite_file(device, queue, &path) {
                Ok(tex) => self.textures.push(Arc::new(tex)),
                Err(e) => {
                    log::warn!(
                        "sprite texture {} unloadable ({}); using fallback",
                        path.display(),
                        e
                    );
                    self.textures.push(Arc::clone(&self.fallback));
                }
            }
        }
    }
}

/// Default search paths for sprite assets. The engine extends this
/// with config-driven dirs.
pub fn default_sprite_dirs() -> Vec<PathBuf> {
    let mut out = Vec::new();
    if let Some(home) = std::env::var_os("HOME") {
        let home = PathBuf::from(home);
        out.push(home.join(".local/share/onedrop/sprites"));
        out.push(home.join(".config/onedrop/sprites"));
        out.push(home.join("Music/milkdrop/sprites"));
    }
    if let Some(xdg) = std::env::var_os("XDG_DATA_HOME") {
        out.push(PathBuf::from(xdg).join("onedrop/sprites"));
    }
    out
}

/// Pick the first existing directory from a search-path list.
pub fn pick_first_existing(dirs: &[PathBuf]) -> Option<PathBuf> {
    dirs.iter().find(|p| p.is_dir()).cloned()
}

/// Sprite GPU pipeline. Owns two `RenderPipeline`s (alpha + additive)
/// and a per-slot uniform buffer; the renderer constructs / drives it
/// once per chain.
pub struct SpritePipeline {
    pipeline_alpha: wgpu::RenderPipeline,
    pipeline_additive: wgpu::RenderPipeline,
    bgl: wgpu::BindGroupLayout,
    /// `MAX_ACTIVE_SPRITES * UNIFORM_STRIDE` bytes; each draw writes
    /// at its slot's offset. UNIFORM_STRIDE is rounded up to the
    /// device's min uniform-buffer offset alignment so dynamic
    /// offsets land on a valid boundary.
    uniform_buffer: wgpu::Buffer,
    uniform_stride: u64,
}

impl SpritePipeline {
    pub fn new(device: &wgpu::Device, format: wgpu::TextureFormat) -> Self {
        let shader = load_wgsl(
            device,
            "Sprite Shader",
            include_str!("../shaders/sprite.wgsl"),
        );

        let bgl = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
            label: Some("Sprite BGL"),
            entries: &[
                wgpu::BindGroupLayoutEntry {
                    binding: 0,
                    visibility: wgpu::ShaderStages::VERTEX_FRAGMENT,
                    ty: wgpu::BindingType::Buffer {
                        ty: wgpu::BufferBindingType::Uniform,
                        has_dynamic_offset: true,
                        min_binding_size: wgpu::BufferSize::new(
                            std::mem::size_of::<SpriteUniform>() as u64,
                        ),
                    },
                    count: None,
                },
                wgpu::BindGroupLayoutEntry {
                    binding: 1,
                    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: 2,
                    visibility: wgpu::ShaderStages::FRAGMENT,
                    ty: wgpu::BindingType::Sampler(wgpu::SamplerBindingType::Filtering),
                    count: None,
                },
            ],
        });

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

        let make_pipeline = |label: &str, additive: bool| -> wgpu::RenderPipeline {
            let blend = 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 pipeline_alpha = make_pipeline("Sprite Pipeline (alpha)", false);
        let pipeline_additive = make_pipeline("Sprite Pipeline (additive)", true);

        // Round each slot to the device's min uniform-buffer alignment
        // (typically 256 bytes; we hard-code the conservative number
        // because querying `Limits` from the pipeline is awkward).
        let uniform_stride = 256u64;
        let uniform_buffer = device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("Sprite Uniforms"),
            size: uniform_stride * MAX_ACTIVE_SPRITES as u64,
            usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });

        Self {
            pipeline_alpha,
            pipeline_additive,
            bgl,
            uniform_buffer,
            uniform_stride,
        }
    }

    /// Record one sprite-pass draw set into the encoder. `cmds`
    /// must be `<= MAX_ACTIVE_SPRITES`; longer slices are truncated
    /// to keep the per-frame allocation budget predictable.
    pub fn record(
        &self,
        encoder: &mut wgpu::CommandEncoder,
        queue: &wgpu::Queue,
        device: &wgpu::Device,
        target: &wgpu::TextureView,
        sampler: &wgpu::Sampler,
        cmds: &[SpriteDrawCmd<'_>],
    ) {
        if cmds.is_empty() {
            return;
        }
        let n = cmds.len().min(MAX_ACTIVE_SPRITES);
        // Stage uniforms into one contiguous write so the GPU only
        // sees one CPU→GPU copy. Stride = 256, so the per-slot data
        // occupies the first `sizeof(SpriteUniform)` bytes and the
        // padding stays zero.
        let mut staging = vec![0u8; self.uniform_stride as usize * n];
        for (i, cmd) in cmds.iter().take(n).enumerate() {
            let start = i * self.uniform_stride as usize;
            let bytes = bytemuck::bytes_of(&cmd.uniform);
            staging[start..start + bytes.len()].copy_from_slice(bytes);
        }
        queue.write_buffer(&self.uniform_buffer, 0, &staging);

        // Build a separate bind group per draw — different textures
        // mean we can't share. Allocations are scoped to this frame
        // (the bind groups drop when `cmds` goes out of scope), and
        // wgpu's BG creation is cheap.
        let mut bind_groups: Vec<wgpu::BindGroup> = Vec::with_capacity(n);
        for cmd in cmds.iter().take(n) {
            bind_groups.push(
                device.create_bind_group(&wgpu::BindGroupDescriptor {
                    label: Some("Sprite BG (per-draw)"),
                    layout: &self.bgl,
                    entries: &[
                        wgpu::BindGroupEntry {
                            binding: 0,
                            resource: wgpu::BindingResource::Buffer(wgpu::BufferBinding {
                                buffer: &self.uniform_buffer,
                                offset: 0,
                                size: wgpu::BufferSize::new(
                                    std::mem::size_of::<SpriteUniform>() as u64
                                ),
                            }),
                        },
                        wgpu::BindGroupEntry {
                            binding: 1,
                            resource: wgpu::BindingResource::TextureView(cmd.texture_view),
                        },
                        wgpu::BindGroupEntry {
                            binding: 2,
                            resource: wgpu::BindingResource::Sampler(sampler),
                        },
                    ],
                }),
            );
        }

        let mut pass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
            label: Some("Sprite Pass"),
            color_attachments: &[Some(wgpu::RenderPassColorAttachment {
                view: target,
                resolve_target: None,
                ops: wgpu::Operations {
                    load: wgpu::LoadOp::Load,
                    store: wgpu::StoreOp::Store,
                },
                depth_slice: None,
            })],
            depth_stencil_attachment: None,
            timestamp_writes: None,
            occlusion_query_set: None,
            multiview_mask: None,
        });

        for (i, cmd) in cmds.iter().take(n).enumerate() {
            let pipeline = match cmd.blend {
                SpriteBlendKind::Alpha => &self.pipeline_alpha,
                SpriteBlendKind::Additive => &self.pipeline_additive,
            };
            pass.set_pipeline(pipeline);
            pass.set_bind_group(
                0,
                Some(&bind_groups[i]),
                &[i as u32 * self.uniform_stride as u32],
            );
            pass.draw(0..6, 0..1);
        }
    }
}

/// Build a [`SpriteUniform`] from raw engine-side numbers + texture
/// dimensions + render aspect ratio. Pulled out so the renderer's
/// `update_sprites` step can produce a single contiguous list of
/// uniforms for upload + the matching draw-cmd list without
/// recomputing aspect ratios per call.
#[allow(clippy::too_many_arguments)]
pub fn build_sprite_uniform(
    x: f32,
    y: f32,
    sx: f32,
    sy: f32,
    rot: f32,
    rgba: [f32; 4],
    texture_w: u32,
    texture_h: u32,
    render_w: u32,
    render_h: u32,
) -> SpriteUniform {
    // MD2 sizes: `sx = sy = 1` should make the sprite fill ~1/3 of
    // the shorter render axis, matching the historical default. We
    // pick `0.33` as the half-extent (so a full sprite occupies
    // 0.66 of the axis); preset-side scales scale around that.
    let base = 0.33;
    let tex_aspect = (texture_w.max(1) as f32) / (texture_h.max(1) as f32);
    let render_aspect = (render_w.max(1) as f32) / (render_h.max(1) as f32);
    // Quad half-size in clip space, baked aspect-correct.
    let size_x = sx * base * tex_aspect / render_aspect;
    let size_y = sy * base;
    SpriteUniform {
        center_size: [x, y, size_x, size_y],
        rot_pad: [rot, 0.0, 0.0, 0.0],
        rgba,
    }
}

fn make_fallback(device: &wgpu::Device, queue: &wgpu::Queue) -> SpriteTexture {
    upload_sprite_rgba8(device, queue, "__fallback_transparent", 1, 1, &[0, 0, 0, 0])
}

fn load_sprite_file(
    device: &wgpu::Device,
    queue: &wgpu::Queue,
    path: &std::path::Path,
) -> Result<SpriteTexture, image::ImageError> {
    let img = image::open(path)?;
    let rgba = img.to_rgba8();
    let (w, h) = rgba.dimensions();
    let name = path
        .file_stem()
        .and_then(|s| s.to_str())
        .map(|s| s.to_ascii_lowercase())
        .unwrap_or_default();
    Ok(upload_sprite_rgba8(device, queue, &name, w, h, &rgba))
}

fn upload_sprite_rgba8(
    device: &wgpu::Device,
    queue: &wgpu::Queue,
    name: &str,
    width: u32,
    height: u32,
    bytes: &[u8],
) -> SpriteTexture {
    let texture = device.create_texture(&wgpu::TextureDescriptor {
        label: Some(&format!("Sprite {name}")),
        size: wgpu::Extent3d {
            width,
            height,
            depth_or_array_layers: 1,
        },
        mip_level_count: 1,
        sample_count: 1,
        dimension: wgpu::TextureDimension::D2,
        format: wgpu::TextureFormat::Rgba8Unorm,
        usage: wgpu::TextureUsages::TEXTURE_BINDING | wgpu::TextureUsages::COPY_DST,
        view_formats: &[],
    });
    queue.write_texture(
        wgpu::TexelCopyTextureInfo {
            texture: &texture,
            mip_level: 0,
            origin: wgpu::Origin3d::ZERO,
            aspect: wgpu::TextureAspect::All,
        },
        bytes,
        wgpu::TexelCopyBufferLayout {
            offset: 0,
            bytes_per_row: Some(width * 4),
            rows_per_image: Some(height),
        },
        wgpu::Extent3d {
            width,
            height,
            depth_or_array_layers: 1,
        },
    );
    let view = texture.create_view(&wgpu::TextureViewDescriptor::default());
    SpriteTexture {
        texture,
        view,
        width,
        height,
        name: name.to_string(),
    }
}

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

    #[test]
    fn build_uniform_packs_aspect_correctly() {
        // 100×50 sprite (2:1) on a 1000×500 render (2:1) at sx=sy=1
        // should land with size_x == size_y * (2/1) / (2/1) == size_y.
        let u = build_sprite_uniform(0.5, 0.5, 1.0, 1.0, 0.0, [1.0; 4], 100, 50, 1000, 500);
        assert!((u.center_size[2] - u.center_size[3]).abs() < 1e-6);
    }

    #[test]
    fn build_uniform_handles_zero_dims() {
        // Degenerate texture/render dims must not blow up `f32::INFINITY`.
        let u = build_sprite_uniform(0.5, 0.5, 1.0, 1.0, 0.0, [1.0; 4], 0, 0, 0, 0);
        assert!(u.center_size[2].is_finite());
        assert!(u.center_size[3].is_finite());
    }

    #[test]
    fn build_uniform_writes_rot_to_first_lane() {
        let u = build_sprite_uniform(0.5, 0.5, 1.0, 1.0, 0.7, [1.0; 4], 100, 100, 100, 100);
        assert!((u.rot_pad[0] - 0.7).abs() < 1e-6);
    }
}