Add convenience methods for constructing and setting storage buffer data

crates/bevy_render/src/storage.rs

@@ -8,6 +8,8 @@ use bevy_ecs::system::SystemParamItem;
 use bevy_reflect::prelude::ReflectDefault;
 use bevy_reflect::Reflect;
 use bevy_utils::default;
+use encase::internal::WriteInto;
+use encase::ShaderType;
 use wgpu::util::BufferInitDescriptor;
 
 /// Adds [`ShaderStorageBuffer`] as an asset that is extracted and uploaded to the GPU.
@@ -72,6 +74,29 @@ impl ShaderStorageBuffer {
         storage.asset_usage = asset_usage;
         storage
     }
+
+    /// Sets the data of the storage buffer to the given [`ShaderType`].
+    pub fn set_data<T>(&mut self, value: T)
+    where
+        T: ShaderType + WriteInto,
+    {
+        let size = value.size().get() as usize;
+        let mut wrapper = encase::StorageBuffer::<Vec<u8>>::new(Vec::with_capacity(size));
+        wrapper.write(&value).unwrap();
+        self.data = Some(wrapper.into_inner());
+    }
+}
+
+impl<T> From<T> for ShaderStorageBuffer
+where
+    T: ShaderType + WriteInto,
+{
+    fn from(value: T) -> Self {
+        let size = value.size().get() as usize;
+        let mut wrapper = encase::StorageBuffer::<Vec<u8>>::new(Vec::with_capacity(size));
+        wrapper.write(&value).unwrap();
+        Self::new(wrapper.as_ref(), RenderAssetUsages::default())
+    }
 }
 
 /// A storage buffer that is prepared as a [`RenderAsset`] and uploaded to the GPU.

examples/shader/storage_buffer.rs

@@ -4,7 +4,6 @@ use bevy::{
     reflect::TypePath,
     render::render_resource::{AsBindGroup, ShaderRef},
 };
-use bevy_render::render_asset::RenderAssetUsages;
 use bevy_render::storage::ShaderStorageBuffer;
 
 const SHADER_ASSET_PATH: &str = "shaders/storage_buffer.wgsl";
@@ -33,10 +32,7 @@ fn setup(
         [0.0, 1.0, 1.0, 1.0],
     ];
 
-    let colors = buffers.add(ShaderStorageBuffer::new(
-        bytemuck::cast_slice(color_data.as_slice()),
-        RenderAssetUsages::default(),
-    ));
+    let colors = buffers.add(ShaderStorageBuffer::from(color_data));
 
     // Create the custom material with the storage buffer
     let custom_material = CustomMaterial { colors };
@@ -72,22 +68,19 @@ fn update(
 ) {
     let material = materials.get_mut(&material_handle.0).unwrap();
     let buffer = buffers.get_mut(&material.colors).unwrap();
-    buffer.data = Some(
-        bytemuck::cast_slice(
-            (0..5)
-                .map(|i| {
-                    let t = time.elapsed_seconds() * 5.0;
-                    [
-                        (t + i as f32).sin() / 2.0 + 0.5,
-                        (t + i as f32 + 2.0).sin() / 2.0 + 0.5,
-                        (t + i as f32 + 4.0).sin() / 2.0 + 0.5,
-                        1.0,
-                    ]
-                })
-                .collect::<Vec<[f32; 4]>>()
-                .as_slice(),
-        )
-        .to_vec(),
+    buffer.set_data(
+        (0..5)
+            .map(|i| {
+                let t = time.elapsed_seconds() * 5.0;
+                [
+                    (t + i as f32).sin() / 2.0 + 0.5,
+                    (t + i as f32 + 2.0).sin() / 2.0 + 0.5,
+                    (t + i as f32 + 4.0).sin() / 2.0 + 0.5,
+                    1.0,
+                ]
+            })
+            .collect::<Vec<[f32; 4]>>()
+            .as_slice(),
     );
 }