SolarEngine/solar_engine/src/light.rs

use bytemuck::{Pod, Zeroable};
use cgmath::{EuclideanSpace, Matrix4, Point3, Transform, Vector3};
use wgpu::util::DeviceExt;

#[repr(u32)]
#[derive(Clone, Copy, Debug)]
pub enum LightType {
    Directional = 0,
    Point = 1,
    Spot = 2
}

pub const CLUSTER_COUNT_X: usize = 16;
pub const CLUSTER_COUNT_Y: usize = 9;
pub const CLUSTER_COUNT_Z: usize = 24;
pub const MAX_LIGHTS_PER_CLUSTER: usize = 32;
pub const TOTAL_CLUSTERS: usize = CLUSTER_COUNT_X * CLUSTER_COUNT_Y * CLUSTER_COUNT_Z;

#[repr(C)]
#[derive(Clone, Copy, Pod, Zeroable, Debug)]
pub struct GpuLight {
    pub position: [f32; 3],
    pub light_type: u32,
    pub color: [f32; 3],
    pub intensity: f32,
    pub direction: [f32; 3],
    pub range: f32,
    pub inner_cutoff: f32,
    pub outer_cutoff: f32,
}

pub struct Light {
    pub light_type: LightType,
    pub position: Vector3<f32>,
    pub direction: Vector3<f32>,
    pub color: Vector3<f32>,
    pub intensity: f32,
    pub range: f32,
    pub inner_cutoff: f32,
    pub outer_cutoff: f32,
}

impl Light {
    pub fn to_gpu(&self) -> GpuLight {
        GpuLight {
            position: self.position.into(),
            light_type: self.light_type as u32,
            color: self.color.into(),
            intensity: self.intensity,
            direction: self.direction.into(),
            range: self.range,
            inner_cutoff: self.inner_cutoff,
            outer_cutoff: self.outer_cutoff,
        }
    }
}

pub struct LightManager {
    pub lights: Vec<Light>,
    pub buffer: wgpu::Buffer,
    pub bind_group: wgpu::BindGroup,
    pub count_buffer: wgpu::Buffer,
    pub layout: wgpu::BindGroupLayout,
    pub cluster_buffers: Option<ClusterBuffers>,
}

#[repr(C)]
#[derive(Clone, Copy, Pod, Zeroable, Debug)]
pub struct LightCount {
    pub count: u32,
}

pub struct ClusterBuffers {
    pub light_indices: wgpu::Buffer,
    pub offsets: wgpu::Buffer,
    pub bind_group: wgpu::BindGroup,
    pub layout: wgpu::BindGroupLayout,
}

#[derive(Debug, Clone)]
pub struct ClusterAssignment {
    pub cluster_light_indices: Vec<u32>,
    pub cluster_offsets: Vec<(u32, u32)>,
}

