bevy_render/
storage.rs

1use crate::{
2    render_asset::{AssetExtractionError, PrepareAssetError, RenderAsset, RenderAssetPlugin},
3    render_resource::{Buffer, BufferUsages},
4    renderer::{RenderDevice, RenderQueue},
5};
6use bevy_app::{App, Plugin};
7use bevy_asset::{Asset, AssetApp, AssetId, RenderAssetUsages};
8use bevy_ecs::system::{lifetimeless::SRes, SystemParamItem};
9use bevy_reflect::{prelude::ReflectDefault, Reflect};
10use bevy_utils::default;
11use encase::{internal::WriteInto, ShaderType};
12use wgpu::util::BufferInitDescriptor;
13
14/// Adds [`ShaderBuffer`] as an asset that is extracted and uploaded to the GPU.
15#[derive(Default)]
16pub struct StoragePlugin;
17
18impl Plugin for StoragePlugin {
19    fn build(&self, app: &mut App) {
20        app.add_plugins(RenderAssetPlugin::<GpuShaderBuffer>::default())
21            .init_asset::<ShaderBuffer>()
22            .register_asset_reflect::<ShaderBuffer>();
23    }
24}
25
26/// A storage buffer that is prepared as a [`RenderAsset`] and uploaded to the GPU.
27#[derive(Asset, Reflect, Debug, Clone)]
28#[reflect(opaque)]
29#[reflect(Default, Debug, Clone)]
30pub struct ShaderBuffer {
31    /// Optional data used to initialize the buffer.
32    pub data: Option<Vec<u8>>,
33    /// The buffer description used to create the buffer.
34    pub buffer_description: wgpu::BufferDescriptor<'static>,
35    /// The asset usage of the storage buffer.
36    pub asset_usage: RenderAssetUsages,
37    /// Whether this buffer should be copied on the GPU when resized.
38    pub copy_on_resize: bool,
39}
40
41impl Default for ShaderBuffer {
42    fn default() -> Self {
43        Self {
44            data: None,
45            buffer_description: wgpu::BufferDescriptor {
46                label: None,
47                size: 0,
48                usage: BufferUsages::STORAGE | BufferUsages::COPY_SRC | BufferUsages::COPY_DST,
49                mapped_at_creation: false,
50            },
51            asset_usage: RenderAssetUsages::default(),
52            copy_on_resize: false,
53        }
54    }
55}
56
57impl ShaderBuffer {
58    /// Creates a new storage buffer with the given data and asset usage.
59    pub fn new(data: &[u8], asset_usage: RenderAssetUsages) -> Self {
60        let mut storage = ShaderBuffer {
61            data: Some(data.to_vec()),
62            ..default()
63        };
64        storage.asset_usage = asset_usage;
65        storage
66    }
67
68    /// Creates a new storage buffer with the given size and asset usage.
69    pub fn with_size(size: usize, asset_usage: RenderAssetUsages) -> Self {
70        let mut storage = ShaderBuffer {
71            data: None,
72            ..default()
73        };
74        storage.buffer_description.size = size as u64;
75        storage.buffer_description.mapped_at_creation = false;
76        storage.asset_usage = asset_usage;
77        storage
78    }
79
80    /// Sets the data of the storage buffer to the given [`ShaderType`].
81    pub fn set_data<T>(&mut self, value: T)
82    where
83        T: ShaderType + WriteInto,
84    {
85        let size = value.size().get() as usize;
86        let mut wrapper = encase::StorageBuffer::<Vec<u8>>::new(Vec::with_capacity(size));
87        wrapper.write(&value).unwrap();
88        self.data = Some(wrapper.into_inner());
89    }
90
91    /// Resizes the buffer to the new size.
92    ///
93    /// If CPU data is present, it will be truncated or zero-extended.
94    /// Does not preserve GPU data when the descriptor changes.
95    pub fn resize(&mut self, size: u64) {
96        self.buffer_description.size = size;
97        if let Some(ref mut data) = self.data {
98            data.resize(size as usize, 0);
99        }
100    }
101
102    /// Resizes the buffer to the new size, preserving existing data.
103    ///
104    /// If CPU data is present, it will be truncated or zero-extended.
105    /// If no CPU data is present, sets `copy_on_resize` to preserve GPU data.
106    pub fn resize_in_place(&mut self, size: u64) {
107        self.buffer_description.size = size;
108        if let Some(ref mut data) = self.data {
109            data.resize(size as usize, 0);
110        } else {
111            self.copy_on_resize = true;
112        }
113    }
114}
115
116impl<T> From<T> for ShaderBuffer
117where
118    T: ShaderType + WriteInto,
119{
120    fn from(value: T) -> Self {
121        let size = value.size().get() as usize;
122        let mut wrapper = encase::StorageBuffer::<Vec<u8>>::new(Vec::with_capacity(size));
123        wrapper.write(&value).unwrap();
124        Self::new(wrapper.as_ref(), RenderAssetUsages::default())
125    }
126}
127
128/// A storage buffer that is prepared as a [`RenderAsset`] and uploaded to the GPU.
129pub struct GpuShaderBuffer {
130    pub buffer: Buffer,
131    pub buffer_descriptor: wgpu::BufferDescriptor<'static>,
132    pub had_data: bool,
133}
134
135impl RenderAsset for GpuShaderBuffer {
136    type SourceAsset = ShaderBuffer;
137    type Param = (SRes<RenderDevice>, SRes<RenderQueue>);
138
139    fn asset_usage(source_asset: &Self::SourceAsset) -> RenderAssetUsages {
140        source_asset.asset_usage
141    }
142
143    fn take_gpu_data(
144        source: &mut Self::SourceAsset,
145        previous_gpu_asset: Option<&Self>,
146    ) -> Result<Self::SourceAsset, AssetExtractionError> {
147        let data = source.data.take();
148
149        let valid_upload = data.is_some() || previous_gpu_asset.is_none_or(|prev| !prev.had_data);
150
151        valid_upload
152            .then(|| Self::SourceAsset {
153                data,
154                ..source.clone()
155            })
156            .ok_or(AssetExtractionError::AlreadyExtracted)
157    }
158
159    fn prepare_asset(
160        source_asset: Self::SourceAsset,
161        _: AssetId<Self::SourceAsset>,
162        (render_device, render_queue): &mut SystemParamItem<Self::Param>,
163        previous_asset: Option<&Self>,
164    ) -> Result<Self, PrepareAssetError<Self::SourceAsset>> {
165        let had_data = source_asset.data.is_some();
166
167        // when cpu data is provided, the actual buffer size is determined by the vec length,
168        // not the descriptor size
169        let actual_size = source_asset
170            .data
171            .as_ref()
172            .map(|d| d.len() as u64)
173            .unwrap_or(source_asset.buffer_description.size);
174
175        let buffer = if let Some(prev) = previous_asset
176            && prev.buffer_descriptor.size == actual_size
177            && prev.buffer_descriptor.usage == source_asset.buffer_description.usage
178            && prev.buffer_descriptor.label == source_asset.buffer_description.label
179            && (!had_data
180                || source_asset
181                    .buffer_description
182                    .usage
183                    .contains(BufferUsages::COPY_DST))
184        {
185            if let Some(ref data) = source_asset.data {
186                render_queue.write_buffer(&prev.buffer, 0, data);
187            }
188            prev.buffer.clone()
189        } else if let Some(ref data) = source_asset.data {
190            render_device.create_buffer_with_data(&BufferInitDescriptor {
191                label: source_asset.buffer_description.label,
192                contents: data,
193                usage: source_asset.buffer_description.usage,
194            })
195        } else {
196            let new_buffer = render_device.create_buffer(&source_asset.buffer_description);
197            if source_asset.copy_on_resize
198                && let Some(previous) = previous_asset
199                && previous
200                    .buffer_descriptor
201                    .usage
202                    .contains(BufferUsages::COPY_SRC)
203                && source_asset
204                    .buffer_description
205                    .usage
206                    .contains(BufferUsages::COPY_DST)
207            {
208                let copy_size = source_asset
209                    .buffer_description
210                    .size
211                    .min(previous.buffer_descriptor.size);
212                let mut encoder =
213                    render_device.create_command_encoder(&wgpu::CommandEncoderDescriptor {
214                        label: Some("copy_buffer_on_resize"),
215                    });
216                encoder.copy_buffer_to_buffer(&previous.buffer, 0, &new_buffer, 0, copy_size);
217                render_queue.submit([encoder.finish()]);
218            }
219            new_buffer
220        };
221
222        Ok(GpuShaderBuffer {
223            buffer,
224            buffer_descriptor: wgpu::BufferDescriptor {
225                size: actual_size,
226                ..source_asset.buffer_description
227            },
228            had_data,
229        })
230    }
231}
bevy_render/storage.rs

bevy_render/
storage.rs
