bevy_render/render_resource/
storage_buffer.rs1use 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
13pub 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 pub fn add_usages(&mut self, usage: BufferUsages) {
115 self.buffer_usage |= usage;
116 self.changed = true;
117 }
118
119 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
153pub 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 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}