impl LightManager {
    pub fn new(device: &wgpu::Device, max_lights: usize) -> Self {
        let buffer = device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("Light Buffer"),
            size: (max_lights * size_of::<GpuLight>()) as wgpu::BufferAddress,
            usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });

        let count_buffer = device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("Light Count Buffer"),
            size: size_of::<LightCount>() as wgpu::BufferAddress,
            usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });

        let layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
            label: Some("Light Bind Group Layout"),
            entries: &[
                wgpu::BindGroupLayoutEntry {
                    binding: 0,
                    visibility: wgpu::ShaderStages::VERTEX_FRAGMENT,
                    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::FRAGMENT,
                    ty: wgpu::BindingType::Buffer {
                        ty: wgpu::BufferBindingType::Uniform,
                        has_dynamic_offset: false,
                        min_binding_size: None,
                    },
                    count: None,
                },
            ],
        });

        let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some("Light Bind Group"),
            layout: &layout,
            entries: &[
                wgpu::BindGroupEntry {
                    binding: 0,
                    resource: buffer.as_entire_binding(),
                },
                wgpu::BindGroupEntry {
                    binding: 1,
                    resource: count_buffer.as_entire_binding(),
                },
            ],
        });

        Self {
            lights: Vec::new(),
            buffer,
            count_buffer,
            bind_group,
            layout,
            cluster_buffers: None,
        }
    }

    pub fn add_light(&mut self, light: Light) {
        self.lights.push(light);
    }

    pub fn clear(&mut self) {
        self.lights.clear();
    }

    pub fn update_gpu(&self, queue: &wgpu::Queue) {
        let gpu_lights: Vec<GpuLight> = self.lights.iter().map(|l| l.to_gpu()).collect();
        queue.write_buffer(&self.buffer, 0, bytemuck::cast_slice(&gpu_lights));

        let count = LightCount {
            count: self.lights.len() as u32,
        };
        queue.write_buffer(&self.count_buffer, 0, bytemuck::bytes_of(&count));
    }

    pub fn bind_group(&self) -> &wgpu::BindGroup {
        &self.bind_group
    }

    pub fn layout(&self) -> &wgpu::BindGroupLayout {
        &self.layout
    }

    pub fn create_cluster_buffers(
        &self,
        device: &wgpu::Device,
        assignment: &ClusterAssignment,
    ) -> ClusterBuffers {
        let cluster_light_indices = if assignment.cluster_light_indices.is_empty() {
            vec![0u32]
        } else {
            assignment.cluster_light_indices.clone()
        };

        let offset_pairs: Vec<[u32; 2]> = if assignment.cluster_offsets.is_empty() {
            vec![[0, 0]]
        } else {
            assignment.cluster_offsets
                .iter()
                .map(|&(o, c)| [o, c])
                .collect()
        };

        let light_index_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
            label: Some("Cluster Light Indices"),
            contents: bytemuck::cast_slice(&cluster_light_indices),
            usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
        });

        let offset_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
            label: Some("Cluster Offsets"),
            contents: bytemuck::cast_slice(&offset_pairs),
            usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
        });

        let layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
            label: Some("Cluster Bind Group Layout"),
            entries: &[
                wgpu::BindGroupLayoutEntry {
                    binding: 0,
                    visibility: wgpu::ShaderStages::FRAGMENT,
                    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::FRAGMENT,
                    ty: wgpu::BindingType::Buffer {
                        ty: wgpu::BufferBindingType::Storage { read_only: true },
                        has_dynamic_offset: false,
                        min_binding_size: None,
                    },
                    count: None,
                },
            ],
        });

        let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some("Cluster Bind Group"),
            layout: &layout,
            entries: &[
                wgpu::BindGroupEntry {
                    binding: 0,
                    resource: light_index_buffer.as_entire_binding(),
                },
                wgpu::BindGroupEntry {
                    binding: 1,
                    resource: offset_buffer.as_entire_binding(),
                },
            ],
        });

        ClusterBuffers {
            light_indices: light_index_buffer,
            offsets: offset_buffer,
            layout,
            bind_group,
        }
    }

    pub fn update_cluster_buffers(
        &mut self,
        device: &wgpu::Device,
        queue: &wgpu::Queue,
        assignment: &ClusterAssignment,
    ) {
        if let Some(buffers) = &self.cluster_buffers {
            queue.write_buffer(&buffers.light_indices, 0, bytemuck::cast_slice(&assignment.cluster_light_indices));

            let offset_pairs: Vec<[u32; 2]> = assignment
                .cluster_offsets
                .iter()
                .map(|&(o, c)| [o, c])
                .collect();

            queue.write_buffer(&buffers.offsets, 0, bytemuck::cast_slice(&offset_pairs));
        } else {
            let buffers = self.create_cluster_buffers(device, assignment);
            self.cluster_buffers = Some(buffers);
        }
    }

    pub fn compute_cluster_assignments(
        &self,
        view_matrix: Matrix4<f32>,
        _projection_matrix: Matrix4<f32>,
        _screen_width: f32,
        _screen_height: f32,
    ) -> ClusterAssignment {
        let mut cluster_light_lists = vec![Vec::new(); TOTAL_CLUSTERS];

        let log_near = 0.1f32.log2();
        let log_far = 1000.0f32.log2();
        let log_range = log_far - log_near;

        for (i, light) in self.lights.iter().enumerate() {
            let pos_view = view_matrix.transform_point(Point3::from_vec(light.position));
            let radius = light.range;

            let z_bounds = [
                (-pos_view.z - radius).max(0.1).log2(),
                (-pos_view.z + radius).max(0.1).log2(),
            ];
            let z_start = ((z_bounds[0].min(z_bounds[1]) - log_near) / log_range * CLUSTER_COUNT_Z as f32).floor() as usize;
            let z_end   = ((z_bounds[0].max(z_bounds[1]) - log_near) / log_range * CLUSTER_COUNT_Z as f32).ceil() as usize;

            for z in z_start.min(CLUSTER_COUNT_Z)..z_end.min(CLUSTER_COUNT_Z) {
                for y in 0..CLUSTER_COUNT_Y {
                    for x in 0..CLUSTER_COUNT_X {
                        let cluster = x + y * CLUSTER_COUNT_X + z * CLUSTER_COUNT_X * CLUSTER_COUNT_Y;

                        if cluster_light_lists[cluster].len() < MAX_LIGHTS_PER_CLUSTER {
                            cluster_light_lists[cluster].push(i as u32);
                        }
                    }
                }
            }
        }

        let mut cluster_light_indices = Vec::new();
        let mut cluster_offsets = Vec::with_capacity(TOTAL_CLUSTERS);

        for lights in cluster_light_lists {
            let offset = cluster_light_indices.len() as u32;
            let count = lights.len() as u32;
            cluster_light_indices.extend(lights);
            cluster_offsets.push((offset, count));
        }

        ClusterAssignment {
            cluster_light_indices,
            cluster_offsets,
        }
    }
}