bevy_pbr: Always write length to StorageVec

Author: superdump

crates/bevy_pbr/src/render/mesh_view_bind_group.wgsl

@@ -72,12 +72,15 @@ 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/storage_vec.rs

@@ -8,29 +8,29 @@ use crate::renderer::{RenderDevice, RenderQueue};
 
 use super::Buffer;
 
-pub struct StorageVec<T: AsStd430 + Default> {
+/// 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>,
     scratch: Vec<u8>,
     storage_buffer: Option<Buffer>,
     capacity: usize,
     item_size: usize,
 }
 
-impl<T: AsStd430 + Default> Default for StorageVec<T> {
+impl<T: AsStd430> Default for StorageVec<T> {
     fn default() -> Self {
         Self {
             values: Vec::new(),
             scratch: Vec::new(),
             storage_buffer: None,
             capacity: 0,
-            // FIXME: Is this correct or even necessary? Maybe only need T::std430_size_static()?
-            item_size: (T::std430_size_static() + <T as AsStd430>::Output::ALIGNMENT - 1)
-                & !(<T as AsStd430>::Output::ALIGNMENT - 1),
+            item_size: T::std430_size_static(),
         }
     }
 }
 
-impl<T: AsStd430 + Default> StorageVec<T> {
+impl<T: AsStd430> StorageVec<T> {
     #[inline]
     pub fn storage_buffer(&self) -> Option<&Buffer> {
         self.storage_buffer.as_ref()
@@ -41,7 +41,7 @@ impl<T: AsStd430 + Default> StorageVec<T> {
         Some(BindingResource::Buffer(BufferBinding {
             buffer: self.storage_buffer()?,
             offset: 0,
-            size: Some(NonZeroU64::new((self.item_size * self.values.len()) as u64).unwrap()),
+            size: Some(NonZeroU64::new((self.size()) as u64).unwrap()),
         }))
     }
 
@@ -71,9 +71,9 @@ impl<T: AsStd430 + Default> StorageVec<T> {
     }
 
     pub fn reserve(&mut self, capacity: usize, device: &RenderDevice) -> bool {
-        if capacity > self.capacity {
+        if self.storage_buffer.is_none() || capacity > self.capacity {
             self.capacity = capacity;
-            let size = self.item_size * capacity;
+            let size = self.size();
             self.scratch.resize(size, 0);
             self.storage_buffer = Some(device.create_buffer(&BufferDescriptor {
                 label: None,
@@ -87,16 +87,29 @@ impl<T: AsStd430 + Default> 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)
+    }
+
     pub fn write_buffer(&mut self, device: &RenderDevice, queue: &RenderQueue) {
-        if self.values.is_empty() {
-            self.values.push(T::default());
-        }
         self.reserve(self.values.len(), device);
         if let Some(storage_buffer) = &self.storage_buffer {
-            let range = 0..self.item_size * self.values.len();
+            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))