bevy_pbr: Rework StorageVec to be StorageBuffer

Supports both a body and variable-size array at the end.

Author: superdump

crates/bevy_crevice/src/std430/traits.rs

@@ -39,6 +39,14 @@ pub unsafe trait Std430: Copy + Zeroable + Pod {
     }
 }
 
+unsafe impl Std430 for () {
+    const ALIGNMENT: usize = 0;
+
+    const PAD_AT_END: bool = false;
+
+    type Padded = ();
+}
+
 /// Trait specifically for Std430::Padded, implements conversions between padded type and base type.
 pub trait Std430Convertible<T: Std430>: Copy {
     /// Convert from self to Std430

crates/bevy_pbr/src/render/light.rs

@@ -95,7 +95,7 @@ pub enum GpuPointLights {
         buffer: UniformVec<[GpuPointLight; MAX_POINT_LIGHTS]>,
     },
     Storage {
-        buffer: StorageVec<GpuPointLight>,
+        buffer: StorageBuffer<(), GpuPointLight>,
     },
 }
 
@@ -116,7 +116,7 @@ impl GpuPointLights {
 
     fn storage() -> Self {
         Self::Storage {
-            buffer: StorageVec::default(),
+            buffer: StorageBuffer::default(),
         }
     }
 
@@ -986,8 +986,8 @@ enum ViewClusterBuffers {
         cluster_offsets_and_counts: UniformVec<[UVec4; ViewClusterBindings::MAX_UNIFORM_ITEMS]>,
     },
     Storage {
-        cluster_light_index_lists: StorageVec<u32>,
-        cluster_offsets_and_counts: StorageVec<UVec2>,
+        cluster_light_index_lists: StorageBuffer<(), u32>,
+        cluster_offsets_and_counts: StorageBuffer<(), UVec2>,
     },
 }
 
@@ -1009,8 +1009,8 @@ impl ViewClusterBuffers {
 
     fn storage() -> Self {
         ViewClusterBuffers::Storage {
-            cluster_light_index_lists: StorageVec::default(),
-            cluster_offsets_and_counts: StorageVec::default(),
+            cluster_light_index_lists: StorageBuffer::default(),
+            cluster_offsets_and_counts: StorageBuffer::default(),
         }
     }
 }
