bevy_render/
storage.rs

1use crate::{
2    render_asset::{PrepareAssetError, RenderAsset, RenderAssetPlugin, RenderAssetUsages},
3    render_resource::{Buffer, BufferUsages},
4    renderer::RenderDevice,
5};
6use bevy_app::{App, Plugin};
7use bevy_asset::{Asset, AssetApp};
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            .register_type::<ShaderStorageBuffer>()
22            .init_asset::<ShaderStorageBuffer>()
23            .register_asset_reflect::<ShaderStorageBuffer>();
24    }
25}
26
27/// A storage buffer that is prepared as a [`RenderAsset`] and uploaded to the GPU.
28#[derive(Asset, Reflect, Debug, Clone)]
29#[reflect(opaque)]
30#[reflect(Default, Debug)]
31pub struct ShaderStorageBuffer {
32    /// Optional data used to initialize the buffer.
33    pub data: Option<Vec<u8>>,
34    /// The buffer description used to create the buffer.
35    pub buffer_description: wgpu::BufferDescriptor<'static>,
36    /// The asset usage of the storage buffer.
37    pub asset_usage: RenderAssetUsages,
38}
39
40impl Default for ShaderStorageBuffer {
41    fn default() -> Self {
42        Self {
43            data: None,
44            buffer_description: wgpu::BufferDescriptor {
45                label: None,
46                size: 0,
47                usage: BufferUsages::STORAGE,
48                mapped_at_creation: false,
49            },
50            asset_usage: RenderAssetUsages::default(),
51        }
52    }
53}
54
55impl ShaderStorageBuffer {
56    /// Creates a new storage buffer with the given data and asset usage.
57    pub fn new(data: &[u8], asset_usage: RenderAssetUsages) -> Self {
58        let mut storage = ShaderStorageBuffer {
59            data: Some(data.to_vec()),
60            ..default()
61        };
62        storage.asset_usage = asset_usage;
63        storage
64    }
65
66    /// Creates a new storage buffer with the given size and asset usage.
67    pub fn with_size(size: usize, asset_usage: RenderAssetUsages) -> Self {
68        let mut storage = ShaderStorageBuffer {
69            data: None,
70            ..default()
71        };
72        storage.buffer_description.size = size as u64;
73        storage.buffer_description.mapped_at_creation = false;
74        storage.asset_usage = asset_usage;
75        storage
76    }
77
78    /// Sets the data of the storage buffer to the given [`ShaderType`].
79    pub fn set_data<T>(&mut self, value: T)
80    where
81        T: ShaderType + WriteInto,
82    {
83        let size = value.size().get() as usize;
84        let mut wrapper = encase::StorageBuffer::<Vec<u8>>::new(Vec::with_capacity(size));
85        wrapper.write(&value).unwrap();
86        self.data = Some(wrapper.into_inner());
87    }
88}
89
90impl<T> From<T> for ShaderStorageBuffer
91where
92    T: ShaderType + WriteInto,
93{
94    fn from(value: T) -> Self {
95        let size = value.size().get() as usize;
96        let mut wrapper = encase::StorageBuffer::<Vec<u8>>::new(Vec::with_capacity(size));
97        wrapper.write(&value).unwrap();
98        Self::new(wrapper.as_ref(), RenderAssetUsages::default())
99    }
100}
101
102/// A storage buffer that is prepared as a [`RenderAsset`] and uploaded to the GPU.
103pub struct GpuShaderStorageBuffer {
104    pub buffer: Buffer,
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 prepare_asset(
116        source_asset: Self::SourceAsset,
117        render_device: &mut SystemParamItem<Self::Param>,
118    ) -> Result<Self, PrepareAssetError<Self::SourceAsset>> {
119        match source_asset.data {
120            Some(data) => {
121                let buffer = render_device.create_buffer_with_data(&BufferInitDescriptor {
122                    label: source_asset.buffer_description.label,
123                    contents: &data,
124                    usage: source_asset.buffer_description.usage,
125                });
126                Ok(GpuShaderStorageBuffer { buffer })
127            }
128            None => {
129                let buffer = render_device.create_buffer(&source_asset.buffer_description);
130                Ok(GpuShaderStorageBuffer { buffer })
131            }
132        }
133    }
134}
bevy_render/storage.rs

bevy_render/
storage.rs
