bevy_render/
storage.rs

1use crate::{
2    render_asset::{AssetExtractionError, PrepareAssetError, RenderAsset, RenderAssetPlugin},
3    render_resource::{Buffer, BufferUsages},
4    renderer::RenderDevice,
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 [`ShaderStorageBuffer`] 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::<GpuShaderStorageBuffer>::default())
21            .init_asset::<ShaderStorageBuffer>()
22            .register_asset_reflect::<ShaderStorageBuffer>();
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 ShaderStorageBuffer {
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}
38
39impl Default for ShaderStorageBuffer {
40    fn default() -> Self {
41        Self {
42            data: None,
43            buffer_description: wgpu::BufferDescriptor {
44                label: None,
45                size: 0,
46                usage: BufferUsages::STORAGE,
47                mapped_at_creation: false,
48            },
49            asset_usage: RenderAssetUsages::default(),
50        }
51    }
52}
53
54impl ShaderStorageBuffer {
55    /// Creates a new storage buffer with the given data and asset usage.
56    pub fn new(data: &[u8], asset_usage: RenderAssetUsages) -> Self {
57        let mut storage = ShaderStorageBuffer {
58            data: Some(data.to_vec()),
59            ..default()
60        };
61        storage.asset_usage = asset_usage;
62        storage
63    }
64
65    /// Creates a new storage buffer with the given size and asset usage.
66    pub fn with_size(size: usize, asset_usage: RenderAssetUsages) -> Self {
67        let mut storage = ShaderStorageBuffer {
68            data: None,
69            ..default()
70        };
71        storage.buffer_description.size = size as u64;
72        storage.buffer_description.mapped_at_creation = false;
73        storage.asset_usage = asset_usage;
74        storage
75    }
76
77    /// Sets the data of the storage buffer to the given [`ShaderType`].
78    pub fn set_data<T>(&mut self, value: T)
79    where
80        T: ShaderType + WriteInto,
81    {
82        let size = value.size().get() as usize;
83        let mut wrapper = encase::StorageBuffer::<Vec<u8>>::new(Vec::with_capacity(size));
84        wrapper.write(&value).unwrap();
85        self.data = Some(wrapper.into_inner());
86    }
87}
88
89impl<T> From<T> for ShaderStorageBuffer
90where
91    T: ShaderType + WriteInto,
92{
93    fn from(value: T) -> Self {
94        let size = value.size().get() as usize;
95        let mut wrapper = encase::StorageBuffer::<Vec<u8>>::new(Vec::with_capacity(size));
96        wrapper.write(&value).unwrap();
97        Self::new(wrapper.as_ref(), RenderAssetUsages::default())
98    }
99}
100
101/// A storage buffer that is prepared as a [`RenderAsset`] and uploaded to the GPU.
102pub struct GpuShaderStorageBuffer {
103    pub buffer: Buffer,
104    pub had_data: bool,
105}
106
107impl RenderAsset for GpuShaderStorageBuffer {
108    type SourceAsset = ShaderStorageBuffer;
109    type Param = SRes<RenderDevice>;
110
111    fn asset_usage(source_asset: &Self::SourceAsset) -> RenderAssetUsages {
112        source_asset.asset_usage
113    }
114
115    fn take_gpu_data(
116        source: &mut Self::SourceAsset,
117        previous_gpu_asset: Option<&Self>,
118    ) -> Result<Self::SourceAsset, AssetExtractionError> {
119        let data = source.data.take();
120
121        let valid_upload = data.is_some() || previous_gpu_asset.is_none_or(|prev| !prev.had_data);
122
123        valid_upload
124            .then(|| Self::SourceAsset {
125                data,
126                ..source.clone()
127            })
128            .ok_or(AssetExtractionError::AlreadyExtracted)
129    }
130
131    fn prepare_asset(
132        source_asset: Self::SourceAsset,
133        _: AssetId<Self::SourceAsset>,
134        render_device: &mut SystemParamItem<Self::Param>,
135        _: Option<&Self>,
136    ) -> Result<Self, PrepareAssetError<Self::SourceAsset>> {
137        match source_asset.data {
138            Some(data) => {
139                let buffer = render_device.create_buffer_with_data(&BufferInitDescriptor {
140                    label: source_asset.buffer_description.label,
141                    contents: &data,
142                    usage: source_asset.buffer_description.usage,
143                });
144                Ok(GpuShaderStorageBuffer {
145                    buffer,
146                    had_data: true,
147                })
148            }
149            None => {
150                let buffer = render_device.create_buffer(&source_asset.buffer_description);
151                Ok(GpuShaderStorageBuffer {
152                    buffer,
153                    had_data: false,
154                })
155            }
156        }
157    }
158}
bevy_render/storage.rs

bevy_render/
storage.rs
