diff --git a/crates/bevy_crevice/src/std430/traits.rs b/crates/bevy_crevice/src/std430/traits.rs index 7f2967f3b4..04f9f526d7 100644 --- a/crates/bevy_crevice/src/std430/traits.rs +++ b/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: Copy { /// Convert from self to Std430 diff --git a/crates/bevy_render/src/render_resource/mod.rs b/crates/bevy_render/src/render_resource/mod.rs index 8089199f79..fd60611685 100644 --- a/crates/bevy_render/src/render_resource/mod.rs +++ b/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::*; diff --git a/crates/bevy_render/src/render_resource/storage_buffer.rs b/crates/bevy_render/src/render_resource/storage_buffer.rs new file mode 100644 index 0000000000..874f1527a1 --- /dev/null +++ b/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 { + body: U, + values: Vec, + scratch: Vec, + storage_buffer: Option, +} + +impl Default for StorageBuffer { + /// 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 StorageBuffer { + // 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::(); + const ITEM_SIZE: usize = std::mem::size_of::(); + + /// 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 { + 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 + ::Output::ALIGNMENT - 1) + & !(::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 + } +}