bevy_render/render_resource/
uniform_buffer.rs1use core::{marker::PhantomData, num::NonZero};
2
3use crate::{
4 render_resource::{make_buffer_label, Buffer},
5 renderer::{RenderDevice, RenderQueue},
6};
7use encase::{
8 internal::{AlignmentValue, BufferMut, WriteInto},
9 DynamicUniformBuffer as DynamicUniformBufferWrapper, ShaderType,
10 UniformBuffer as UniformBufferWrapper,
11};
12use wgpu::{
13 util::BufferInitDescriptor, BindingResource, BufferBinding, BufferDescriptor, BufferUsages,
14};
15
16use super::IntoBinding;
17
18pub struct UniformBuffer<T: ShaderType> {
40 value: T,
41 scratch: UniformBufferWrapper<Vec<u8>>,
42 buffer: Option<Buffer>,
43 label: Option<String>,
44 changed: bool,
45 buffer_usage: BufferUsages,
46}
47
48impl<T: ShaderType> From<T> for UniformBuffer<T> {
49 fn from(value: T) -> Self {
50 Self {
51 value,
52 scratch: UniformBufferWrapper::new(Vec::new()),
53 buffer: None,
54 label: None,
55 changed: false,
56 buffer_usage: BufferUsages::COPY_DST | BufferUsages::UNIFORM,
57 }
58 }
59}
60
61impl<T: ShaderType + Default> Default for UniformBuffer<T> {
62 fn default() -> Self {
63 Self {
64 value: T::default(),
65 scratch: UniformBufferWrapper::new(Vec::new()),
66 buffer: None,
67 label: None,
68 changed: false,
69 buffer_usage: BufferUsages::COPY_DST | BufferUsages::UNIFORM,
70 }
71 }
72}
73
74impl<T: ShaderType + WriteInto> UniformBuffer<T> {
75 #[inline]
76 pub fn buffer(&self) -> Option<&Buffer> {
77 self.buffer.as_ref()
78 }
79
80 #[inline]
81 pub fn binding(&self) -> Option<BindingResource<'_>> {
82 Some(BindingResource::Buffer(
83 self.buffer()?.as_entire_buffer_binding(),
84 ))
85 }
86
87 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 pub fn add_usages(&mut self, usage: BufferUsages) {
120 self.buffer_usage |= usage;
121 self.changed = true;
122 }
123
124 pub fn write_buffer(&mut self, device: &RenderDevice, queue: &RenderQueue) {
130 self.scratch.write(&self.value).unwrap();
131
132 if self.changed || self.buffer.is_none() {
133 self.buffer = Some(device.create_buffer_with_data(&BufferInitDescriptor {
134 label: make_buffer_label::<Self>(&self.label),
135 usage: self.buffer_usage,
136 contents: self.scratch.as_ref(),
137 }));
138 self.changed = false;
139 } else if let Some(buffer) = &self.buffer {
140 queue.write_buffer(buffer, 0, self.scratch.as_ref());
141 }
142 }
143}
144
145impl<'a, T: ShaderType + WriteInto> IntoBinding<'a> for &'a UniformBuffer<T> {
146 #[inline]
147 fn into_binding(self) -> BindingResource<'a> {
148 self.buffer()
149 .expect("Failed to get buffer")
150 .as_entire_buffer_binding()
151 .into_binding()
152 }
153}
154
155pub struct DynamicUniformBuffer<T: ShaderType> {
177 scratch: DynamicUniformBufferWrapper<Vec<u8>>,
178 buffer: Option<Buffer>,
179 label: Option<String>,
180 changed: bool,
181 buffer_usage: BufferUsages,
182 _marker: PhantomData<fn() -> T>,
183}
184
185impl<T: ShaderType> Default for DynamicUniformBuffer<T> {
186 fn default() -> Self {
187 Self {
188 scratch: DynamicUniformBufferWrapper::new(Vec::new()),
189 buffer: None,
190 label: None,
191 changed: false,
192 buffer_usage: BufferUsages::COPY_DST | BufferUsages::UNIFORM,
193 _marker: PhantomData,
194 }
195 }
196}
197
198impl<T: ShaderType> DynamicUniformBuffer<T> {
199 pub fn new_with_alignment(alignment: u64) -> Self {
200 Self {
201 scratch: DynamicUniformBufferWrapper::new_with_alignment(Vec::new(), alignment),
202 buffer: None,
203 label: None,
204 changed: false,
205 buffer_usage: BufferUsages::COPY_DST | BufferUsages::UNIFORM,
206 _marker: PhantomData,
207 }
208 }
209
210 #[inline]
211 pub fn buffer(&self) -> Option<&Buffer> {
212 self.buffer.as_ref()
213 }
214
215 #[inline]
216 pub fn binding(&self) -> Option<BindingResource<'_>> {
217 Some(BindingResource::Buffer(BufferBinding {
218 buffer: self.buffer()?,
219 offset: 0,
220 size: Some(T::min_size()),
221 }))
222 }
223
224 #[inline]
225 pub fn is_empty(&self) -> bool {
226 self.scratch.as_ref().is_empty()
227 }
228
229 pub fn set_label(&mut self, label: Option<&str>) {
230 let label = label.map(str::to_string);
231
232 if label != self.label {
233 self.changed = true;
234 }
235
236 self.label = label;
237 }
238
239 pub fn get_label(&self) -> Option<&str> {
240 self.label.as_deref()
241 }
242
243 pub fn add_usages(&mut self, usage: BufferUsages) {
249 self.buffer_usage |= usage;
250 self.changed = true;
251 }
252
253 #[inline]
269 pub fn get_writer<'a>(
270 &'a mut self,
271 max_count: usize,
272 device: &RenderDevice,
273 queue: &'a RenderQueue,
274 ) -> Option<DynamicUniformBufferWriter<T>> {
275 let alignment = if cfg!(target_abi = "sim") {
276 AlignmentValue::new(256)
281 } else {
282 AlignmentValue::new(device.limits().min_uniform_buffer_offset_alignment as u64)
283 };
284
285 let mut capacity = self.buffer.as_deref().map(wgpu::Buffer::size).unwrap_or(0);
286 let size = alignment
287 .round_up(T::min_size().get())
288 .checked_mul(max_count as u64)
289 .unwrap();
290
291 if capacity < size || (self.changed && size > 0) {
292 let buffer = device.create_buffer(&BufferDescriptor {
293 label: make_buffer_label::<Self>(&self.label),
294 usage: self.buffer_usage,
295 size,
296 mapped_at_creation: false,
297 });
298 capacity = buffer.size();
299 self.buffer = Some(buffer);
300 self.changed = false;
301 }
302
303 if let Some(buffer) = self.buffer.as_deref() {
304 let buffer_view =
305 queue.write_buffer_with(buffer, 0, NonZero::<u64>::new(buffer.size())?)?;
306 Some(DynamicUniformBufferWriter {
307 buffer: encase::DynamicUniformBuffer::new_with_alignment(
308 QueueWriteBufferViewWrapper {
309 capacity: capacity as usize,
310 buffer_view,
311 },
312 alignment.get(),
313 ),
314 _marker: PhantomData,
315 })
316 } else {
317 None
318 }
319 }
320
321 #[inline]
327 pub fn write_buffer(&mut self, device: &RenderDevice, queue: &RenderQueue) {
328 let capacity = self.buffer.as_deref().map(wgpu::Buffer::size).unwrap_or(0);
329 let size = self.scratch.as_ref().len() as u64;
330
331 if capacity < size || (self.changed && size > 0) {
332 self.buffer = Some(device.create_buffer_with_data(&BufferInitDescriptor {
333 label: self.label.as_deref(),
334 usage: self.buffer_usage,
335 contents: self.scratch.as_ref(),
336 }));
337 self.changed = false;
338 } else if let Some(buffer) = &self.buffer {
339 queue.write_buffer(buffer, 0, self.scratch.as_ref());
340 }
341 }
342
343 #[inline]
344 pub fn clear(&mut self) {
345 self.scratch.as_mut().clear();
346 self.scratch.set_offset(0);
347 }
348}
349
350impl<T: ShaderType + WriteInto> DynamicUniformBuffer<T> {
351 #[inline]
353 pub fn push(&mut self, value: &T) -> u32 {
354 self.scratch.write(value).unwrap() as u32
355 }
356}
357
358pub struct DynamicUniformBufferWriter<T> {
362 buffer: encase::DynamicUniformBuffer<QueueWriteBufferViewWrapper>,
363 _marker: PhantomData<fn() -> T>,
364}
365
366impl<T: ShaderType + WriteInto> DynamicUniformBufferWriter<T> {
367 pub fn write(&mut self, value: &T) -> u32 {
368 self.buffer.write(value).unwrap() as u32
369 }
370}
371
372struct QueueWriteBufferViewWrapper {
375 buffer_view: wgpu::QueueWriteBufferView,
376 capacity: usize,
379}
380
381impl BufferMut for QueueWriteBufferViewWrapper {
382 #[inline]
383 fn capacity(&self) -> usize {
384 self.capacity
385 }
386
387 #[inline]
388 fn write<const N: usize>(&mut self, offset: usize, val: &[u8; N]) {
389 self.buffer_view
390 .slice(offset..offset + val.len())
391 .copy_from_slice(val);
392 }
393
394 #[inline]
395 fn write_slice(&mut self, offset: usize, val: &[u8]) {
396 self.buffer_view
397 .slice(offset..offset + val.len())
398 .copy_from_slice(val);
399 }
400}
401
402impl<'a, T: ShaderType + WriteInto> IntoBinding<'a> for &'a DynamicUniformBuffer<T> {
403 #[inline]
404 fn into_binding(self) -> BindingResource<'a> {
405 self.binding().unwrap()
406 }
407}