path: root/src/renderer/lighting.rs

use glam::{IVec2, Vec2};
use wgpu::TextureViewDescriptor;

use crate::{
    camera::Camera,
    content::materials::{MaterialDef, MaterialId},
    sim::sim_manager::SimManager,
};

// light resolution = light pixels per cell
const LIGHT_RESOLUTION: f32 = 1.0;
const LAYERS: u32 = 4;
const WORKGROUP_SIZE: u32 = 8;

#[repr(C)]
#[derive(Copy, Clone, bytemuck::Pod, bytemuck::Zeroable)]
struct LightingUniform {
    // world position of the lighting texture's (0, 0) texel
    origin: [f32; 2],
    _padding: [f32; 2],
}

fn create_scene_texture(device: &wgpu::Device, size: IVec2) -> wgpu::Texture {
    device.create_texture(&wgpu::TextureDescriptor {
        label: Some("Scene texture"),
        size: wgpu::Extent3d {
            width: size.x as u32,
            height: size.y as u32,
            depth_or_array_layers: 1,
        },
        usage: wgpu::TextureUsages::STORAGE_BINDING
            | wgpu::TextureUsages::TEXTURE_BINDING
            | wgpu::TextureUsages::COPY_DST,
        format: wgpu::TextureFormat::Rgba8Unorm,
        dimension: wgpu::TextureDimension::D2,
        mip_level_count: 1,
        sample_count: 1,
        view_formats: &[],
    })
}

fn create_lighting_texture(device: &wgpu::Device, layers: u32, size: IVec2) -> wgpu::Texture {
    device.create_texture(&wgpu::TextureDescriptor {
        label: Some("Lighting texture"),
        size: wgpu::Extent3d {
            width: size.x as u32,
            height: size.y as u32,
            depth_or_array_layers: layers,
        },
        usage: wgpu::TextureUsages::STORAGE_BINDING | wgpu::TextureUsages::TEXTURE_BINDING,
        format: wgpu::TextureFormat::Rgba16Float,
        dimension: wgpu::TextureDimension::D2,
        mip_level_count: 1,
        sample_count: 1,
        view_formats: &[],
    })
}

fn create_lighting_view(texture: &wgpu::Texture) -> wgpu::TextureView {
    texture.create_view(&TextureViewDescriptor {
        dimension: Some(wgpu::TextureViewDimension::D2Array),
        array_layer_count: Some(LAYERS),
        ..TextureViewDescriptor::default()
    })
}

fn create_lighting_bind_group(
    device: &wgpu::Device,
    layout: &wgpu::BindGroupLayout,
    scene_texture: &wgpu::Texture,
    src: &wgpu::Texture,
    dst: &wgpu::Texture,
) -> wgpu::BindGroup {
    device.create_bind_group(&wgpu::BindGroupDescriptor {
        label: Some("Lighting bind group"),
        layout,
        entries: &[
            wgpu::BindGroupEntry {
                binding: 0,
                resource: wgpu::BindingResource::TextureView(&scene_texture.create_view(
                    &TextureViewDescriptor {
                        dimension: Some(wgpu::TextureViewDimension::D2),
                        ..TextureViewDescriptor::default()
                    },
                )),
            },
            wgpu::BindGroupEntry {
                binding: 1,
                resource: wgpu::BindingResource::TextureView(&create_lighting_view(src)),
            },
            wgpu::BindGroupEntry {
                binding: 2,
                resource: wgpu::BindingResource::TextureView(&create_lighting_view(dst)),
            },
        ],
    })
}

pub struct LightingRendererState {
    lighting_pipeline: wgpu::ComputePipeline,
    lighting_bind_group_layout: wgpu::BindGroupLayout,
    uniform_buffer: wgpu::Buffer,
    scene_texture: wgpu::Texture,
    lighting_texture_a: wgpu::Texture,
    lighting_texture_b: wgpu::Texture,
    lighting_bind_group_a: wgpu::BindGroup,
    lighting_bind_group_b: wgpu::BindGroup,
}

impl LightingRendererState {
    pub fn lighting_view(&self) -> wgpu::TextureView {
        create_lighting_view(&self.lighting_texture_a)
    }

    pub fn uniform_buffer(&self) -> &wgpu::Buffer {
        &self.uniform_buffer
    }

