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/// * [`BufferVec`](crate::render_resource::BufferVec)
26/// * [`DynamicStorageBuffer`]
27/// * [`DynamicUniformBuffer`](crate::render_resource::DynamicUniformBuffer)
28/// * [`GpuArrayBuffer`](crate::render_resource::GpuArrayBuffer)
29/// * [`RawBufferVec`](crate::render_resource::RawBufferVec)
30/// * [`Texture`](crate::render_resource::Texture)
31/// * [`UniformBuffer`](crate::render_resource::UniformBuffer)
32///
33/// [std430 alignment/padding requirements]: https://www.w3.org/TR/WGSL/#address-spaces-storage
34pub struct StorageBuffer<T: ShaderType> {
35    value: T,
36    scratch: StorageBufferWrapper<Vec<u8>>,
37    buffer: Option<Buffer>,
38    label: Option<String>,
39    changed: bool,
40    buffer_usage: BufferUsages,
41    last_written_size: Option<BufferSize>,
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            last_written_size: None,
54        }
55    }
56}
57
58impl<T: ShaderType + Default> Default for StorageBuffer<T> {
59    fn default() -> Self {
60        Self {
61            value: T::default(),
62            scratch: StorageBufferWrapper::new(Vec::new()),
63            buffer: None,
64            label: None,
65            changed: false,
66            buffer_usage: BufferUsages::COPY_DST | BufferUsages::STORAGE,
67            last_written_size: None,
68        }
69    }
70}
71
72impl<T: ShaderType + WriteInto> StorageBuffer<T> {
73    #[inline]
74    pub fn buffer(&self) -> Option<&Buffer> {
75        self.buffer.as_ref()
76    }
77
78    #[inline]
79    pub fn binding(&self) -> Option<BindingResource> {
80        Some(BindingResource::Buffer(BufferBinding {
81            buffer: self.buffer()?,
82            offset: 0,
83            size: self.last_written_size,
84        }))
85    }
86
87    pub fn set(&mut self, value: T) {
88        self.value = value;
89    }
90
91    pub fn get(&self) -> &T {
92        &self.value
93    }
94
95    pub fn get_mut(&mut self) -> &mut T {
96        &mut self.value
97    }
98
99    pub fn set_label(&mut self, label: Option<&str>) {
100        let label = label.map(str::to_string);
101
102        if label != self.label {
103            self.changed = true;
104        }
105
106        self.label = label;
107    }
108
109    pub fn get_label(&self) -> Option<&str> {
110        self.label.as_deref()
111    }
112
113    /// Add more [`BufferUsages`] to the buffer.
114    ///
115    /// This method only allows addition of flags to the default usage flags.
116    ///
117    /// The default values for buffer usage are `BufferUsages::COPY_DST` and `BufferUsages::STORAGE`.
118    pub fn add_usages(&mut self, usage: BufferUsages) {
119        self.buffer_usage |= usage;
120        self.changed = true;
121    }
122
123    /// Queues writing of data from system RAM to VRAM using the [`RenderDevice`]
124    /// and the provided [`RenderQueue`].
125    ///
126    /// If there is no GPU-side buffer allocated to hold the data currently stored, or if a GPU-side buffer previously
127    /// allocated does not have enough capacity, a new GPU-side buffer is created.
128    pub fn write_buffer(&mut self, device: &RenderDevice, queue: &RenderQueue) {
129        self.scratch.write(&self.value).unwrap();
130
131        let capacity = self.buffer.as_deref().map(wgpu::Buffer::size).unwrap_or(0);
132        let size = self.scratch.as_ref().len() as u64;
133
134        if capacity < size || self.changed {
135            self.buffer = Some(device.create_buffer_with_data(&BufferInitDescriptor {
136                label: self.label.as_deref(),
137                usage: self.buffer_usage,
138                contents: self.scratch.as_ref(),
139            }));
140            self.changed = false;
141        } else if let Some(buffer) = &self.buffer {
142            queue.write_buffer(buffer, 0, self.scratch.as_ref());
143        }
144
145        self.last_written_size = BufferSize::new(size);
146    }
147}
148
149impl<'a, T: ShaderType + WriteInto> IntoBinding<'a> for &'a StorageBuffer<T> {
150    #[inline]
151    fn into_binding(self) -> BindingResource<'a> {
152        self.binding().expect("Failed to get buffer")
153    }
154}
155
156/// Stores data to be transferred to the GPU and made accessible to shaders as a dynamic storage buffer.
157///
158/// This is just a [`StorageBuffer`], but also allows you to set dynamic offsets.
159///
160/// Dynamic storage buffers can be made available to shaders in some combination of read/write mode, and can store large amounts
161/// of data. Note however that WebGL2 does not support storage buffers, so consider alternative options in this case. Dynamic
162/// storage buffers support multiple separate bindings at dynamic byte offsets and so have a
163/// [`push`](DynamicStorageBuffer::push) method.
164///
165/// The contained data is stored in system RAM. [`write_buffer`](DynamicStorageBuffer::write_buffer)
166/// queues copying of the data from system RAM to VRAM. The data within a storage buffer binding must conform to
167/// [std430 alignment/padding requirements]. `DynamicStorageBuffer` takes care of serializing the inner type to conform to
168/// these requirements. Each item [`push`](DynamicStorageBuffer::push)ed into this structure
169/// will additionally be aligned to meet dynamic offset alignment requirements.
170///
171/// Other options for storing GPU-accessible data are:
172/// * [`BufferVec`](crate::render_resource::BufferVec)
173/// * [`DynamicUniformBuffer`](crate::render_resource::DynamicUniformBuffer)
174/// * [`GpuArrayBuffer`](crate::render_resource::GpuArrayBuffer)
175/// * [`RawBufferVec`](crate::render_resource::RawBufferVec)
176/// * [`StorageBuffer`]
177/// * [`Texture`](crate::render_resource::Texture)
178/// * [`UniformBuffer`](crate::render_resource::UniformBuffer)
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
