dear-imgui-wgpu 0.16.0-alpha.2

//! Uniform buffer management for the WGPU renderer
//!
//! This module handles the uniform data structure and buffer management,
//! corresponding to the Uniforms struct in imgui_impl_wgpu.cpp

use bytemuck::{Pod, Zeroable};
use wgpu::*;

/// Memory alignment function (equivalent to MEMALIGN macro in C++)
/// Aligns size to the specified alignment boundary
fn align_size(size: usize, alignment: usize) -> usize {
    (size + alignment - 1) & !(alignment - 1)
}

/// Uniform data structure
///
/// This corresponds to the Uniforms struct in the C++ implementation.
/// Contains the MVP matrix and gamma correction value.
#[repr(C)]
#[derive(Debug, Copy, Clone, Pod, Zeroable)]
pub struct Uniforms {
    /// Model-View-Projection matrix (4x4 f32 matrix)
    pub mvp: [[f32; 4]; 4],
    /// Gamma correction value
    pub gamma: f32,
    /// Padding to ensure proper alignment
    pub _padding: [f32; 3],
}

impl Uniforms {
    /// Create new uniforms with identity matrix and default gamma
    pub fn new() -> Self {
        Self {
            mvp: [
                [1.0, 0.0, 0.0, 0.0],
                [0.0, 1.0, 0.0, 0.0],
                [0.0, 0.0, 1.0, 0.0],
                [0.0, 0.0, 0.0, 1.0],
            ],
            gamma: 1.0,
            _padding: [0.0; 3],
        }
    }

    /// Create orthographic projection matrix for Dear ImGui
    ///
    /// This matches the matrix calculation in ImGui_ImplWGPU_SetupRenderState
    pub fn create_orthographic_matrix(
        display_pos: [f32; 2],
        display_size: [f32; 2],
    ) -> [[f32; 4]; 4] {
        let l = display_pos[0];
        let r = display_pos[0] + display_size[0];
        let t = display_pos[1];
        let b = display_pos[1] + display_size[1];

        [
            [2.0 / (r - l), 0.0, 0.0, 0.0],
            [0.0, 2.0 / (t - b), 0.0, 0.0],
            [0.0, 0.0, 0.5, 0.0],
            [(r + l) / (l - r), (t + b) / (b - t), 0.5, 1.0],
        ]
    }

    /// Determine gamma value based on texture format
    ///
    /// This matches the gamma detection logic in ImGui_ImplWGPU_SetupRenderState
    /// from the official C++ implementation, supporting all sRGB texture formats.
    pub fn gamma_for_format(format: TextureFormat) -> f32 {
        match format {
            // sRGB formats need gamma correction (gamma = 2.2)
            // This matches the complete list from imgui_impl_wgpu.cpp

            // ASTC sRGB formats
            TextureFormat::Astc {
                block: AstcBlock::B4x4,
                channel: AstcChannel::UnormSrgb,
            }
            | TextureFormat::Astc {
                block: AstcBlock::B5x4,
                channel: AstcChannel::UnormSrgb,
            }
            | TextureFormat::Astc {
                block: AstcBlock::B5x5,
                channel: AstcChannel::UnormSrgb,
            }
            | TextureFormat::Astc {
                block: AstcBlock::B6x5,
                channel: AstcChannel::UnormSrgb,
            }
            | TextureFormat::Astc {
                block: AstcBlock::B6x6,
                channel: AstcChannel::UnormSrgb,
            }
            | TextureFormat::Astc {
                block: AstcBlock::B8x5,
                channel: AstcChannel::UnormSrgb,
            }
            | TextureFormat::Astc {
                block: AstcBlock::B8x6,
                channel: AstcChannel::UnormSrgb,
            }
            | TextureFormat::Astc {
                block: AstcBlock::B8x8,
                channel: AstcChannel::UnormSrgb,
            }
            | TextureFormat::Astc {
                block: AstcBlock::B10x5,
                channel: AstcChannel::UnormSrgb,
            }
            | TextureFormat::Astc {
                block: AstcBlock::B10x6,
                channel: AstcChannel::UnormSrgb,
            }
            | TextureFormat::Astc {
                block: AstcBlock::B10x8,
                channel: AstcChannel::UnormSrgb,
            }
            | TextureFormat::Astc {
                block: AstcBlock::B10x10,
                channel: AstcChannel::UnormSrgb,
            }
            | TextureFormat::Astc {
                block: AstcBlock::B12x10,
                channel: AstcChannel::UnormSrgb,
            }
            | TextureFormat::Astc {
                block: AstcBlock::B12x12,
                channel: AstcChannel::UnormSrgb,
            }
            // BC (Block Compression) sRGB formats
            | TextureFormat::Bc1RgbaUnormSrgb
            | TextureFormat::Bc2RgbaUnormSrgb
            | TextureFormat::Bc3RgbaUnormSrgb
            | TextureFormat::Bc7RgbaUnormSrgb
            // Standard sRGB formats
            | TextureFormat::Rgba8UnormSrgb
            | TextureFormat::Bgra8UnormSrgb
            // ETC2 sRGB formats
            | TextureFormat::Etc2Rgb8UnormSrgb
            | TextureFormat::Etc2Rgb8A1UnormSrgb
            | TextureFormat::Etc2Rgba8UnormSrgb => 2.2,
            // Linear formats don't need gamma correction (gamma = 1.0)
            _ => 1.0,
        }
    }