    pub fn update_scene_texture(
        &mut self,
        device: &wgpu::Device,
        queue: &wgpu::Queue,
        sim: &mut SimManager,
        camera: &Camera,
    ) -> bool {
        puffin::profile_function!();
        let (lower, upper) = camera.viewport_bounds_world();
        let origin = (lower * LIGHT_RESOLUTION).floor() / LIGHT_RESOLUTION;
        let size = ((upper - origin) * LIGHT_RESOLUTION).ceil().as_ivec2();

        let recreated = self.scene_texture.size().width != size.x as u32
            || self.scene_texture.size().height != size.y as u32;

        if recreated {
            self.scene_texture = create_scene_texture(device, size);
            self.lighting_texture_a = create_lighting_texture(device, LAYERS, size);
            self.lighting_texture_b = create_lighting_texture(device, LAYERS, size);
            self.lighting_bind_group_a = create_lighting_bind_group(
                device,
                &self.lighting_bind_group_layout,
                &self.scene_texture,
                &self.lighting_texture_a,
                &self.lighting_texture_b,
            );
            self.lighting_bind_group_b = create_lighting_bind_group(
                device,
                &self.lighting_bind_group_layout,
                &self.scene_texture,
                &self.lighting_texture_b,
                &self.lighting_texture_a,
            );
        }

        queue.write_buffer(
            &self.uniform_buffer,
            0,
            bytemuck::bytes_of(&LightingUniform {
                origin: origin.to_array(),
                _padding: [0.0; 2],
            }),
        );

        let mut data: Vec<u8> = vec![
            0;
            (self.scene_texture.size().width * self.scene_texture.size().height * 4)
                as usize
        ];

        for y in 0..self.scene_texture.size().height {
            for x in 0..self.scene_texture.size().width {
                let local = Vec2::new(x as f32, y as f32);
                let world = origin + local * LIGHT_RESOLUTION;
                let cell = sim
                    .cell_manager
                    .get_cell_from_game_position(world.x.round() as i32, world.y.round() as i32);
                let material = cell.map_or(MaterialId::Void.def(), |c| c.material.def());

                let idx = (x + y * self.scene_texture.size().width) as usize * 4;
                data[idx] = material.emission.0;
                data[idx + 1] = material.emission.1;
                data[idx + 2] = material.emission.2;
                data[idx + 3] = 0xFF - (material.opacity / 2);
            }
        }

        queue.write_texture(
            wgpu::TexelCopyTextureInfoBase {
                texture: &self.scene_texture,
                mip_level: 0,
                origin: wgpu::Origin3d::ZERO,
                aspect: wgpu::TextureAspect::All,
            },
            &data,
            wgpu::TexelCopyBufferLayout {
                bytes_per_row: Some(4 * self.scene_texture.size().width),
                rows_per_image: Some(self.scene_texture.size().height),
                offset: 0,
            },
            self.scene_texture.size(),
        );

        recreated
    }

    pub fn compute(&self, compute_pass: &mut wgpu::ComputePass) {
        compute_pass.set_pipeline(&self.lighting_pipeline);
        // must be even number of iterations so the result ends up in texture a
        for i in 0..4 {
            compute_pass.set_bind_group(
                0,
                if i % 2 == 0 {
                    &self.lighting_bind_group_a
                } else {
                    &self.lighting_bind_group_b
                },
                &[],
            );
            compute_pass.dispatch_workgroups(
                self.lighting_texture_a
                    .size()
                    .width
                    .div_ceil(WORKGROUP_SIZE),
                self.lighting_texture_a
                    .size()
                    .height
                    .div_ceil(WORKGROUP_SIZE),
                1,
            );
        }
    }

    pub fn new(device: &wgpu::Device) -> Self {
        let lpv_shader = device.create_shader_module(wgpu::ShaderModuleDescriptor {
            label: Some("Light propagation volume shader"),
            source: wgpu::ShaderSource::Wgsl(include_str!("../shader/lpv.wgsl").into()),
        });

        let lighting_bind_group_layout =
            device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
                label: None,
                entries: &[
                    // scene
                    wgpu::BindGroupLayoutEntry {
                        ty: wgpu::BindingType::Texture {
                            multisampled: false,
                            sample_type: wgpu::TextureSampleType::Float { filterable: false },
                            view_dimension: wgpu::TextureViewDimension::D2,
                        },
                        binding: 0,
                        count: None,
                        visibility: wgpu::ShaderStages::COMPUTE,
                    },
                    // src
                    wgpu::BindGroupLayoutEntry {
                        ty: wgpu::BindingType::Texture {
                            multisampled: false,
                            sample_type: wgpu::TextureSampleType::Float { filterable: false },
                            view_dimension: wgpu::TextureViewDimension::D2Array,
                        },
                        binding: 1,
                        count: None,
                        visibility: wgpu::ShaderStages::COMPUTE,
                    },
                    // dst
                    wgpu::BindGroupLayoutEntry {
                        ty: wgpu::BindingType::StorageTexture {
                            access: wgpu::StorageTextureAccess::WriteOnly,
                            format: wgpu::TextureFormat::Rgba16Float,
                            view_dimension: wgpu::TextureViewDimension::D2Array,
                        },
                        binding: 2,
                        count: None,
                        visibility: wgpu::ShaderStages::COMPUTE,
                    },
                ],
            });

        let lighting_pipeline_layout =
            device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
                label: None,
                immediate_size: 0,
                bind_group_layouts: &[Some(&lighting_bind_group_layout)],
            });

        let lighting_pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
            label: None,
            layout: Some(&lighting_pipeline_layout),
            entry_point: Some("propagate"),
            compilation_options: wgpu::PipelineCompilationOptions::default(),
            module: &lpv_shader,
            cache: None,
        });

        let uniform_buffer = device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("Lighting uniform buffer"),
            size: size_of::<LightingUniform>() as u64,
            usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });

        let size = IVec2::ONE;

        let scene_texture = create_scene_texture(device, size);
        let lighting_texture_a = create_lighting_texture(device, LAYERS, size);
        let lighting_texture_b = create_lighting_texture(device, LAYERS, size);

        let lighting_bind_group_a = create_lighting_bind_group(
            device,
            &lighting_bind_group_layout,
            &scene_texture,
            &lighting_texture_a,
            &lighting_texture_b,
        );
        let lighting_bind_group_b = create_lighting_bind_group(
            device,
            &lighting_bind_group_layout,
            &scene_texture,
            &lighting_texture_b,
            &lighting_texture_a,
        );

        LightingRendererState {
            lighting_pipeline,
            lighting_bind_group_layout,
            uniform_buffer,
            lighting_bind_group_a,
            lighting_bind_group_b,
            scene_texture,
            lighting_texture_a,
            lighting_texture_b,
        }
    }
}