storage_buffer/
storage_buffer.rs

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