lightingHEADmaster

1 files changed, 333 insertions, 0 deletions

diff --git a/src/renderer/lighting.rs b/src/renderer/lighting.rs
new file mode 100644
index 0000000..94c7fb2
--- /dev/null
+++ b/src/renderer/lighting.rs
@@ -0,0 +1,333 @@
+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,
+        }
+    }
+}