bevy_render/render_resource/
storage_buffer.rs

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