storage_buffer/
storage_buffer.rs

1//! This example demonstrates how to use a storage buffer with `AsBindGroup` in a custom material.
2use bevy::{
3    mesh::MeshTag,
4    prelude::*,
5    reflect::TypePath,
6    render::{render_resource::AsBindGroup, storage::ShaderStorageBuffer},
7    shader::ShaderRef,
8};
9
10const SHADER_ASSET_PATH: &str = "shaders/storage_buffer.wgsl";
11
12fn main() {
13    App::new()
14        .add_plugins((DefaultPlugins, MaterialPlugin::<CustomMaterial>::default()))
15        .add_systems(Startup, setup)
16        .add_systems(Update, update)
17        .run();
18}
19
20/// set up a simple 3D scene
21fn setup(
22    mut commands: Commands,
23    mut meshes: ResMut<Assets<Mesh>>,
24    mut buffers: ResMut<Assets<ShaderStorageBuffer>>,
25    mut materials: ResMut<Assets<CustomMaterial>>,
26) {
27    // Example data for the storage buffer
28    let color_data: Vec<[f32; 4]> = vec![
29        [1.0, 0.0, 0.0, 1.0],
30        [0.0, 1.0, 0.0, 1.0],
31        [0.0, 0.0, 1.0, 1.0],
32        [1.0, 1.0, 0.0, 1.0],
33        [0.0, 1.0, 1.0, 1.0],
34    ];
35
36    let colors = buffers.add(ShaderStorageBuffer::from(color_data));
37
38    let mesh_handle = meshes.add(Cuboid::from_size(Vec3::splat(0.3)));
39    // Create the custom material with the storage buffer
40    let material_handle = materials.add(CustomMaterial {
41        colors: colors.clone(),
42    });
43
44    commands.insert_resource(CustomMaterialHandle(material_handle.clone()));
45
46    // Spawn cubes with the custom material
47    let mut current_color_id: u32 = 0;
48    for i in -6..=6 {
49        for j in -3..=3 {
50            commands.spawn((
51                Mesh3d(mesh_handle.clone()),
52                MeshMaterial3d(material_handle.clone()),
53                MeshTag(current_color_id % 5),
54                Transform::from_xyz(i as f32, j as f32, 0.0),
55            ));
56            current_color_id += 1;
57        }
58    }
59
60    // Camera
61    commands.spawn((
62        Camera3d::default(),
63        Transform::from_xyz(0.0, 0.0, 10.0).looking_at(Vec3::ZERO, Vec3::Y),
64    ));
65}
66
67// Update the material color by time
68fn update(
69    time: Res<Time>,
70    material_handles: Res<CustomMaterialHandle>,
71    mut materials: ResMut<Assets<CustomMaterial>>,
72    mut buffers: ResMut<Assets<ShaderStorageBuffer>>,
73) {
74    let material = materials.get_mut(&material_handles.0).unwrap();
75
76    let buffer = buffers.get_mut(&material.colors).unwrap();
77    buffer.set_data(
78        (0..5)
79            .map(|i| {
80                let t = time.elapsed_secs() * 5.0;
81                [
82                    ops::sin(t + i as f32) / 2.0 + 0.5,
83                    ops::sin(t + i as f32 + 2.0) / 2.0 + 0.5,
84                    ops::sin(t + i as f32 + 4.0) / 2.0 + 0.5,
85                    1.0,
86                ]
87            })
88            .collect::<Vec<[f32; 4]>>(),
89    );
90}
91
92// Holds handles to the custom materials
93#[derive(Resource)]
94struct CustomMaterialHandle(Handle<CustomMaterial>);
95
96// This struct defines the data that will be passed to your shader
97#[derive(Asset, TypePath, AsBindGroup, Debug, Clone)]
98struct CustomMaterial {
99    #[storage(0, read_only)]
100    colors: Handle<ShaderStorageBuffer>,
101}
102
103impl Material for CustomMaterial {
104    fn vertex_shader() -> ShaderRef {
105        SHADER_ASSET_PATH.into()
106    }
107
108    fn fragment_shader() -> ShaderRef {
109        SHADER_ASSET_PATH.into()
110    }
111}