Add a helper for storage buffers similar to UniformVec (#4079)

# Objective

- Add a helper for storage buffers similar to `UniformVec`

## Solution

- Add a `StorageBuffer<T, U>` where `T` is the main body of the shader struct without any final variable-sized array member, and `U` is the type of the items in a variable-sized array.
- Use `()` as the type for unwanted parts, e.g. `StorageBuffer<(), Vec4>::default()` would construct a binding that would work with `struct MyType { data: array<vec4<f32>>; }` in WGSL and `StorageBuffer<MyType, ()>::default()` would work with `struct MyType { ... }` in WGSL as long as there are no variable-sized arrays.
- Std430 requires that there is at most one variable-sized array in a storage buffer, that if there is one it is the last member of the binding, and that it has at least one item. `StorageBuffer` handles all of these constraints.

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_render/src/render_resource/mod.rs

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

crates/bevy_render/src/render_resource/storage_buffer.rs

@@ -0,0 +1,129 @@
+use std::num::NonZeroU64;
+
+use bevy_crevice::std430::{self, AsStd430, Std430};
+use bevy_utils::tracing::warn;
+use wgpu::{BindingResource, BufferBinding, BufferDescriptor, BufferUsages};
+
+use crate::renderer::{RenderDevice, RenderQueue};
+
+use super::Buffer;
+
+/// 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: U,
+    values: Vec<T>,
+    scratch: Vec<u8>,
+    storage_buffer: Option<Buffer>,
+}
+
+impl<T: AsStd430, U: AsStd430 + Default> Default for StorageBuffer<T, U> {
+    /// Creates a new [`StorageBuffer`]
+    ///
+    /// This does not immediately allocate system/video RAM buffers.
+    fn default() -> Self {
+        Self {
+            body: U::default(),
+            values: Vec::new(),
+            scratch: Vec::new(),
+            storage_buffer: None,
+        }
+    }
+}
+
+impl<T: AsStd430, U: AsStd430> StorageBuffer<T, U> {
+    // NOTE: AsStd430::std430_size_static() uses size_of internally but trait functions cannot be
+    // marked as const functions
+    const BODY_SIZE: usize = std::mem::size_of::<U>();
+    const ITEM_SIZE: usize = std::mem::size_of::<T>();
+
+    /// Gets the reference to the underlying buffer, if one has been allocated.
+    #[inline]
+    pub fn buffer(&self) -> Option<&Buffer> {
+        self.storage_buffer.as_ref()
+    }
+
+    #[inline]
+    pub fn binding(&self) -> Option<BindingResource> {
+        Some(BindingResource::Buffer(BufferBinding {
+            buffer: self.buffer()?,
+            offset: 0,
+            size: Some(NonZeroU64::new((self.size()) as u64).unwrap()),
+        }))
+    }
+
+    #[inline]
+    pub fn set_body(&mut self, body: U) {
+        self.body = body;
+    }
+
+    fn reserve_buffer(&mut self, device: &RenderDevice) -> bool {
+        let size = self.size();
+        if self.storage_buffer.is_none() || size > self.scratch.len() {
+            self.scratch.resize(size, 0);
+            self.storage_buffer = Some(device.create_buffer(&BufferDescriptor {
+                label: None,
+                size: size as wgpu::BufferAddress,
+                usage: BufferUsages::COPY_DST | BufferUsages::STORAGE,
+                mapped_at_creation: false,
+            }));
+            true
+        } else {
+            false
+        }
+    }
+
+    fn size(&self) -> usize {
+        let mut size = 0;
+        size += Self::BODY_SIZE;
+        if Self::ITEM_SIZE > 0 {
+            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.values.len().max(1);
+        }
+        size
+    }
+
+    pub fn write_buffer(&mut self, device: &RenderDevice, queue: &RenderQueue) {
+        self.reserve_buffer(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()]);
+            let mut offset = 0;
+            // First write the struct body if there is one
+            if Self::BODY_SIZE > 0 {
+                if let Ok(new_offset) = writer.write(&self.body).map_err(|e| warn!("{:?}", e)) {
+                    offset = new_offset;
+                }
+            }
+            if Self::ITEM_SIZE > 0 {
+                if self.values.is_empty() {
+                    // Zero-out the padding and dummy array item in the case of the array being 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]);
+        }
+    }
+
+    pub fn values(&self) -> &[T] {
+        &self.values
+    }
+
+    pub fn values_mut(&mut self) -> &mut [T] {
+        &mut self.values
+    }
+}