wsg/lib/src/resources/uniform.rs

//! # Uniform Module — GPU Buffer Data Types
//!
//! Defines the CPU-side `Pod` (plain old data) structs that are uploaded to GPU uniform buffers.
//! Their memory layout must match **exactly** the WGSL uniforms declared in `standard_shader.wgsl`
//! (see the "Uniform Contract" section of that file) — 16-byte alignment (std140), no padding.
//!
//! Two bind groups are shared by every pipeline (single-layout decision, Étape 3) :
//! - `@group(0) @binding(0)` : `FrameUniforms` (per-frame : camera + lights) → 192 bytes
//! - `@group(1) @binding(0)` : `ObjectUniform` (per-entity model matrix) → 64 bytes
//!
//! ## Interaction with Other Modules
//! - `pipeline_cache::build_pipeline()` creates the two bind group layouts matching these types.
//! - `Renderer` allocates the buffers and `BindGroup`s from these types and writes them each frame.
//! - `standard_shader.wgsl` consumes them (layout identical to these structs).

use bytemuck::{Pod, Zeroable};
use glam::{Mat4, Vec4};

/// Byte size of the per-frame uniform buffer (`FrameUniforms`).
pub const FRAME_UNIFORMS_SIZE: u64 = std::mem::size_of::<FrameUniforms>() as u64;
/// Byte size of the per-object uniform buffer (`ObjectUniform`).
pub const OBJECT_UNIFORM_SIZE: u64 = std::mem::size_of::<ObjectUniform>() as u64;

/// Maximum number of lights stored in the per-frame uniform buffer.
/// Bounded capacity: adding more than this returns `WsgError` (no dynamic UBO allocation).
pub const MAX_LIGHTS: usize = 8;

/// A single light, stored in the per-frame uniform array. One struct serves both types; the
/// *position in the array* disambiguates: indices `0..num_directional` are directional
/// (`position_dir.xyz` = direction **from the surface toward the light**), indices
/// `num_directional..` are point (`position_dir.xyz` = world position). No type flag in the struct.
///
/// 3 × Vec4 = 48 bytes, 16-byte aligned (std140-compatible with the WGSL `struct Light`).
#[repr(C)]
#[derive(Clone, Copy, Pod, Zeroable, PartialEq)]
pub struct Light {
    /// xyz = direction from surface toward the light (directional) or world position (point); w = 0.
    pub position_dir: Vec4,
    /// rgb = color; a = intensity (multiplier).
    pub color: Vec4,
    /// x = attenuation radius (point lights); 0 for directional.
    pub radius: Vec4,
}

/// Per-frame GPU uniforms : camera matrices + ambient + global light list + options.
///
/// Mirrors the WGSL `FrameUniforms` struct in `standard_shader.wgsl` (offset table there).
/// 192 + 48·MAX_LIGHTS bytes (16-byte aligned) — `Pod` for direct `bytes_of` upload. The bind-group
/// layout uses `min_binding_size: None`, so extending this struct is transparent (no relayout).
#[repr(C)]
#[derive(Clone, Copy, Pod, Zeroable)]
pub struct FrameUniforms {
    /// Camera view matrix (world → view space). Offset 0.
    pub view: Mat4,
    /// Camera projection matrix (view → clip space). Offset 64.
    pub proj: Mat4,
    /// Camera world position (`.xyz` used). Offset 128.
    pub cam_pos: Vec4,
    /// Ambient hemisphere color (`.rgb` used). Offset 144.
    pub ambient: Vec4,
    /// Global light list: `0..num_directional` directional, then `num_point` point. Offset 160.
    pub lights: [Light; MAX_LIGHTS],
    /// Number of active directional lights (indices `0..num_directional`). Offset 160 + 48·MAX_LIGHTS.
    pub num_directional: u32,
    /// Number of active point lights (indices after the directionals).
    pub num_point: u32,
    /// Padding so `options` lands on a 16-byte boundary — matching the WGSL `vec4<u32>`
    /// (alignment 16), which Rust's `repr(C)` would otherwise place at offset 552.
    pub _pad: [u32; 2],
    /// Options. `options[0]` = unlit flag (1 → flat color, no lighting).
    pub options: [u32; 4],
}

impl Default for FrameUniforms {
    /// Sensible defaults : identity camera, white ambient, a single white directional light along
    /// +Z (from surface toward light), *lit* mode — reproduces the pre-multi-light look exactly.
    fn default() -> Self {
        Self {
            view: Mat4::IDENTITY,
            proj: Mat4::IDENTITY,
            cam_pos: Vec4::ZERO,
            ambient: Vec4::ONE,
            lights: [Light {
                position_dir: Vec4::new(0.0, 0.0, 1.0, 0.0), // from surface toward light = +Z
                color: Vec4::ONE,
                radius: Vec4::ZERO,
            }; MAX_LIGHTS],
            num_directional: 1,
            num_point: 0,
            _pad: [0, 0],
            options: [0, 0, 0, 0],
        }
    }
}

/// Per-object GPU uniforms : the entity's world-space model matrix.
///
/// Mirrors the WGSL `ObjectUniform` struct. 64 bytes, `Pod`.
#[repr(C)]
#[derive(Clone, Copy, Pod, Zeroable, Default)]
pub struct ObjectUniform {
    /// Model matrix (object → world space). Offset 0.
    pub model: Mat4,
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::mem::offset_of;
    use std::mem::{align_of, size_of};

    #[test]
    fn frame_uniforms_layout_matches_wgsl() {
        // The offsets below must match the offset table in standard_shader.wgsl.
        // Header (view..ambient) = 160, lights = 48·MAX_LIGHTS, then counters(8) + pad(8) +
        // options(16) = 32. Total = 160 + 48·MAX_LIGHTS + 32 = 576 bytes.
        assert_eq!(size_of::<FrameUniforms>(), 160 + 48 * MAX_LIGHTS + 32);
        assert_eq!(align_of::<FrameUniforms>(), 16);

        let f = FrameUniforms::default();
        assert_eq!(offset_of!(FrameUniforms, view), 0);
        assert_eq!(offset_of!(FrameUniforms, proj), 64);
        assert_eq!(offset_of!(FrameUniforms, cam_pos), 128);
        assert_eq!(offset_of!(FrameUniforms, ambient), 144);
        assert_eq!(offset_of!(FrameUniforms, lights), 160);
        assert_eq!(
            offset_of!(FrameUniforms, num_directional),
            160 + 48 * MAX_LIGHTS
        );
        assert_eq!(
            offset_of!(FrameUniforms, options),
            160 + 48 * MAX_LIGHTS + 16
        );
        // Default is lit mode (unlit flag cleared), one directional light, no point lights.
        assert_eq!(f.options[0], 0);
        assert_eq!(f.num_directional, 1);
        assert_eq!(f.num_point, 0);
    }

    #[test]
    fn object_uniform_layout_matches_wgsl() {
        assert_eq!(size_of::<ObjectUniform>(), 64);
        assert_eq!(align_of::<ObjectUniform>(), 16);
        assert_eq!(offset_of!(ObjectUniform, model), 0);
    }
}