@@ -1049,6 +1049,7 @@ impl ViewClusterBindings {
             ViewClusterBuffers::Storage {
                 cluster_light_index_lists,
                 cluster_offsets_and_counts,
+                ..
             } => {
                 cluster_light_index_lists.clear();
                 cluster_offsets_and_counts.clear();
@@ -1059,8 +1060,8 @@ impl ViewClusterBindings {
     pub fn push_offset_and_count(&mut self, offset: usize, count: usize) {
         match &mut self.buffers {
             ViewClusterBuffers::Uniform {
-                cluster_light_index_lists: _,
                 cluster_offsets_and_counts,
+                ..
             } => {
                 let array_index = self.n_offsets >> 2; // >> 2 is equivalent to / 4
                 if array_index >= Self::MAX_UNIFORM_ITEMS {
@@ -1073,8 +1074,8 @@ impl ViewClusterBindings {
                 cluster_offsets_and_counts.get_mut(0)[array_index][component] = packed;
             }
             ViewClusterBuffers::Storage {
-                cluster_light_index_lists: _,
                 cluster_offsets_and_counts,
+                ..
             } => {
                 cluster_offsets_and_counts.push(UVec2::new(offset as u32, count as u32));
             }
@@ -1091,7 +1092,7 @@ impl ViewClusterBindings {
         match &mut self.buffers {
             ViewClusterBuffers::Uniform {
                 cluster_light_index_lists,
-                cluster_offsets_and_counts: _,
+                ..
             } => {
                 let array_index = self.n_indices >> 4; // >> 4 is equivalent to / 16
                 let component = (self.n_indices >> 2) & ((1 << 2) - 1);
@@ -1103,7 +1104,7 @@ impl ViewClusterBindings {
             }
             ViewClusterBuffers::Storage {
                 cluster_light_index_lists,
-                cluster_offsets_and_counts: _,
+                ..
             } => {
                 cluster_light_index_lists.push(index as u32);
             }
@@ -1135,24 +1136,24 @@ impl ViewClusterBindings {
         match &self.buffers {
             ViewClusterBuffers::Uniform {
                 cluster_light_index_lists,
-                cluster_offsets_and_counts: _,
+                ..
             } => cluster_light_index_lists.binding(),
             ViewClusterBuffers::Storage {
                 cluster_light_index_lists,
-                cluster_offsets_and_counts: _,
+                ..
             } => cluster_light_index_lists.binding(),
         }
     }
 
     pub fn offsets_and_counts_binding(&self) -> Option<BindingResource> {
         match &self.buffers {
             ViewClusterBuffers::Uniform {
-                cluster_light_index_lists: _,
                 cluster_offsets_and_counts,
+                ..
             } => cluster_offsets_and_counts.binding(),
             ViewClusterBuffers::Storage {
-                cluster_light_index_lists: _,
                 cluster_offsets_and_counts,
+                ..
             } => cluster_offsets_and_counts.binding(),
         }
     }

crates/bevy_pbr/src/render/mesh_view_bind_group.wgsl

@@ -72,15 +72,12 @@ struct ClusterOffsetsAndCounts {
 };
 #else
 struct PointLights {
-    length: u32;
     data: array<PointLight>;
 };
 struct ClusterLightIndexLists {
-    length: u32;
     data: array<u32>;
 };
 struct ClusterOffsetsAndCounts {
-    length: u32;
     data: array<vec2<u32>>;
 };
 #endif

crates/bevy_render/src/render_resource/mod.rs

@@ -6,7 +6,7 @@ mod pipeline;
 mod pipeline_cache;
 mod pipeline_specializer;
 mod shader;
-mod storage_vec;
+mod storage_buffer;
 mod texture;
 mod uniform_vec;
 
@@ -18,7 +18,7 @@ pub use pipeline::*;
 pub use pipeline_cache::*;
 pub use pipeline_specializer::*;
 pub use shader::*;
-pub use storage_vec::*;
+pub use storage_buffer::*;
 pub use texture::*;
 pub use uniform_vec::*;
 

crates/bevy_render/src/render_resource/storage_vec.rs renamed to crates/bevy_render/src/render_resource/storage_buffer.rs

@@ -8,29 +8,30 @@ use crate::renderer::{RenderDevice, RenderQueue};
 
 use super::Buffer;
 
-/// A helper for a storage buffer binding with a variable-sized array
-/// Use a struct with length: u32, data: array<T> in the shader.
-pub struct StorageVec<T: AsStd430> {
-    values: Vec<T>,
+/// A helper for a storage buffer binding with a body, or a variable-sized array, or both.
+pub struct StorageBuffer<T: AsStd430, U: AsStd430> {
+    body: T,
+    values: Vec<U>,
     scratch: Vec<u8>,
     storage_buffer: Option<Buffer>,
     capacity: usize,
     item_size: usize,
 }
 
-impl<T: AsStd430> Default for StorageVec<T> {
+impl<T: AsStd430 + Default, U: AsStd430> Default for StorageBuffer<T, U> {
     fn default() -> Self {
         Self {
+            body: T::default(),
             values: Vec::new(),
             scratch: Vec::new(),
             storage_buffer: None,
             capacity: 0,
-            item_size: T::std430_size_static(),
+            item_size: U::std430_size_static(),
         }
     }
 }
 
-impl<T: AsStd430> StorageVec<T> {
+impl<T: AsStd430, U: AsStd430> StorageBuffer<T, U> {
     #[inline]
     pub fn storage_buffer(&self) -> Option<&Buffer> {
         self.storage_buffer.as_ref()
@@ -55,22 +56,26 @@ impl<T: AsStd430> StorageVec<T> {
         self.values.is_empty()
     }
 
+    pub fn set_body(&mut self, body: T) {
+        self.body = body;
+    }
+
     #[inline]
     pub fn capacity(&self) -> usize {
         self.capacity
     }
 
-    pub fn push(&mut self, value: T) -> usize {
+    pub fn push(&mut self, value: U) -> usize {
         let index = self.values.len();
         self.values.push(value);
         index
     }
 
-    pub fn get_mut(&mut self, index: usize) -> &mut T {
+    pub fn get_mut(&mut self, index: usize) -> &mut U {
         &mut self.values[index]
     }
 
-    pub fn reserve(&mut self, capacity: usize, device: &RenderDevice) -> bool {
+    fn reserve(&mut self, capacity: usize, device: &RenderDevice) -> bool {
         if self.storage_buffer.is_none() || capacity > self.capacity {
             self.capacity = capacity;
             let size = self.size();
@@ -87,33 +92,44 @@ impl<T: AsStd430> StorageVec<T> {
         }
     }
 
-    pub fn size(&self) -> usize {
-        // 4 bytes for a u32 array length plus the necessary padding according to T's alignment
-        ((4 + <T as AsStd430>::Output::ALIGNMENT - 1) & !(<T as AsStd430>::Output::ALIGNMENT - 1))
-            // Variable size arrays must have at least 1 element
-            + self.item_size * self.capacity.max(1)
+    fn size(&self) -> usize {
+        let mut size = 0;
+        size += T::std430_size_static();
+        if size > 0 {
+            // Pad according to the array item type's alignment
+            size = (size + <U as AsStd430>::Output::ALIGNMENT - 1)
+                & !(<U as AsStd430>::Output::ALIGNMENT - 1);
+        }
+        // Variable size arrays must have at least 1 element
+        size += self.item_size * self.capacity.max(1);
+        size
     }
 
     pub fn write_buffer(&mut self, device: &RenderDevice, queue: &RenderQueue) {
         self.reserve(self.values.len(), device);
         if let Some(storage_buffer) = &self.storage_buffer {
             let range = 0..self.size();
             let mut writer = std430::Writer::new(&mut self.scratch[range.clone()]);
-            // NOTE: Can only represent up to u32::MAX values
-            debug_assert!(self.values.len() <= u32::MAX as usize);
-            // NOTE: Failing to write should be non-fatal. It would likely cause visual
-            // artifacts but a warning message should suffice.
-            // First write the length of the array
-            writer
-                .write(&(self.values.len() as u32))
-                .map_err(|e| warn!("{:?}", e))
-                .ok();
-            // Then write the array. Note that padding bytes may be added between the array length
-            // and the array in order to align the array to the alignment requirements of T
-            writer
-                .write(self.values.as_slice())
-                .map_err(|e| warn!("{:?}", e))
-                .ok();
+            let mut offset = 0;
+            // First write the struct body if there is one
+            if T::std430_size_static() > 0 {
+                if let Ok(new_offset) = writer.write(&self.body).map_err(|e| warn!("{:?}", e)) {
+                    offset = new_offset;
+                }
+            }
+            if self.values.is_empty() {
+                for i in offset..self.size() {
+                    self.scratch[i] = 0;
+                }
+            } else {
+                // Then write the array. Note that padding bytes may be added between the body
+                // and the array in order to align the array to the alignment requirements of its
+                // items
+                writer
+                    .write(self.values.as_slice())
+                    .map_err(|e| warn!("{:?}", e))
+                    .ok();
+            }
             queue.write_buffer(storage_buffer, 0, &self.scratch[range]);
         }
     }
@@ -122,7 +138,7 @@ impl<T: AsStd430> StorageVec<T> {
         self.values.clear();
     }
 
-    pub fn values(&self) -> &[T] {
+    pub fn values(&self) -> &[U] {
         &self.values
     }
 }