Move brickmap cache staging and uploading

src/voxel/brickworld/brickmap.rs

@@ -25,15 +25,6 @@ struct WorldState {
     _pad: u32,
 }
 
-#[repr(C)]
-#[derive(Debug, Copy, Clone, bytemuck::Pod, bytemuck::Zeroable)]
-struct BrickmapUnpackElement {
-    cache_idx: u32,
-    brickmap: Brickmap,
-    shading_element_count: u32,
-    shading_elements: [u32; 512],
-}
-
 #[derive(Debug)]
 pub struct BrickmapManager {
     state_uniform: WorldState,
@@ -45,8 +36,6 @@ pub struct BrickmapManager {
     feedback_buffer: wgpu::Buffer,
     feedback_result_buffer: wgpu::Buffer,
     unpack_max_count: usize,
-    brickmap_staged: Vec<BrickmapUnpackElement>,
-    brickmap_unpack_buffer: wgpu::Buffer,
 }
 
 // TODO:
@@ -66,8 +55,11 @@ impl BrickmapManager {
         };
 
         let brickgrid = Brickgrid::new(context, brickgrid_dims, max_uploaded_brickmaps as usize);
-
+        let brickmap_cache = BrickmapCache::new(
-        let brickmap_cache = BrickmapCache::new(context, brickmap_cache_size);
+            context,
+            brickmap_cache_size,
+            max_uploaded_brickmaps as usize,
+        );
 
         let shading_table_allocator = ShadingTableAllocator::new(4, shading_table_bucket_size);
         let shading_table = vec![0u32; shading_table_allocator.total_elements as usize];
@@ -78,13 +70,11 @@ impl BrickmapManager {
 
         let mut brickmap_upload_data = vec![0u32; 4 + 532 * max_uploaded_brickmaps as usize];
         brickmap_upload_data[0] = max_uploaded_brickmaps;
-        let brickmap_staged = Vec::new();
 
         let mut buffers = gfx::BulkBufferBuilder::new()
             .with_init_buffer_bm("Brick World State", &[state_uniform])
             .set_usage(wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST)
             .with_init_buffer_bm("Shading Table", &shading_table)
-            .with_init_buffer_bm("Brickmap Unpack", &brickmap_upload_data)
             .set_usage(
                 wgpu::BufferUsages::STORAGE
                     | wgpu::BufferUsages::COPY_DST
@@ -101,11 +91,9 @@ impl BrickmapManager {
             brickmap_cache,
             shading_table_allocator,
             unpack_max_count: max_uploaded_brickmaps as usize,
-            brickmap_staged,
 
             state_buffer: buffers.remove(0),
             shading_table_buffer: buffers.remove(0),
-            brickmap_unpack_buffer: buffers.remove(0),
             feedback_buffer: buffers.remove(0),
             feedback_result_buffer: buffers.remove(0),
         }
@@ -136,7 +124,7 @@ impl BrickmapManager {
     }
 
     pub fn get_brickmap_unpack_buffer(&self) -> &wgpu::Buffer {
-        &self.brickmap_unpack_buffer
+        self.brickmap_cache.get_upload_buffer()
     }
 
     pub fn get_brickgrid_unpack_buffer(&self) -> &wgpu::Buffer {
@@ -209,7 +197,10 @@ impl BrickmapManager {
             .try_alloc(albedo_data.len() as u32)
             .unwrap() as usize;
 
-        if let Some(entry) = self.brickmap_cache.add_entry(grid_idx, shading_idx as u32) {
+        if let Some(entry) =
+            self.brickmap_cache
+                .add_entry(grid_idx, shading_idx as u32, bitmask_data, albedo_data)
+        {
             self.update_brickgrid_element(entry.grid_idx, 1);
         }
 
@@ -231,25 +222,6 @@ impl BrickmapManager {
                 log::warn!("{}", e)
             }
         }
-
-        // Update the brickmap
-        let brickmap = Brickmap {
-            bitmask: bitmask_data,
-            shading_table_offset: shading_idx as u32,
-            lod_color: 0,
-        };
-
-        let shading_element_count = albedo_data.len();
-        let mut shading_elements = [0u32; 512];
-        shading_elements[..shading_element_count].copy_from_slice(&albedo_data);
-
-        let staged_brickmap = BrickmapUnpackElement {
-            cache_idx: self.brickmap_cache.index as u32,
-            brickmap,
-            shading_element_count: shading_element_count as u32,
-            shading_elements,
-        };
-        self.brickmap_staged.push(staged_brickmap);
     }
 
     fn update_brickgrid_element(&mut self, index: usize, data: u32) -> Option<BrickmapCacheEntry> {
@@ -265,30 +237,8 @@ impl BrickmapManager {
         brickmap_cache_entry
     }
 
-    // TODO: Tidy this up more
     fn upload_unpack_buffers(&mut self, context: &gfx::Context) {
         self.brickgrid.upload(context);
-
+        self.brickmap_cache.upload(context);
-        // Brickmap
-        let end = self.unpack_max_count.min(self.brickmap_staged.len());
-        let iter = self.brickmap_staged.drain(0..end);
-        let data = iter.as_slice();
-        context.queue.write_buffer(
-            &self.brickmap_unpack_buffer,
-            4,
-            bytemuck::cast_slice(&[end]),
-        );
-        context
-            .queue
-            .write_buffer(&self.brickmap_unpack_buffer, 16, bytemuck::cast_slice(data));
-        drop(iter);
-
-        if end > 0 {
-            log::info!(
-                "Uploading {} brickmap entries. ({} remaining)",
-                end,
-                self.brickmap_staged.len()
-            );
-        }
     }
 }

src/voxel/brickworld/brickmap_cache.rs

@@ -8,28 +8,50 @@ pub struct BrickmapCacheEntry {
     pub shading_table_offset: u32,
 }
 
+#[repr(C)]
+#[derive(Debug, Copy, Clone, bytemuck::Pod, bytemuck::Zeroable)]
+// TODO: Rename struct
+struct BrickmapUnpackElement {
+    cache_idx: u32, // TODO: Change to usize?
+    brickmap: Brickmap,
+    shading_element_count: u32,
+    shading_elements: [u32; 512], // TODO: Replace u32 with custom type?
+}
+
 #[derive(Debug)]
 pub struct BrickmapCache {
-    buffer: wgpu::Buffer,
     cache: Vec<Option<BrickmapCacheEntry>>,
     pub index: usize,
     pub num_loaded: u32,
+    staged: Vec<BrickmapUnpackElement>,
+    max_upload_count: usize,
+    buffer: wgpu::Buffer,
+    upload_buffer: wgpu::Buffer,
 }
 
 impl BrickmapCache {
-    pub fn new(context: &Context, size: usize) -> Self {
+    pub fn new(context: &Context, size: usize, max_upload_count: usize) -> Self {
-        let buffer_data = vec![Brickmap::default(); size];
+        let data = vec![Brickmap::default(); size];
-        let buffer = BulkBufferBuilder::new()
+
+        // TODO: change type of upload data. Will need some messyness with bytemucking probably
+        // but should lead to clearer data definitions
+        let mut upload_data = vec![0u32; 4 + 532 * max_upload_count as usize];
+        upload_data[0] = max_upload_count as u32;
+
+        let mut buffers = BulkBufferBuilder::new()
             .set_usage(wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST)
-            .with_init_buffer_bm("Brickmap Cache", &buffer_data)
+            .with_init_buffer_bm("Brickmap Cache", &data)
-            .build(context)
+            .with_init_buffer_bm("Brickmap Unpack", &upload_data)
-            .remove(0);
+            .build(context);
 
         Self {
-            buffer,
             cache: vec![None; size],
             index: 0,
             num_loaded: 0,
+            staged: vec![],
+            max_upload_count,
+            buffer: buffers.remove(0),
+            upload_buffer: buffers.remove(0),
         }
     }
 
@@ -37,11 +59,17 @@ impl BrickmapCache {
         &self.buffer
     }
 
+    pub fn get_upload_buffer(&self) -> &wgpu::Buffer {
+        &self.upload_buffer
+    }
+
     /// Adds a brickmap entry and returns the entry that was overwritten.
     pub fn add_entry(
         &mut self,
         grid_idx: usize,
         shading_table_offset: u32,
+        bitmask: [u32; 16],
+        albedo_data: Vec<u32>,
     ) -> Option<BrickmapCacheEntry> {
         // We do this first because we want this to be the index of the most recently added entry
         // This has the side effect of meaning that on the first loop through the cache the first
@@ -58,6 +86,25 @@ impl BrickmapCache {
             shading_table_offset,
         });
 
+        // Need to stage this entry
+        let brickmap = Brickmap {
+            bitmask,
+            shading_table_offset,
+            lod_color: 0,
+        };
+
+        let shading_element_count = albedo_data.len();
+        let mut shading_elements = [0u32; 512];
+        shading_elements[..shading_element_count].copy_from_slice(&albedo_data);
+
+        let staged_brickmap = BrickmapUnpackElement {
+            cache_idx: self.index as u32,
+            brickmap,
+            shading_element_count: shading_element_count as u32,
+            shading_elements,
+        };
+        self.staged.push(staged_brickmap);
+
         existing_entry
     }
 
@@ -75,4 +122,29 @@ impl BrickmapCache {
     pub fn get_entry(&self, index: usize) -> Option<BrickmapCacheEntry> {
         self.cache[index]
     }
+
+    pub fn upload(&mut self, context: &Context) {
+        // Takes up to max_upload_count upload elements
+        let count = usize::min(self.max_upload_count, self.staged.len());
+        let iter = self.staged.drain(0..count);
+        let upload_data = iter.as_slice();
+
+        // Upload buffer is {max_count, count, pad, pad, maps[]}. So we need to add
+        // the count and pads, and upload at an offset to skip max_count
+        let data: Vec<u8> = [
+            bytemuck::cast_slice(&[count as u32, 0, 0]),
+            bytemuck::cast_slice(upload_data),
+        ]
+        .concat();
+        context.queue.write_buffer(&self.upload_buffer, 4, &data);
+        drop(iter);
+
+        if count > 0 {
+            log::info!(
+                "Uploading {} brickmap entries. ({} remaining)",
+                count,
+                self.staged.len()
+            );
+        }
+    }
 }