bevy_render/render_resource/
storage_buffer.rs

1use core::marker::PhantomData;
2
3use super::Buffer;
4use crate::{
5    render_resource::make_buffer_label,
6    renderer::{RenderDevice, RenderQueue},
7};
8use encase::{
9    internal::WriteInto, DynamicStorageBuffer as DynamicStorageBufferWrapper, ShaderType,
10    StorageBuffer as StorageBufferWrapper,
11};
12use wgpu::{util::BufferInitDescriptor, BindingResource, BufferBinding, BufferSize, BufferUsages};
13
14use super::IntoBinding;
15
16/// Stores data to be transferred to the GPU and made accessible to shaders as a storage buffer.
17///
18/// Storage buffers can be made available to shaders in some combination of read/write mode, and can store large amounts of data.
19/// Note however that WebGL2 does not support storage buffers, so consider alternative options in this case.
20///
21/// Storage buffers can store runtime-sized arrays, but only if they are the last field in a structure.
22///
23/// The contained data is stored in system RAM. [`write_buffer`](StorageBuffer::write_buffer) queues
24/// copying of the data from system RAM to VRAM. Storage buffers must conform to [std430 alignment/padding requirements], which
25/// is automatically enforced by this structure.
26///
27/// Other options for storing GPU-accessible data are:
28/// * [`BufferVec`](crate::render_resource::BufferVec)
29/// * [`DynamicStorageBuffer`]
30/// * [`DynamicUniformBuffer`](crate::render_resource::DynamicUniformBuffer)
31/// * [`GpuArrayBuffer`](crate::render_resource::GpuArrayBuffer)
32/// * [`RawBufferVec`](crate::render_resource::RawBufferVec)
33/// * [`Texture`](crate::render_resource::Texture)
34/// * [`UniformBuffer`](crate::render_resource::UniformBuffer)
35///
36/// [std430 alignment/padding requirements]: https://www.w3.org/TR/WGSL/#address-spaces-storage
37pub struct StorageBuffer<T: ShaderType> {
38    value: T,
39    scratch: StorageBufferWrapper<Vec<u8>>,
40    buffer: Option<Buffer>,
41    label: Option<String>,
42    changed: bool,
43    buffer_usage: BufferUsages,
44    last_written_size: Option<BufferSize>,
45}
46
47impl<T: ShaderType> From<T> for StorageBuffer<T> {
48    fn from(value: T) -> Self {
49        Self {
50            value,
51            scratch: StorageBufferWrapper::new(Vec::new()),
52            buffer: None,
53            label: None,
54            changed: false,
55            buffer_usage: BufferUsages::COPY_DST | BufferUsages::STORAGE,
56            last_written_size: None,
57        }
58    }
59}
60
61impl<T: ShaderType + Default> Default for StorageBuffer<T> {
62    fn default() -> Self {
63        Self {
64            value: T::default(),
65            scratch: StorageBufferWrapper::new(Vec::new()),
66            buffer: None,
67            label: None,
68            changed: false,
69            buffer_usage: BufferUsages::COPY_DST | BufferUsages::STORAGE,
70            last_written_size: None,
71        }
72    }
73}
74
75impl<T: ShaderType + WriteInto> StorageBuffer<T> {
76    #[inline]
77    pub fn buffer(&self) -> Option<&Buffer> {
78        self.buffer.as_ref()
79    }
80
81    #[inline]
82    pub fn binding(&self) -> Option<BindingResource<'_>> {
83        Some(BindingResource::Buffer(BufferBinding {
84            buffer: self.buffer()?,
85            offset: 0,
86            size: self.last_written_size,
87        }))
88    }
89
90    pub fn set(&mut self, value: T) {
91        self.value = value;
92    }
93
94    pub fn get(&self) -> &T {
95        &self.value
96    }
97
98    pub fn get_mut(&mut self) -> &mut T {
99        &mut self.value
100    }
101
102    pub fn set_label(&mut self, label: Option<&str>) {
103        let label = label.map(str::to_string);
104
105        if label != self.label {
106            self.changed = true;
107        }
108
109        self.label = label;
110    }
111
112    pub fn get_label(&self) -> Option<&str> {
113        self.label.as_deref()
114    }
115
116    /// Add more [`BufferUsages`] to the buffer.
117    ///
118    /// This method only allows addition of flags to the default usage flags.
119    ///
120    /// The default values for buffer usage are `BufferUsages::COPY_DST` and `BufferUsages::STORAGE`.
121    pub fn add_usages(&mut self, usage: BufferUsages) {
122        self.buffer_usage |= usage;
123        self.changed = true;
124    }
125
126    /// Queues writing of data from system RAM to VRAM using the [`RenderDevice`]
127    /// and the provided [`RenderQueue`].
128    ///
129    /// If there is no GPU-side buffer allocated to hold the data currently stored, or if a GPU-side buffer previously
130    /// allocated does not have enough capacity, a new GPU-side buffer is created.
131    pub fn write_buffer(&mut self, device: &RenderDevice, queue: &RenderQueue) {
132        self.scratch.write(&self.value).unwrap();
133
134        let capacity = self.buffer.as_deref().map(wgpu::Buffer::size).unwrap_or(0);
135        let size = self.scratch.as_ref().len() as u64;
136
137        if capacity < size || self.changed {
138            self.buffer = Some(device.create_buffer_with_data(&BufferInitDescriptor {
139                label: make_buffer_label::<Self>(&self.label),
140                usage: self.buffer_usage,
141                contents: self.scratch.as_ref(),
142            }));
143            self.changed = false;
144        } else if let Some(buffer) = &self.buffer {
145            queue.write_buffer(buffer, 0, self.scratch.as_ref());
146        }
147
148        self.last_written_size = BufferSize::new(size);
149    }
150}
151
152impl<'a, T: ShaderType + WriteInto> IntoBinding<'a> for &'a StorageBuffer<T> {
153    #[inline]
154    fn into_binding(self) -> BindingResource<'a> {
155        self.binding().expect("Failed to get buffer")
156    }
157}
158
159/// Stores data to be transferred to the GPU and made accessible to shaders as a dynamic storage buffer.
160///
161/// This is just a [`StorageBuffer`], but also allows you to set dynamic offsets.
162///
163/// Dynamic storage buffers can be made available to shaders in some combination of read/write mode, and can store large amounts
164/// of data. Note however that WebGL2 does not support storage buffers, so consider alternative options in this case. Dynamic
165/// storage buffers support multiple separate bindings at dynamic byte offsets and so have a
166/// [`push`](DynamicStorageBuffer::push) method.
167///
168/// The contained data is stored in system RAM. [`write_buffer`](DynamicStorageBuffer::write_buffer)
169/// queues copying of the data from system RAM to VRAM. The data within a storage buffer binding must conform to
170/// [std430 alignment/padding requirements]. `DynamicStorageBuffer` takes care of serializing the inner type to conform to
171/// these requirements. Each item [`push`](DynamicStorageBuffer::push)ed into this structure
172/// will additionally be aligned to meet dynamic offset alignment requirements.
173///
174/// Other options for storing GPU-accessible data are:
175/// * [`BufferVec`](crate::render_resource::BufferVec)
176/// * [`DynamicUniformBuffer`](crate::render_resource::DynamicUniformBuffer)
177/// * [`GpuArrayBuffer`](crate::render_resource::GpuArrayBuffer)
178/// * [`RawBufferVec`](crate::render_resource::RawBufferVec)
179/// * [`StorageBuffer`]
180/// * [`Texture`](crate::render_resource::Texture)
181/// * [`UniformBuffer`](crate::render_resource::UniformBuffer)
182///
183/// [std430 alignment/padding requirements]: https://www.w3.org/TR/WGSL/#address-spaces-storage
184pub struct DynamicStorageBuffer<T: ShaderType> {
185    scratch: DynamicStorageBufferWrapper<Vec<u8>>,
186    buffer: Option<Buffer>,
187    label: Option<String>,
188    changed: bool,
189    buffer_usage: BufferUsages,
190    last_written_size: Option<BufferSize>,
191    _marker: PhantomData<fn() -> T>,
192}
193
194impl<T: ShaderType> Default for DynamicStorageBuffer<T> {
195    fn default() -> Self {
196        Self {
197            scratch: DynamicStorageBufferWrapper::new(Vec::new()),
198            buffer: None,
199            label: None,
200            changed: false,
201            buffer_usage: BufferUsages::COPY_DST | BufferUsages::STORAGE,
202            last_written_size: None,
203            _marker: PhantomData,
204        }
205    }
206}
207
208impl<T: ShaderType + WriteInto> DynamicStorageBuffer<T> {
209    #[inline]
210    pub fn buffer(&self) -> Option<&Buffer> {
211        self.buffer.as_ref()
212    }
213
214    #[inline]
215    pub fn binding(&self) -> Option<BindingResource<'_>> {
216        Some(BindingResource::Buffer(BufferBinding {
217            buffer: self.buffer()?,
218            offset: 0,
219            size: self.last_written_size,
220        }))
221    }
222
223    #[inline]
224    pub fn is_empty(&self) -> bool {
225        self.scratch.as_ref().is_empty()
226    }
227
228    #[inline]
229    pub fn push(&mut self, value: T) -> u32 {
230        self.scratch.write(&value).unwrap() as u32
231    }
232
233    pub fn set_label(&mut self, label: Option<&str>) {
234        let label = label.map(str::to_string);
235
236        if label != self.label {
237            self.changed = true;
238        }
239
240        self.label = label;
241    }
242
243    pub fn get_label(&self) -> Option<&str> {
244        self.label.as_deref()
245    }
246
247    /// Add more [`BufferUsages`] to the buffer.
248    ///
249    /// This method only allows addition of flags to the default usage flags.
250    ///
251    /// The default values for buffer usage are `BufferUsages::COPY_DST` and `BufferUsages::STORAGE`.
252    pub fn add_usages(&mut self, usage: BufferUsages) {
253        self.buffer_usage |= usage;
254        self.changed = true;
255    }
256
257    #[inline]
258    pub fn write_buffer(&mut self, device: &RenderDevice, queue: &RenderQueue) {
259        let capacity = self.buffer.as_deref().map(wgpu::Buffer::size).unwrap_or(0);
260        let size = self.scratch.as_ref().len() as u64;
261
262        if capacity < size || (self.changed && size > 0) {
263            self.buffer = Some(device.create_buffer_with_data(&BufferInitDescriptor {
264                label: make_buffer_label::<Self>(&self.label),
265                usage: self.buffer_usage,
266                contents: self.scratch.as_ref(),
267            }));
268            self.changed = false;
269        } else if let Some(buffer) = &self.buffer {
270            queue.write_buffer(buffer, 0, self.scratch.as_ref());
271        }
272
273        self.last_written_size = BufferSize::new(size);
274    }
275
276    #[inline]
277    pub fn clear(&mut self) {
278        self.scratch.as_mut().clear();
279        self.scratch.set_offset(0);
280    }
281}
282
283impl<'a, T: ShaderType + WriteInto> IntoBinding<'a> for &'a DynamicStorageBuffer<T> {
284    #[inline]
285    fn into_binding(self) -> BindingResource<'a> {
286        self.binding().expect("Failed to get buffer")
287    }
288}
bevy_render/render_resource/storage_buffer.rs

bevy_render/render_resource/
storage_buffer.rs
