bevy_render/
storage.rs

1use crate::{
2    render_asset::{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}
105
106impl RenderAsset for GpuShaderStorageBuffer {
107    type SourceAsset = ShaderStorageBuffer;
108    type Param = SRes<RenderDevice>;
109
110    fn asset_usage(source_asset: &Self::SourceAsset) -> RenderAssetUsages {
111        source_asset.asset_usage
112    }
113
114    fn prepare_asset(
115        source_asset: Self::SourceAsset,
116        _: AssetId<Self::SourceAsset>,
117        render_device: &mut SystemParamItem<Self::Param>,
118        _: Option<&Self>,
119    ) -> Result<Self, PrepareAssetError<Self::SourceAsset>> {
120        match source_asset.data {
121            Some(data) => {
122                let buffer = render_device.create_buffer_with_data(&BufferInitDescriptor {
123                    label: source_asset.buffer_description.label,
124                    contents: &data,
125                    usage: source_asset.buffer_description.usage,
126                });
127                Ok(GpuShaderStorageBuffer { buffer })
128            }
129            None => {
130                let buffer = render_device.create_buffer(&source_asset.buffer_description);
131                Ok(GpuShaderStorageBuffer { buffer })
132            }
133        }
134    }
135}
bevy_render/storage.rs

bevy_render/
storage.rs