    /// Update both MVP and gamma
    pub fn update(&mut self, mvp: [[f32; 4]; 4], gamma: f32) {
        self.mvp = mvp;
        self.gamma = gamma;
    }
}

impl Default for Uniforms {
    fn default() -> Self {
        Self::new()
    }
}

/// Uniform buffer manager
pub struct UniformBuffer {
    buffer: Buffer,
    bind_group: BindGroup,
}

impl UniformBuffer {
    pub(crate) fn create_bind_group_layout(device: &Device) -> BindGroupLayout {
        device.create_bind_group_layout(&BindGroupLayoutDescriptor {
            label: Some("Dear ImGui Common Bind Group Layout"),
            entries: &[
                BindGroupLayoutEntry {
                    binding: 0,
                    visibility: ShaderStages::VERTEX | ShaderStages::FRAGMENT,
                    ty: BindingType::Buffer {
                        ty: BufferBindingType::Uniform,
                        has_dynamic_offset: false,
                        min_binding_size: None,
                    },
                    count: None,
                },
                BindGroupLayoutEntry {
                    binding: 1,
                    visibility: ShaderStages::FRAGMENT,
                    ty: BindingType::Sampler(SamplerBindingType::Filtering),
                    count: None,
                },
            ],
        })
    }

    pub(crate) fn new_with_layout(
        device: &Device,
        sampler: &Sampler,
        bind_group_layout: &BindGroupLayout,
    ) -> Self {
        let buffer_size = align_size(std::mem::size_of::<Uniforms>(), 16);
        let buffer = device.create_buffer(&BufferDescriptor {
            label: Some("Dear ImGui Uniform Buffer"),
            size: buffer_size as u64,
            usage: BufferUsages::UNIFORM | BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });
        let bind_group = device.create_bind_group(&BindGroupDescriptor {
            label: Some("Dear ImGui Common Bind Group"),
            layout: bind_group_layout,
            entries: &[
                BindGroupEntry {
                    binding: 0,
                    resource: buffer.as_entire_binding(),
                },
                BindGroupEntry {
                    binding: 1,
                    resource: BindingResource::Sampler(sampler),
                },
            ],
        });

        Self { buffer, bind_group }
    }

    /// Update the uniform buffer with new data
    pub fn update(&self, queue: &Queue, uniforms: &Uniforms) {
        queue.write_buffer(&self.buffer, 0, bytemuck::bytes_of(uniforms));
    }

    /// Get the bind group for rendering
    pub fn bind_group(&self) -> &BindGroup {
        &self.bind_group
    }

    /// Get the buffer
    pub fn buffer(&self) -> &Buffer {
        &self.buffer
    }
}

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

    #[test]
    fn uniform_layout_matches_the_shader_contract() {
        assert_eq!(size_of::<Uniforms>(), 80);
        assert_eq!(align_of::<Uniforms>(), 4);

        let uniforms = Uniforms::new();
        assert_eq!(uniforms.gamma, 1.0);
        assert_eq!(
            uniforms.mvp,
            [
                [1.0, 0.0, 0.0, 0.0],
                [0.0, 1.0, 0.0, 0.0],
                [0.0, 0.0, 1.0, 0.0],
                [0.0, 0.0, 0.0, 1.0],
            ]
        );
    }

    #[test]
    fn gamma_policy_distinguishes_srgb_and_linear_targets() {
        for format in [
            TextureFormat::Rgba8UnormSrgb,
            TextureFormat::Bgra8UnormSrgb,
            TextureFormat::Bc7RgbaUnormSrgb,
            TextureFormat::Etc2Rgba8UnormSrgb,
            TextureFormat::Astc {
                block: AstcBlock::B12x12,
                channel: AstcChannel::UnormSrgb,
            },
        ] {
            assert_eq!(Uniforms::gamma_for_format(format), 2.2);
        }
        for format in [
            TextureFormat::Rgba8Unorm,
            TextureFormat::Bgra8Unorm,
            TextureFormat::R8Unorm,
            TextureFormat::Rg8Unorm,
        ] {
            assert_eq!(Uniforms::gamma_for_format(format), 1.0);
        }
    }

    #[test]
    fn orthographic_matrix_matches_imgui_wgpu_projection() {
        let display_pos = [10.0, 20.0];
        let display_size = [800.0, 600.0];
        let matrix = Uniforms::create_orthographic_matrix(display_pos, display_size);
        let [left, top] = display_pos;
        let right = left + display_size[0];
        let bottom = top + display_size[1];

        assert_eq!(matrix[0][0], 2.0 / (right - left));
        assert_eq!(matrix[1][1], 2.0 / (top - bottom));
        assert_eq!(matrix[2][2], 0.5);
        assert_eq!(matrix[3][0], (right + left) / (left - right));
        assert_eq!(matrix[3][1], (top + bottom) / (bottom - top));
        assert_eq!(matrix[3][2], 0.5);
        assert_eq!(matrix[3][3], 1.0);
    }
}