roast2d_internal 0.4.0

//! GPU-instanced mesh shader for efficient rendering of many copies.

use std::borrow::Cow;

use glam::Mat4;

use crate::color::Color;
use crate::engine::Engine;
use crate::platform::types::TextureResource;

use super::mesh::Mesh3D;
use super::shader_standard::Mesh3DShader;
use super::vertex::{InstanceData, Mesh3DVertex};

/// Per-frame uniform data for instanced rendering (shared across all instances)
#[repr(C)]
#[derive(Clone, Copy, bytemuck::Pod, bytemuck::Zeroable, Debug)]
struct InstancedFrameUniform {
    view_proj: [[f32; 4]; 4],
    camera_pos: [f32; 4],
    // Lighting
    ambient: [f32; 4],
    light_dir: [f32; 4], // direction for directional/spot lights
    light_color: [f32; 4],
    light_pos: [f32; 4], // position for point/spot, w = light type (0=dir, 1=point, 2=spot)
    light_params: [f32; 4], // x=radius, y=inner_angle, z=outer_angle, w=intensity
    // Material (shared by all instances)
    diffuse: [f32; 4],
    specular: [f32; 4],
    emissive: [f32; 4],
    shininess: f32,
    has_texture: f32,
    _padding: [f32; 2],
}

/// Maximum number of instances per draw call
const MAX_INSTANCES: usize = 4096;

/// Shader for instanced 3D mesh rendering.
///
/// Uses GPU instancing to efficiently render many copies of the same mesh
/// with different transforms and colors in a single draw call.
pub struct Mesh3DInstancedShader {
    pipeline: wgpu::RenderPipeline,
    frame_uniform: wgpu::Buffer,
    frame_bind_group: wgpu::BindGroup,
    texture_bind_group_layout: wgpu::BindGroupLayout,
    #[allow(dead_code)]
    default_texture: TextureResource,
    default_texture_bind_group: wgpu::BindGroup,
    /// Instance data buffer
    instance_buffer: wgpu::Buffer,
    instance_buffer_capacity: usize,
}

