Add a basic brickmap shading table

assets/shaders/voxel_volume.wgsl

@@ -1,6 +1,11 @@
 @group(0) @binding(0) var output: texture_storage_2d<rgba8unorm, write>;
 @group(0) @binding(1) var<storage, read> brickmap: Brickmap;
-@group(0) @binding(2) var<uniform> camera: Camera;
+@group(0) @binding(2) var<storage, read> shading_table: array<ShadingElement>;
+@group(0) @binding(3) var<uniform> camera: Camera;
+
+struct ShadingElement {
+    albedo: u32,
+}
 
 struct Brickmap {
     bitmask: array<u32, 16>,
@@ -26,6 +31,19 @@ struct AabbHitInfo {
     distance: f32,
 };
 
+// TODO: Include brickmap base shading table offset
+fn get_shading_offset(p: vec3<i32>) -> u32 {
+    let local_index = u32(p.x + p.y * 8 + p.z * 64);
+    let bitmask_index = local_index / 32u;
+    var shading_offset = 0u;
+    for (var i: i32 = 0; i < i32(bitmask_index); i++) {
+        shading_offset += countOneBits(brickmap.bitmask[i]);
+    }
+    let extracted_bits = extractBits(brickmap.bitmask[bitmask_index], 0u, (local_index % 32u));
+    shading_offset += countOneBits(extracted_bits);
+    return shading_offset;
+}
+
 fn ray_intersect_aabb(ray_pos: vec3<f32>, ray_dir: vec3<f32>) -> AabbHitInfo {
     let ray_dir_inv = 1.0 / ray_dir;
     let t1 = (vec3<f32>(0.0) - ray_pos) * ray_dir_inv;
@@ -142,6 +160,12 @@ fn compute(@builtin(global_invocation_id) global_id: vec3<u32>) {
         else {
             color = vec4<f32>(1.0);
         }
+        let offset = get_shading_offset(hit_info.hit_pos);
+        let raw_color = shading_table[offset].albedo;
+        color.x = f32((raw_color >> 24u) & 255u) / 255.0;
+        color.y = f32((raw_color >> 16u) & 255u) / 255.0;
+        color.z = f32((raw_color >> 8u) & 255u) / 255.0;
+        color.w = f32(raw_color & 255u) / 255.0;
         // color = textureLoad(voxels_t, hit_info.hit_pos, 0);
     }
 

src/voxel/brickmap.rs

@@ -1,3 +1,5 @@
+use std::mem::size_of;
+
 use wgpu::util::DeviceExt;
 
 use crate::render;
@@ -24,6 +26,8 @@ impl BrickmapUniform {
 pub struct BrickmapManager {
     uniform: BrickmapUniform,
     buffer: wgpu::Buffer,
+    shading_table: Vec<u32>,
+    shading_table_buffer: wgpu::Buffer,
 }
 
 impl BrickmapManager {
@@ -37,20 +41,46 @@ impl BrickmapManager {
                 usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
             });
 
-        Self { uniform, buffer }
+        let shading_table_buffer =
+            context
+                .device
+                .create_buffer_init(&wgpu::util::BufferInitDescriptor {
+                    label: None,
+                    contents: &[0; 25000000],
+                    usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
+                });
+
+        let shading_table = Vec::<u32>::new();
+
+        Self {
+            uniform,
+            buffer,
+            shading_table,
+            shading_table_buffer,
+        }
     }
 
-    pub fn set_data(&mut self, data: &[u32; 16]) {
+    // TODO: Ideally this should take a generic voxel format and do the data mapping here
+    pub fn set_data(&mut self, data: &[u32; 16], colours: &[u32]) {
         self.uniform.bitmask = *data;
+        self.shading_table = colours.to_vec();
     }
 
     pub fn update_buffer(&self, context: &render::Context) {
-        context
+        let queue = &context.queue;
-            .queue
+        queue.write_buffer(&self.buffer, 0, bytemuck::cast_slice(&[self.uniform]));
-            .write_buffer(&self.buffer, 0, bytemuck::cast_slice(&[self.uniform]));
+        queue.write_buffer(
+            &self.shading_table_buffer,
+            0,
+            bytemuck::cast_slice(&self.shading_table),
+        );
     }
 
     pub fn get_buffer(&self) -> &wgpu::Buffer {
         &self.buffer
     }
+
+    pub fn get_shading_buffer(&self) -> &wgpu::Buffer {
+        &self.shading_table_buffer
+    }
 }

src/voxel/voxel_renderer.rs

@@ -61,7 +61,8 @@ impl VoxelRenderer {
 
         log::info!("Creating brickmap manager...");
         let mut brickmap_manager = super::brickmap::BrickmapManager::new(context);
-        let mut data = [0xFFFFFFFF as u32; 16];
+        let mut bitmask_data = [0xFFFFFFFF as u32; 16];
+        let mut albedo_data = Vec::<u32>::new();
         for z in 0..8 {
             let mut entry = 0u64;
             for y in 0..8 {
@@ -70,13 +71,19 @@ impl VoxelRenderer {
                     let pos = glam::vec3(x as f32, y as f32, z as f32) - glam::vec3(3.5, 3.5, 3.5);
                     if pos.length_squared() <= (u32::pow(4, 2) as f32) {
                         entry += 1 << idx;
+                        let mut albedo = 0u32;
+                        albedo += ((x + 1) * 32 - 1) << 24;
+                        albedo += ((y + 1) * 32 - 1) << 16;
+                        albedo += ((z + 1) * 32 - 1) << 8;
+                        albedo += 255;
+                        albedo_data.push(albedo);
                     }
                 }
             }
-            data[2 * z] = (entry & 0xFFFFFFFF).try_into().unwrap();
+            bitmask_data[2 * z as usize] = (entry & 0xFFFFFFFF).try_into().unwrap();
-            data[2 * z + 1] = ((entry >> 32) & 0xFFFFFFFF).try_into().unwrap();
+            bitmask_data[2 * z as usize + 1] = ((entry >> 32) & 0xFFFFFFFF).try_into().unwrap();
         }
-        brickmap_manager.set_data(&data);
+        brickmap_manager.set_data(&bitmask_data, &albedo_data);
         brickmap_manager.update_buffer(context);
 
         log::info!("Creating compute pipeline...");
@@ -101,6 +108,15 @@ impl VoxelRenderer {
                 },
                 None,
             )
+            .with_entry(
+                wgpu::ShaderStages::COMPUTE,
+                wgpu::BindingType::Buffer {
+                    ty: wgpu::BufferBindingType::Storage { read_only: true },
+                    has_dynamic_offset: false,
+                    min_binding_size: None,
+                },
+                None,
+            )
             .with_entry(
                 wgpu::ShaderStages::COMPUTE,
                 wgpu::BindingType::Buffer {
@@ -115,6 +131,7 @@ impl VoxelRenderer {
             .with_layout(&compute_layout)
             .with_entry(wgpu::BindingResource::TextureView(&render_texture.view))
             .with_entry(brickmap_manager.get_buffer().as_entire_binding())
+            .with_entry(brickmap_manager.get_shading_buffer().as_entire_binding())
             .with_entry(camera_controller.get_buffer().as_entire_binding())
             .build(context);
         let compute_pipeline =