bevy_render/render_resource/
storage_buffer.rs

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

bevy_render/render_resource/
storage_buffer.rs