impl Mesh3DInstancedShader {
    /// Create a new instanced mesh shader.
    pub fn new(g: &Engine) -> Self {
        let state = g.backend_state();
        let device = &state.device;
        let format = state.surface_view_format;

        let shader = device.create_shader_module(wgpu::ShaderModuleDescriptor {
            label: Some("mesh3d_instanced.shader"),
            source: wgpu::ShaderSource::Wgsl(Cow::Borrowed(include_str!(
                "../../assets/shaders/mesh3d_instanced.wgsl"
            ))),
        });

        // Frame uniform buffer (single block, no dynamic offset)
        let frame_uniform = device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("mesh3d_instanced.frame_uniform"),
            size: std::mem::size_of::<InstancedFrameUniform>() as u64,
            usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });

        let frame_bind_group_layout =
            device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
                label: Some("mesh3d_instanced.frame_bgl"),
                entries: &[wgpu::BindGroupLayoutEntry {
                    binding: 0,
                    visibility: wgpu::ShaderStages::VERTEX | wgpu::ShaderStages::FRAGMENT,
                    ty: wgpu::BindingType::Buffer {
                        ty: wgpu::BufferBindingType::Uniform,
                        has_dynamic_offset: false,
                        min_binding_size: wgpu::BufferSize::new(std::mem::size_of::<
                            InstancedFrameUniform,
                        >() as u64),
                    },
                    count: None,
                }],
            });

        let frame_bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some("mesh3d_instanced.frame_bg"),
            layout: &frame_bind_group_layout,
            entries: &[wgpu::BindGroupEntry {
                binding: 0,
                resource: frame_uniform.as_entire_binding(),
            }],
        });

        // Texture bind group layout (same as non-instanced)
        let texture_bind_group_layout =
            device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
                label: Some("mesh3d_instanced.texture_bgl"),
                entries: &[
                    wgpu::BindGroupLayoutEntry {
                        binding: 0,
                        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: 1,
                        visibility: wgpu::ShaderStages::FRAGMENT,
                        ty: wgpu::BindingType::Sampler(wgpu::SamplerBindingType::Filtering),
                        count: None,
                    },
                ],
            });

        // Default texture
        let default_texture = Mesh3DShader::create_default_texture(device, &state.queue);
        let default_texture_bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some("mesh3d_instanced.default_texture_bg"),
            layout: &texture_bind_group_layout,
            entries: &[
                wgpu::BindGroupEntry {
                    binding: 0,
                    resource: wgpu::BindingResource::TextureView(&default_texture.view),
                },
                wgpu::BindGroupEntry {
                    binding: 1,
                    resource: wgpu::BindingResource::Sampler(&default_texture.sampler),
                },
            ],
        });

        let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
            label: Some("mesh3d_instanced.pl"),
            bind_group_layouts: &[&frame_bind_group_layout, &texture_bind_group_layout],
            push_constant_ranges: &[],
        });

        // Instance buffer layout
        let instance_buffer_layout = wgpu::VertexBufferLayout {
            array_stride: std::mem::size_of::<InstanceData>() as wgpu::BufferAddress,
            step_mode: wgpu::VertexStepMode::Instance,
            attributes: &[
                // model matrix column 0
                wgpu::VertexAttribute {
                    format: wgpu::VertexFormat::Float32x4,
                    offset: 0,
                    shader_location: 3,
                },
                // model matrix column 1
                wgpu::VertexAttribute {
                    format: wgpu::VertexFormat::Float32x4,
                    offset: 16,
                    shader_location: 4,
                },
                // model matrix column 2
                wgpu::VertexAttribute {
                    format: wgpu::VertexFormat::Float32x4,
                    offset: 32,
                    shader_location: 5,
                },
                // model matrix column 3
                wgpu::VertexAttribute {
                    format: wgpu::VertexFormat::Float32x4,
                    offset: 48,
                    shader_location: 6,
                },
                // color
                wgpu::VertexAttribute {
                    format: wgpu::VertexFormat::Float32x4,
                    offset: 64,
                    shader_location: 7,
                },
            ],
        };

        let pipeline = device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
            label: Some("mesh3d_instanced.pipeline"),
            layout: Some(&pipeline_layout),
            cache: None,
            vertex: wgpu::VertexState {
                module: &shader,
                entry_point: Some("vs_main"),
                buffers: &[
                    // Vertex buffer (per-vertex)
                    wgpu::VertexBufferLayout {
                        array_stride: std::mem::size_of::<Mesh3DVertex>() as wgpu::BufferAddress,
                        step_mode: wgpu::VertexStepMode::Vertex,
                        attributes: &wgpu::vertex_attr_array![0 => Float32x3, 1 => Float32x3, 2 => Float32x2],
                    },
                    // Instance buffer (per-instance)
                    instance_buffer_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::TriangleList,
                cull_mode: Some(wgpu::Face::Back),
                front_face: wgpu::FrontFace::Ccw,
                ..Default::default()
            },
            depth_stencil: Some(wgpu::DepthStencilState {
                format: wgpu::TextureFormat::Depth32Float,
                depth_write_enabled: true,
                depth_compare: wgpu::CompareFunction::Greater,
                stencil: wgpu::StencilState::default(),
                bias: wgpu::DepthBiasState::default(),
            }),
            multisample: wgpu::MultisampleState::default(),
            multiview: None,
        });

        // Initial instance buffer
        let instance_buffer_capacity = MAX_INSTANCES;
        let instance_buffer = device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("mesh3d_instanced.instance_buffer"),
            size: (std::mem::size_of::<InstanceData>() * instance_buffer_capacity) as u64,
            usage: wgpu::BufferUsages::VERTEX | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });

        Self {
            pipeline,
            frame_uniform,
            frame_bind_group,
            texture_bind_group_layout,
            default_texture,
            default_texture_bind_group,
            instance_buffer,
            instance_buffer_capacity,
        }
    }

    /// Create a texture bind group for a given texture resource.
    fn create_texture_bind_group(
        &self,
        device: &wgpu::Device,
        texture: &TextureResource,
    ) -> wgpu::BindGroup {
        device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some("mesh3d_instanced.texture_bg"),
            layout: &self.texture_bind_group_layout,
            entries: &[
                wgpu::BindGroupEntry {
                    binding: 0,
                    resource: wgpu::BindingResource::TextureView(&texture.view),
                },
                wgpu::BindGroupEntry {
                    binding: 1,
                    resource: wgpu::BindingResource::Sampler(&texture.sampler),
                },
            ],
        })
    }

    /// Record instanced draw commands into the encoder.
    #[allow(clippy::type_complexity)]
    #[allow(clippy::too_many_arguments)]
    pub fn record(
        &mut self,
        g: &Engine,
        color_target: &wgpu::Texture,
        depth_view: &wgpu::TextureView,
        view_proj: Mat4,
        draws: &[(
            &Mesh3D,
            &[InstanceData],
            Option<&TextureResource>,
            &crate::light3d::Material3D,
        )],
        encoder: &mut wgpu::CommandEncoder,
        should_clear: bool,
    ) {
        if draws.is_empty() {
            return;
        }

        let state = g.backend_state();
        let lighting = g.lighting();
        let camera = g.camera3d();

        // Get the first enabled light of any type
        let (light_dir, light_color, light_pos, light_params) = lighting
            .lights
            .iter()
            .flatten()
            .find(|l| l.enabled)
            .map(|l| {
                use crate::light3d::LightKind;
                let (light_type, radius, inner_angle, outer_angle) = match l.kind {
                    LightKind::Directional => (0.0, 0.0, 0.0, 0.0),
                    LightKind::Point { radius } => (1.0, radius, 0.0, 0.0),
                    LightKind::Spot {
                        radius,
                        inner_angle,
                        outer_angle,
                    } => (2.0, radius, inner_angle, outer_angle),
                };
                (
                    l.direction,
                    l.color,
                    [l.position.x, l.position.y, l.position.z, light_type],
                    [radius, inner_angle, outer_angle, l.intensity],
                )
            })
            .unwrap_or_else(|| {
                (
                    glam::Vec3::new(0.5, -1.0, -0.5).normalize(),
                    Color::rgb(1.0, 1.0, 1.0),
                    [0.0, 0.0, 0.0, 0.0], // Directional light by default
                    [0.0, 0.0, 0.0, 1.0], // Intensity = 1.0
                )
            });

        let color_view = color_target.create_view(&wgpu::TextureViewDescriptor::default());

        let (color_load, depth_load) = if should_clear {
            (
                wgpu::LoadOp::Clear(wgpu::Color::TRANSPARENT),
                wgpu::LoadOp::Clear(0.0),
            )
        } else {
            (wgpu::LoadOp::Load, wgpu::LoadOp::Load)
        };

        let mut pass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
            label: Some("mesh3d_instanced.pass"),
            color_attachments: &[Some(wgpu::RenderPassColorAttachment {
                view: &color_view,
                resolve_target: None,
                ops: wgpu::Operations {
                    load: color_load,
                    store: wgpu::StoreOp::Store,
                },
                depth_slice: None,
            })],
            depth_stencil_attachment: Some(wgpu::RenderPassDepthStencilAttachment {
                view: depth_view,
                depth_ops: Some(wgpu::Operations {
                    load: depth_load,
                    store: wgpu::StoreOp::Store,
                }),
                stencil_ops: None,
            }),
            occlusion_query_set: None,
            timestamp_writes: None,
        });

        pass.set_pipeline(&self.pipeline);
        pass.set_bind_group(0, &self.frame_bind_group, &[]);

        for (mesh, instances, texture_res, material) in draws {
            if instances.is_empty() {
                continue;
            }

            let instance_count = instances.len().min(MAX_INSTANCES);

            // Ensure instance buffer is large enough
            if instance_count > self.instance_buffer_capacity {
                self.instance_buffer = state.device.create_buffer(&wgpu::BufferDescriptor {
                    label: Some("mesh3d_instanced.instance_buffer"),
                    size: (std::mem::size_of::<InstanceData>() * instance_count) as u64,
                    usage: wgpu::BufferUsages::VERTEX | wgpu::BufferUsages::COPY_DST,
                    mapped_at_creation: false,
                });
                self.instance_buffer_capacity = instance_count;
            }

            // Write instance data
            state.queue.write_buffer(
                &self.instance_buffer,
                0,
                bytemuck::cast_slice(&instances[..instance_count]),
            );

            // Update frame uniform for this draw batch
            let has_texture = if texture_res.is_some() {
                1.0f32
            } else {
                0.0f32
            };
            let frame_data = InstancedFrameUniform {
                view_proj: view_proj.to_cols_array_2d(),
                camera_pos: [camera.eye.x, camera.eye.y, camera.eye.z, 0.0],
                ambient: [
                    lighting.ambient.r,
                    lighting.ambient.g,
                    lighting.ambient.b,
                    lighting.ambient.a,
                ],
                light_dir: [light_dir.x, light_dir.y, light_dir.z, 0.0],
                light_color: [light_color.r, light_color.g, light_color.b, light_color.a],
                light_pos,
                light_params,
                diffuse: [
                    material.diffuse.r,
                    material.diffuse.g,
                    material.diffuse.b,
                    material.diffuse.a,
                ],
                specular: [
                    material.specular.r,
                    material.specular.g,
                    material.specular.b,
                    material.specular.a,
                ],
                emissive: [
                    material.emissive.r,
                    material.emissive.g,
                    material.emissive.b,
                    material.emissive.a,
                ],
                shininess: material.shininess,
                has_texture,
                _padding: [0.0; 2],
            };
            state
                .queue
                .write_buffer(&self.frame_uniform, 0, bytemuck::bytes_of(&frame_data));

            // Set texture bind group
            if let Some(tex_res) = texture_res {
                let tex_bg = self.create_texture_bind_group(&state.device, tex_res);
                pass.set_bind_group(1, &tex_bg, &[]);
            } else {
                pass.set_bind_group(1, &self.default_texture_bind_group, &[]);
            }

            // Set buffers and draw
            pass.set_vertex_buffer(0, mesh.vertex.slice(..));
            pass.set_vertex_buffer(1, self.instance_buffer.slice(..));
            pass.set_index_buffer(mesh.index.slice(..), wgpu::IndexFormat::Uint16);
            pass.draw_indexed(0..mesh.index_count, 0, 0..instance_count as u32);
        }
    }
}