docs/bevy/examples/shader/storage_buffer.rs 3 K · 108 lines · raw · history

1 //! This example demonstrates how to use a storage buffer with `AsBindGroup` in a custom material.
2 use bevy::{
3 mesh::MeshTag,
4 prelude::*,
5 reflect::TypePath,
6 render::{render_resource::AsBindGroup, storage::ShaderBuffer},
7 shader::ShaderRef,
8 };
9
10 const SHADER_ASSET_PATH: &str = "shaders/storage_buffer.wesl";
11
12 fn 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
21 fn setup(
22 mut commands: Commands,
23 mut meshes: ResMut<Assets<Mesh>>,
24 mut buffers: ResMut<Assets<ShaderBuffer>>,
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(ShaderBuffer::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::new(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
68 fn update(
69 time: Res<Time>,
70 material_handles: Res<CustomMaterialHandle>,
71 mut materials: ResMut<Assets<CustomMaterial>>,
72 mut buffers: ResMut<Assets<ShaderBuffer>>,
73 ) {
74 let material = materials.get_mut(&material_handles.0).unwrap();
75
76 let mut buffer = buffers.get_mut(&material.colors).unwrap();
77 buffer.clear();
78 buffer.extend((0..5).map(|i| {
79 let t = time.elapsed_secs() * 5.0;
80 [
81 ops::sin(t + i as f32) / 2.0 + 0.5,
82 ops::sin(t + i as f32 + 2.0) / 2.0 + 0.5,
83 ops::sin(t + i as f32 + 4.0) / 2.0 + 0.5,
84 1.0,
85 ]
86 }));
87 }
88
89 // Holds handles to the custom materials
90 #[derive(Resource)]
91 struct CustomMaterialHandle(Handle<CustomMaterial>);
92
93 // This struct defines the data that will be passed to your shader
94 #[derive(Asset, TypePath, AsBindGroup, Debug, Clone)]
95 struct CustomMaterial {
96 #[storage(0, read_only)]
97 colors: Handle<ShaderBuffer>,
98 }
99
100 impl Material for CustomMaterial {
101 fn vertex_shader() -> ShaderRef {
102 SHADER_ASSET_PATH.into()
103 }
104
105 fn fragment_shader() -> ShaderRef {
106 SHADER_ASSET_PATH.into()
107 }
108 }