bevy_render/
storage.rs

1//! [`ShaderBuffer`], an asset that encapsulates arbitrary data that will be
2//! extracted and uploaded to the GPU for use in shaders.
3
4use alloc::borrow::Cow;
5
6use crate::{
7    render_asset::{AssetExtractionError, PrepareAssetError, RenderAsset, RenderAssetPlugin},
8    render_resource::{Buffer, BufferUsages},
9    renderer::{RenderDevice, RenderQueue},
10    Render, RenderApp, RenderSystems,
11};
12use bevy_app::{App, Plugin};
13use bevy_asset::{Asset, AssetApp, AssetId, RenderAssetUsages};
14use bevy_derive::{Deref, DerefMut};
15use bevy_ecs::{
16    resource::Resource,
17    schedule::IntoScheduleConfigs as _,
18    system::{
19        lifetimeless::{SRes, SResMut},
20        ResMut, SystemParamItem,
21    },
22};
23use bevy_platform::collections::{AlignedVec, HashSet};
24use bevy_reflect::{prelude::ReflectDefault, Reflect};
25use bevy_utils::default;
26use wgpu_types::BufferDescriptor;
27
28/// Adds a [`ShaderBuffer`] as an asset that is extracted and uploaded to the
29/// GPU.
30#[derive(Default)]
31pub struct StoragePlugin;
32
33impl Plugin for StoragePlugin {
34    fn build(&self, app: &mut App) {
35        app.add_plugins(RenderAssetPlugin::<GpuShaderBuffer>::default())
36            .init_asset::<ShaderBuffer>()
37            .register_asset_reflect::<ShaderBuffer>();
38
39        let Some(render_app) = app.get_sub_app_mut(RenderApp) else {
40            return;
41        };
42        render_app
43            .init_resource::<RenderChangedShaderBuffers>()
44            .add_systems(
45                Render,
46                clear_changed_shader_buffers.in_set(RenderSystems::Cleanup),
47            );
48    }
49}
50
51/// A storage buffer that is prepared as a [`RenderAsset`] and uploaded to the GPU.
52///
53/// This buffer primarily exists in order to be embedded into a material that
54/// implements the [`bevy_render_macros::AsBindGroup`] trait. Compared to
55/// embedding a raw [`Buffer`], [`ShaderBuffer`] has the advantage that the
56/// buffer can be resized without regenerating the materials that embed it.
57#[derive(Asset, Reflect, Debug, Clone)]
58#[reflect(opaque)]
59#[reflect(Default, Debug, Clone)]
60pub struct ShaderBuffer {
61    /// Optional data used to initialize the buffer, as well as the buffer's size.
62    pub data: ShaderBufferData,
63    /// A label that can be used to identify this buffer in a debugger.
64    pub label: Cow<'static, str>,
65    /// How this buffer can legally be used.
66    pub buffer_usage: BufferUsages,
67    /// The asset usage of the storage buffer.
68    pub asset_usage: RenderAssetUsages,
69    /// Whether this buffer should be copied on the GPU when resized.
70    /// The buffer should have `BufferUsages::COPY_SRC | BufferUsages::COPY_DST` usages to be copyable.
71    ///
72    /// This has no effect if data is [`ShaderBufferData::Initialized`], where GPU buffer is always populated
73    /// using CPU data when GPU buffer is resized.
74    pub copy_on_resize: bool,
75}
76
77/// Optional data used to initialize a [`ShaderBuffer`].
78///
79/// This also includes the buffer's size in bytes.
80/// The buffer size must be a multiple of 4 as required by wgpu.
81/// Zero size is allowed but can't be used as binding resource.
82#[derive(Reflect, Debug, Clone)]
83#[reflect(Default, Debug, Clone)]
84#[reflect(opaque)]
85pub enum ShaderBufferData {
86    /// The buffer will be uninitialized when created and has the given size in
87    /// bytes.
88    Uninitialized(wgpu_types::BufferAddress),
89    /// The buffer will be initialized with the given data.
90    ///
91    /// The size of the buffer is equal to `buffer_size`, not the size of `data`.
92    Initialized {
93        data: AlignedVec,
94        buffer_size: wgpu_types::BufferAddress,
95    },
96}
97
98impl Default for ShaderBuffer {
99    fn default() -> Self {
100        Self {
101            data: ShaderBufferData::Uninitialized(0),
102            label: Cow::Borrowed("shader buffer"),
103            buffer_usage: BufferUsages::STORAGE | BufferUsages::COPY_SRC | BufferUsages::COPY_DST,
104            asset_usage: RenderAssetUsages::default(),
105            copy_on_resize: false,
106        }
107    }
108}
109
110impl Default for ShaderBufferData {
111    fn default() -> Self {
112        ShaderBufferData::Uninitialized(0)
113    }
114}
115
116impl ShaderBuffer {
117    /// Creates a new initialized storage buffer with the given data and asset usage, with alignment `align_of::<T>()`.
118    pub fn new<T: bytemuck::NoUninit>(data: Vec<T>, asset_usage: RenderAssetUsages) -> Self {
119        let data = AlignedVec::from(data);
120        let buffer_size = data.len() as u64;
121        ShaderBuffer {
122            data: ShaderBufferData::Initialized { data, buffer_size },
123            asset_usage,
124            ..default()
125        }
126    }
127
128    /// Creates a new uninitialized storage buffer with the given size and asset usage.
129    pub fn with_size(size: u64, asset_usage: RenderAssetUsages) -> Self {
130        ShaderBuffer {
131            data: ShaderBufferData::Uninitialized(size),
132            asset_usage,
133            ..default()
134        }
135    }
136
137    /// Clear [`Self::data`] if it is [`ShaderBufferData::Initialized`] with its capacity and buffer size reserved.
138    pub fn clear(&mut self) {
139        if let ShaderBufferData::Initialized { data, .. } = &mut self.data {
140            data.clear();
141        }
142    }
143
144    /// Extends the data with a slice of [`bytemuck::NoUninit`].
145    /// If [`Self::data`] is uninitialized, it will be initialized with alignment `align_of::<T>()`
146    ///
147    /// [`ShaderBufferData::Initialized::buffer_size`] will be set to the data length.
148    pub fn extend_from_slice<T>(&mut self, values: &[T])
149    where
150        T: bytemuck::NoUninit,
151    {
152        let data = core::mem::take(&mut self.data);
153        let mut data = match data {
154            ShaderBufferData::Uninitialized(_) => AlignedVec::new(align_of::<T>()),
155            ShaderBufferData::Initialized { data, .. } => data,
156        };
157        data.extend_from_slice(bytemuck::cast_slice(values));
158        let buffer_size = data.len() as u64;
159        self.data = ShaderBufferData::Initialized { data, buffer_size };
160    }
161
162    /// Extends the data with an iterator of [`bytemuck::NoUninit`].
163    /// If [`Self::data`] is uninitialized, it will be initialized with alignment `align_of::<T>()`.
164    ///
165    /// [`ShaderBufferData::Initialized::buffer_size`] will be set to the data length.
166    pub fn extend<T>(&mut self, values: impl IntoIterator<Item = T>)
167    where
168        T: bytemuck::NoUninit,
169    {
170        let values = values.into_iter();
171        let data = core::mem::take(&mut self.data);
172        let mut data = match data {
173            ShaderBufferData::Uninitialized(_) => AlignedVec::new(align_of::<T>()),
174            ShaderBufferData::Initialized { data, .. } => data,
175        };
176        data.reserve(values.size_hint().0 * size_of::<T>());
177        for value in values {
178            data.extend_from_slice(bytemuck::bytes_of(&value));
179        }
180        let buffer_size = data.len() as u64;
181        self.data = ShaderBufferData::Initialized { data, buffer_size };
182    }
183
184    /// Casts and returns a slice of `T` of [`ShaderBufferData::Initialized`],
185    /// or returns `None` if it's [`ShaderBufferData::Uninitialized`]
186    ///
187    /// Panics:
188    /// * If `T` has a greater alignment requirement and the `AlignedVec` isn't aligned.
189    /// * If the size of `AlignedVec` is not a multiple of `size_of::<T>()`
190    pub fn cast_slice<T: bytemuck::AnyBitPattern>(&self) -> Option<&[T]> {
191        match &self.data {
192            ShaderBufferData::Uninitialized(_) => None,
193            ShaderBufferData::Initialized { data, .. } => Some(data.cast_slice()),
194        }
195    }
196
197    /// Casts and returns a mutable slice of `T` of [`ShaderBufferData::Initialized`],
198    /// or returns `None` if it's [`ShaderBufferData::Uninitialized`]
199    ///
200    /// Panics:
201    /// * If `T` has a greater alignment requirement than the `AlignedVec`.
202    /// * If the size of `AlignedVec` is not a multiple of `size_of::<T>()`
203    pub fn cast_slice_mut<T: bytemuck::NoUninit + bytemuck::AnyBitPattern>(
204        &mut self,
205    ) -> Option<&mut [T]> {
206        match &mut self.data {
207            ShaderBufferData::Uninitialized(_) => None,
208            ShaderBufferData::Initialized { data, .. } => Some(data.cast_slice_mut()),
209        }
210    }
211
212    /// Resizes the CPU data and buffer to the new size.
213    ///
214    /// If CPU data is present, the GPU buffer will be re-populated using CPU data.
215    /// Any remaining part that lacks CPU data is implicitly zero-initialized by wgpu.
216    ///
217    /// If CPU data is not present, the entire GPU buffer will be reallocated and implicitly zero-initialized by wgpu.
218    /// If `copy_on_resize` is true, previous buffer will attempt to be copied to this buffer.
219    ///
220    /// CPU data is truncated or zero-extended, too.
221    pub fn resize(&mut self, new_size: wgpu_types::BufferAddress) {
222        match self.data {
223            ShaderBufferData::Initialized {
224                ref mut data,
225                ref mut buffer_size,
226            } => {
227                data.resize(new_size as usize, 0);
228                *buffer_size = new_size;
229            }
230            ShaderBufferData::Uninitialized(ref mut size) => {
231                *size = new_size;
232            }
233        }
234    }
235
236    /// Resizes the GPU buffer to the new size.
237    ///
238    /// If CPU data is present, the GPU buffer will be re-populated using CPU data.
239    /// Any remaining part that lacks CPU data is implicitly zero-initialized by wgpu.
240    ///
241    /// If CPU data is not present, the entire GPU buffer will be reallocated and implicitly zero-initialized by wgpu.
242    /// If `copy_on_resize` is true, previous buffer will attempt to be copied to this buffer.
243    ///
244    /// CPU data is unchanged.
245    pub fn resize_buffer(&mut self, new_size: wgpu_types::BufferAddress) {
246        match self.data {
247            ShaderBufferData::Initialized {
248                ref mut buffer_size,
249                ..
250            } => {
251                *buffer_size = new_size;
252            }
253            ShaderBufferData::Uninitialized(ref mut size) => {
254                *size = new_size;
255            }
256        }
257    }
258
259    /// Returns the size of the buffer in bytes.
260    pub fn buffer_size(&self) -> wgpu_types::BufferAddress {
261        match self.data {
262            ShaderBufferData::Initialized { buffer_size, .. } => buffer_size,
263            ShaderBufferData::Uninitialized(len) => len,
264        }
265    }
266}
267
268impl<T: bytemuck::NoUninit> From<Vec<T>> for ShaderBuffer {
269    /// Creates a new initialized storage buffer with the given data, with alignment `align_of::<T>()`.
270    fn from(value: Vec<T>) -> Self {
271        Self::new(value, Default::default())
272    }
273}
274
275/// A render-world resource that stores the IDs of [`ShaderBuffer`]s that have
276/// been updated to point at a different buffer.
277///
278/// The raw underlying buffer that a [`ShaderBuffer`] points to may change from
279/// frame to frame. This will happen, for example, if the buffer represents a
280/// CPU-managed vector that might grow. When this happens, the material bind
281/// group allocator must invalidate any cached bind groups that referred to the
282/// old buffer. This resource tracks those modified buffers to enable this
283/// invalidation to happen.
284///
285/// Note that a [`ShaderBuffer`] will only be in this set if the *identity* of
286/// the buffer that it wraps changed. If only the *contents* of the buffer
287/// changed since last frame, then bind groups don't need to be updated, and the
288/// shader buffer won't be present in this set.
289#[derive(Resource, Default, Deref, DerefMut)]
290pub struct RenderChangedShaderBuffers(pub HashSet<AssetId<ShaderBuffer>>);
291
292/// A storage buffer that is prepared as a [`RenderAsset`] and uploaded to the GPU.
293pub struct GpuShaderBuffer {
294    /// The raw GPU buffer.
295    pub buffer: Buffer,
296    /// A debugging label to identify the buffer.
297    pub label: Cow<'static, str>,
298    /// The allowable render usages of the buffer.
299    pub buffer_usage: BufferUsages,
300    /// Whether the buffer contains data that must be preserved.
301    pub had_data: bool,
302}
303
304impl RenderAsset for GpuShaderBuffer {
305    type SourceAsset = ShaderBuffer;
306    type Param = (
307        SRes<RenderDevice>,
308        SRes<RenderQueue>,
309        SResMut<RenderChangedShaderBuffers>,
310    );
311
312    fn asset_usage(source_asset: &Self::SourceAsset) -> RenderAssetUsages {
313        source_asset.asset_usage
314    }
315
316    fn take_gpu_data(
317        source: &mut Self::SourceAsset,
318        previous_gpu_asset: Option<&Self>,
319    ) -> Result<Self::SourceAsset, AssetExtractionError> {
320        let len = source.buffer_size();
321        let data = core::mem::replace(&mut source.data, ShaderBufferData::Uninitialized(len));
322
323        let valid_upload = matches!(data, ShaderBufferData::Initialized { .. })
324            || previous_gpu_asset.is_none_or(|prev| !prev.had_data);
325
326        valid_upload
327            .then(|| Self::SourceAsset {
328                data,
329                ..source.clone()
330            })
331            .ok_or(AssetExtractionError::AlreadyExtracted)
332    }
333
334    fn prepare_asset(
335        source_asset: Self::SourceAsset,
336        asset_id: AssetId<Self::SourceAsset>,
337        &mut (
338            ref render_device,
339            ref render_queue,
340            ref mut changed_shader_buffers,
341        ): &mut SystemParamItem<Self::Param>,
342        previous_asset: Option<&Self>,
343    ) -> Result<Self, PrepareAssetError<Self::SourceAsset>> {
344        let had_data = matches!(source_asset.data, ShaderBufferData::Initialized { .. });
345
346        let buffer = if let Some(prev) = previous_asset
347            && prev.buffer.size() == source_asset.buffer_size()
348            && prev.buffer.usage() == source_asset.buffer_usage
349            && *prev.label == *source_asset.label
350            && (!had_data || source_asset.buffer_usage.contains(BufferUsages::COPY_DST))
351        {
352            if let ShaderBufferData::Initialized { ref data, .. } = source_asset.data {
353                render_queue.write_buffer(
354                    &prev.buffer,
355                    0,
356                    &data[..((source_asset.buffer_size() as usize).min(data.len()))],
357                );
358            }
359            prev.buffer.clone()
360        } else if let ShaderBufferData::Initialized { data, buffer_size } = source_asset.data {
361            changed_shader_buffers.insert(asset_id);
362            let desc = BufferDescriptor {
363                label: Some(&*source_asset.label),
364                usage: source_asset.buffer_usage,
365                size: buffer_size,
366                mapped_at_creation: buffer_size != 0,
367            };
368
369            let buffer = render_device.create_buffer(&desc);
370
371            // Skip mapping if the buffer is zero sized
372            if buffer_size != 0 {
373                // Upload at most `buffer_size` bytes. If the data is shorter, the
374                // remaining bytes stay zero-initialized; if it's longer, the tail
375                // is truncated.
376                let upload_len = (buffer_size as usize).min(data.len());
377                buffer
378                    .get_mapped_range_mut(..upload_len as u64)
379                    .unwrap()
380                    .copy_from_slice(&data[..upload_len]);
381                buffer.unmap();
382            }
383            buffer
384        } else {
385            changed_shader_buffers.insert(asset_id);
386            let new_buffer = render_device.create_buffer(&BufferDescriptor {
387                label: Some(&*source_asset.label),
388                size: source_asset.buffer_size(),
389                usage: source_asset.buffer_usage,
390                mapped_at_creation: false,
391            });
392            if source_asset.copy_on_resize
393                && let Some(previous) = previous_asset
394                && previous.buffer.usage().contains(BufferUsages::COPY_SRC)
395                && source_asset.buffer_usage.contains(BufferUsages::COPY_DST)
396            {
397                let copy_size = source_asset.buffer_size().min(previous.buffer.size());
398                let mut encoder =
399                    render_device.create_command_encoder(&wgpu::CommandEncoderDescriptor {
400                        label: Some("copy_buffer_on_resize"),
401                    });
402                encoder.copy_buffer_to_buffer(&previous.buffer, 0, &new_buffer, 0, copy_size);
403                render_queue.submit([encoder.finish()]);
404            }
405            new_buffer
406        };
407
408        Ok(GpuShaderBuffer {
409            buffer,
410            label: source_asset.label,
411            buffer_usage: source_asset.buffer_usage,
412            had_data,
413        })
414    }
415}
416
417/// A render-world system that clears out the [`RenderChangedShaderBuffers`]
418/// resource in preparation for a new frame.
419fn clear_changed_shader_buffers(mut changed_shader_buffers: ResMut<RenderChangedShaderBuffers>) {
420    changed_shader_buffers.clear();
421}
422
423#[cfg(test)]
424mod tests {
425    use super::*;
426    use alloc::vec;
427
428    use bevy_ecs::{
429        system::{lifetimeless::SRes, SystemState},
430        world::World,
431    };
432
433    use bevy_asset::Assets;
434
435    use crate::test_utils::create_dummy_device;
436
437    /// Runs the extraction step of the [`RenderAsset`] pipeline on `source` and
438    /// returns the extracted asset.
439    fn extract(
440        source: &mut ShaderBuffer,
441        previous_gpu_asset: Option<&GpuShaderBuffer>,
442    ) -> ShaderBuffer {
443        GpuShaderBuffer::take_gpu_data(source, previous_gpu_asset)
444            .expect("shader buffer should be extractable")
445    }
446
447    /// Creates a GPU buffer from an extracted [`ShaderBuffer`] using the given
448    /// noop wgpu device (no real GPU required), optionally reusing
449    /// `previous_asset`. The same device must be used across prepares when GPU
450    /// buffers from a previous prepare are passed in.
451    ///
452    /// Returns the prepared GPU buffer along with the
453    /// [`RenderChangedShaderBuffers`] resource, so tests can verify which asset
454    /// ids were recorded as needing their bind groups invalidated.
455    fn prepare(
456        asset_id: AssetId<ShaderBuffer>,
457        extracted: ShaderBuffer,
458        previous_asset: Option<&GpuShaderBuffer>,
459        device: &RenderDevice,
460        queue: &RenderQueue,
461    ) -> (GpuShaderBuffer, RenderChangedShaderBuffers) {
462        let mut world = World::new();
463        world.insert_resource(device.clone());
464        world.insert_resource(queue.clone());
465        world.insert_resource(RenderChangedShaderBuffers::default());
466        let mut system_state = SystemState::<(
467            SRes<RenderDevice>,
468            SRes<RenderQueue>,
469            SResMut<RenderChangedShaderBuffers>,
470        )>::new(&mut world);
471        let gpu_buffer = {
472            let mut params = system_state.get_mut(&mut world).expect(
473                "RenderDevice, RenderQueue and RenderChangedShaderBuffers resources should be present",
474            );
475            GpuShaderBuffer::prepare_asset(extracted, asset_id, &mut params, previous_asset)
476                .expect("shader buffer should be prepared successfully")
477        };
478        let changed_buffer = world
479            .remove_resource::<RenderChangedShaderBuffers>()
480            .expect("RenderChangedShaderBuffers resource should be present");
481        (gpu_buffer, changed_buffer)
482    }
483
484    /// Runs the full extract + prepare pipeline on a device shared with the
485    /// given `previous_asset`.
486    fn extract_and_prepare(
487        source: &mut ShaderBuffer,
488        previous_asset: Option<&GpuShaderBuffer>,
489        device: &RenderDevice,
490        queue: &RenderQueue,
491    ) -> GpuShaderBuffer {
492        let extracted = extract(source, previous_asset);
493        let mut assets = Assets::<ShaderBuffer>::default();
494        let asset_id = assets.add(extracted.clone()).id();
495        prepare(asset_id, extracted, previous_asset, device, queue).0
496    }
497
498    /// Runs the full pipeline on a fresh dummy device, for tests that don't
499    /// chain multiple prepares together.
500    fn extract_and_prepare_on_new_device(
501        source: &mut ShaderBuffer,
502        previous_asset: Option<&GpuShaderBuffer>,
503    ) -> GpuShaderBuffer {
504        let (device, queue) = create_dummy_device();
505        extract_and_prepare(source, previous_asset, &device, &queue)
506    }
507
508    /// Extracts a buffer with data and uploads it to the GPU, verifying that a
509    /// buffer of `buffer_size()` bytes is created and that extraction leaves the
510    /// source asset uninitialized with its size preserved.
511    #[test]
512    fn extract_and_create_gpu_buffer_from_data() {
513        let mut source = ShaderBuffer::new(
514            vec![1u32, 2, 3],
515            RenderAssetUsages::MAIN_WORLD | RenderAssetUsages::RENDER_WORLD,
516        );
517        let size = 3 * size_of::<u32>() as u64;
518
519        // The buffer size defaults to the data length (3 * 4 = 12 bytes).
520        assert_eq!(source.buffer_size(), size);
521        assert_eq!(source.cast_slice::<u32>(), Some(&[1, 2, 3][..]));
522
523        let gpu = extract_and_prepare_on_new_device(&mut source, None);
524
525        // Extraction moved the CPU data out of the source asset, leaving it
526        // uninitialized but keeping its buffer size.
527        assert!(matches!(source.data, ShaderBufferData::Uninitialized(s) if s == size));
528
529        // The GPU buffer has the buffer size and all extracted data uploaded.
530        assert_eq!(gpu.buffer.size(), size);
531        assert!(gpu.had_data);
532        assert_eq!(gpu.label, source.label);
533        assert_eq!(gpu.buffer_usage, source.buffer_usage);
534    }
535
536    /// Verifies that an explicitly larger buffer size is respected: the GPU
537    /// buffer is created with `buffer_size` bytes, even though the CPU data is
538    /// shorter.
539    #[test]
540    fn create_gpu_buffer_with_buffer_size_larger_than_data() {
541        let mut source = ShaderBuffer::new(vec![1u32], RenderAssetUsages::default());
542        // Grow the GPU buffer without touching the CPU data.
543        source.resize_buffer(64);
544        assert_eq!(source.buffer_size(), 64);
545        assert_eq!(source.cast_slice::<u32>(), Some(&[1][..]));
546
547        let gpu = extract_and_prepare_on_new_device(&mut source, None);
548
549        assert_eq!(gpu.buffer.size(), 64);
550        assert!(gpu.had_data);
551    }
552
553    /// Verifies that an explicitly smaller buffer size is respected: the GPU
554    /// buffer is created with `buffer_size` bytes and only the first
555    /// `buffer_size` bytes of the CPU data are uploaded.
556    #[test]
557    fn create_gpu_buffer_with_buffer_size_smaller_than_data() {
558        let mut source = ShaderBuffer::new(vec![1u32, 2, 3, 4], RenderAssetUsages::default());
559        source.resize_buffer(8);
560        assert_eq!(source.buffer_size(), 8);
561
562        let gpu = extract_and_prepare_on_new_device(&mut source, None);
563
564        assert_eq!(gpu.buffer.size(), 8);
565        assert!(gpu.had_data);
566    }
567
568    /// Verifies that zero-sized buffers are created without attempting to map
569    /// them, both for initialized and uninitialized sources.
570    #[test]
571    fn create_zero_sized_gpu_buffer() {
572        // An initialized buffer whose data is empty.
573        let mut source = ShaderBuffer::new(Vec::<u32>::new(), RenderAssetUsages::default());
574        assert_eq!(source.buffer_size(), 0);
575
576        let gpu = extract_and_prepare_on_new_device(&mut source, None);
577        assert_eq!(gpu.buffer.size(), 0);
578        assert!(gpu.had_data);
579
580        // An uninitialized buffer with size zero.
581        let mut source = ShaderBuffer::with_size(0, RenderAssetUsages::default());
582        let gpu = extract_and_prepare_on_new_device(&mut source, None);
583        assert_eq!(gpu.buffer.size(), 0);
584        assert!(!gpu.had_data);
585    }
586
587    /// Verifies that an uninitialized buffer creates an uninitialized GPU buffer
588    /// of the requested size and is reported as having no data.
589    #[test]
590    fn create_uninitialized_gpu_buffer() {
591        let mut source = ShaderBuffer::with_size(1024, RenderAssetUsages::default());
592        assert_eq!(source.buffer_size(), 1024);
593
594        let gpu = extract_and_prepare_on_new_device(&mut source, None);
595
596        assert_eq!(gpu.buffer.size(), 1024);
597        assert!(!gpu.had_data);
598    }
599
600    /// Verifies the extraction guards: a buffer whose CPU data has already been
601    /// moved to the render world can still be extracted if there is no previous
602    /// GPU asset carrying data, but is rejected once a previous GPU asset
603    /// contained data.
604    #[test]
605    fn extraction_rejects_buffer_whose_data_would_be_lost() {
606        let mut source = ShaderBuffer::new(vec![1u32], RenderAssetUsages::default());
607        assert!(GpuShaderBuffer::take_gpu_data(&mut source, None).is_ok());
608        assert!(matches!(source.data, ShaderBufferData::Uninitialized(4)));
609
610        // The source no longer holds data, but with no previous GPU asset the
611        // GPU buffer can simply be created uninitialized.
612        assert!(GpuShaderBuffer::take_gpu_data(&mut source, None).is_ok());
613
614        // An uninitialized buffer is rejected when the previous GPU asset had
615        // data, since re-preparing it would silently drop that data.
616        let mut source = ShaderBuffer::with_size(4, RenderAssetUsages::default());
617        let previous = GpuShaderBuffer {
618            buffer: create_dummy_device().0.create_buffer(&BufferDescriptor {
619                label: Some("previous"),
620                size: 4,
621                usage: BufferUsages::STORAGE,
622                mapped_at_creation: false,
623            }),
624            label: Cow::Borrowed("shader buffer"),
625            buffer_usage: BufferUsages::STORAGE,
626            had_data: true,
627        };
628        assert!(matches!(
629            GpuShaderBuffer::take_gpu_data(&mut source, Some(&previous)),
630            Err(AssetExtractionError::AlreadyExtracted)
631        ));
632    }
633
634    /// Verifies that an unchanged buffer reuses the existing GPU buffer instead
635    /// of allocating a new one, and that changing the buffer size or losing the
636    /// data invalidates the reuse.
637    ///
638    /// Also verifies that [`RenderChangedShaderBuffers`] only records
639    /// [`ShaderBuffer`]s whose GPU buffer identity changed: creating or
640    /// reallocating the buffer records the asset id, while reusing the existing
641    /// buffer does not.
642    #[test]
643    fn reuses_gpu_buffer_when_unchanged() {
644        let (device, queue) = create_dummy_device();
645
646        let mut assets = Assets::<ShaderBuffer>::default();
647        let handle = assets.add(ShaderBuffer::new(
648            vec![1u32, 2, 3],
649            RenderAssetUsages::default(),
650        ));
651        let asset_id = handle.id();
652
653        // Creating the buffer records it as changed.
654        let extracted = extract(&mut assets.get_mut(asset_id).unwrap(), None);
655        let (first, changed) = prepare(asset_id, extracted, None, &device, &queue);
656        assert!(changed.contains(&asset_id));
657
658        // Same size/usage/label: the existing buffer is reused and its contents
659        // are updated in place, so nothing is recorded as changed.
660        assets.get_mut(asset_id).unwrap().extend([4u32, 5, 6]);
661        let extracted = extract(&mut assets.get_mut(asset_id).unwrap(), Some(&first));
662        let (second, changed) = prepare(asset_id, extracted, Some(&first), &device, &queue);
663        assert_eq!(second.buffer.id(), first.buffer.id());
664        assert!(changed.is_empty());
665
666        // A different buffer size forces a new allocation, which is recorded.
667        assets.get_mut(asset_id).unwrap().extend([7u32]);
668        let extracted = extract(&mut assets.get_mut(asset_id).unwrap(), Some(&first));
669        let (resized, changed) = prepare(asset_id, extracted, Some(&first), &device, &queue);
670        assert_ne!(resized.buffer.id(), first.buffer.id());
671        assert_eq!(resized.buffer.size(), 4);
672        assert!(changed.contains(&asset_id));
673    }
674}
bevy_render/storage.rs

bevy_render/
storage.rs
