Support storage buffers in derive AsBindGroup (bevyengine#6129)

# Objective

- Storage buffers are useful and not currently supported by the `AsBindGroup` derive which means you need to expand the macro if you need a storage buffer

## Solution

- Add a new `#[storage]` attribute to the derive `AsBindGroup` macro.
	- Support and optional `read_only` parameter that defaults to false when not present.
	- Support visibility parameters like the texture and sampler attributes.

---

## Changelog

- Add a new `#[storage(index)]` attribute to the derive `AsBindGroup` macro.


Co-authored-by: IceSentry <[email protected]>

crates/bevy_render/macros/src/as_bind_group.rs

@@ -11,13 +11,15 @@ use syn::{
 const UNIFORM_ATTRIBUTE_NAME: Symbol = Symbol("uniform");
 const TEXTURE_ATTRIBUTE_NAME: Symbol = Symbol("texture");
 const SAMPLER_ATTRIBUTE_NAME: Symbol = Symbol("sampler");
+const STORAGE_ATTRIBUTE_NAME: Symbol = Symbol("storage");
 const BIND_GROUP_DATA_ATTRIBUTE_NAME: Symbol = Symbol("bind_group_data");
 
 #[derive(Copy, Clone, Debug)]
 enum BindingType {
     Uniform,
     Texture,
     Sampler,
+    Storage,
 }
 
 #[derive(Clone)]
@@ -55,7 +57,6 @@ pub fn derive_as_bind_group(ast: syn::DeriveInput) -> Result<TokenStream> {
                 }
             } else if attr_ident == UNIFORM_ATTRIBUTE_NAME {
                 let (binding_index, converted_shader_type) = get_uniform_binding_attr(attr)?;
-
                 binding_impls.push(quote! {{
                     use #render_path::render_resource::AsBindGroupShaderType;
                     let mut buffer = #render_path::render_resource::encase::UniformBuffer::new(Vec::new());
@@ -126,6 +127,8 @@ pub fn derive_as_bind_group(ast: syn::DeriveInput) -> Result<TokenStream> {
                 BindingType::Texture
             } else if attr_ident == SAMPLER_ATTRIBUTE_NAME {
                 BindingType::Sampler
+            } else if attr_ident == STORAGE_ATTRIBUTE_NAME {
+                BindingType::Storage
             } else {
                 continue;
             };
@@ -190,7 +193,45 @@ pub fn derive_as_bind_group(ast: syn::DeriveInput) -> Result<TokenStream> {
             }
 
             match binding_type {
-                BindingType::Uniform => { /* uniform codegen is deferred to account for combined uniform bindings */
+                BindingType::Uniform => {
+                    // uniform codegen is deferred to account for combined uniform bindings
+                }
+                BindingType::Storage => {
+                    let StorageAttrs {
+                        visibility,
+                        read_only,
+                    } = get_storage_binding_attr(nested_meta_items)?;
+                    let visibility =
+                        visibility.hygenic_quote(&quote! { #render_path::render_resource });
+
+                    let field_name = field.ident.as_ref().unwrap();
+                    let field_ty = &field.ty;
+
+                    binding_impls.push(quote! {{
+                        use #render_path::render_resource::AsBindGroupShaderType;
+                        let mut buffer = #render_path::render_resource::encase::StorageBuffer::new(Vec::new());
+                        buffer.write(&self.#field_name).unwrap();
+                        #render_path::render_resource::OwnedBindingResource::Buffer(render_device.create_buffer_with_data(
+                            &#render_path::render_resource::BufferInitDescriptor {
+                                label: None,
+                                usage: #render_path::render_resource::BufferUsages::COPY_DST | #render_path::render_resource::BufferUsages::STORAGE,
+                                contents: buffer.as_ref(),
+                            },
+                        ))
+                    }});
+
+                    binding_layouts.push(quote! {
+                        #render_path::render_resource::BindGroupLayoutEntry {
+                            binding: #binding_index,
+                            visibility: #visibility,
+                            ty: #render_path::render_resource::BindingType::Buffer {
+                                ty: #render_path::render_resource::BufferBindingType::Storage { read_only: #read_only },
+                                has_dynamic_offset: false,
+                                min_binding_size: Some(<#field_ty as #render_path::render_resource::ShaderType>::min_size()),
+                            },
+                            count: None,
+                        }
+                    });
                 }
                 BindingType::Texture => {
                     let TextureAttrs {
@@ -861,3 +902,40 @@ fn get_sampler_binding_type_value(lit_str: &LitStr) -> Result<SamplerBindingType
         )),
     }
 }
+
+#[derive(Default)]
+struct StorageAttrs {
+    visibility: ShaderStageVisibility,
+    read_only: bool,
+}
+
+const READ_ONLY: Symbol = Symbol("read_only");
+
+fn get_storage_binding_attr(metas: Vec<NestedMeta>) -> Result<StorageAttrs> {
+    let mut visibility = ShaderStageVisibility::vertex_fragment();
+    let mut read_only = false;
+
+    for meta in metas {
+        use syn::{Meta::List, Meta::Path, NestedMeta::Meta};
+        match meta {
+            // Parse #[storage(0, visibility(...))].
+            Meta(List(m)) if m.path == VISIBILITY => {
+                visibility = get_visibility_flag_value(&m.nested)?;
+            }
+            Meta(Path(path)) if path == READ_ONLY => {
+                read_only = true;
+            }
+            _ => {
+                return Err(Error::new_spanned(
+                    meta,
+                    "Not a valid attribute. Available attributes: `read_only`, `visibility`",
+                ));
+            }
+        }
+    }
+
+    Ok(StorageAttrs {
+        visibility,
+        read_only,
+    })
+}

crates/bevy_render/macros/src/lib.rs

@@ -17,7 +17,10 @@ pub fn derive_extract_resource(input: TokenStream) -> TokenStream {
     extract_resource::derive_extract_resource(input)
 }
 
-#[proc_macro_derive(AsBindGroup, attributes(uniform, texture, sampler, bind_group_data))]
+#[proc_macro_derive(
+    AsBindGroup,
+    attributes(uniform, texture, sampler, bind_group_data, storage)
+)]
 pub fn derive_as_bind_group(input: TokenStream) -> TokenStream {
     let input = parse_macro_input!(input as DeriveInput);
 

crates/bevy_render/src/render_resource/bind_group.rs

@@ -82,6 +82,8 @@ impl Deref for BindGroup {
 ///     #[texture(1)]
 ///     #[sampler(2)]
 ///     color_texture: Handle<Image>,
+///     #[storage(3, read_only)]
+///     values: Vec<f32>,
 /// }
 /// ```
 ///
@@ -94,6 +96,8 @@ impl Deref for BindGroup {
 /// var color_texture: texture_2d<f32>;
 /// @group(1) @binding(2)
 /// var color_sampler: sampler;
+/// @group(1) @binding(3)
+/// var<storage> values: array<f32>;
 /// ```
 /// Note that the "group" index is determined by the usage context. It is not defined in [`AsBindGroup`]. For example, in Bevy material bind groups
 /// are generally bound to group 1.
@@ -132,6 +136,15 @@ impl Deref for BindGroup {
 /// | `sampler_type` = "..." | `"filtering"`, `"non_filtering"`, `"comparison"`.                       |  `"filtering"`         |
 /// | `visibility(...)`      | `all`, `none`, or a list-combination of `vertex`, `fragment`, `compute` |   `vertex`, `fragment` |
 ///
+/// * `storage(BINDING_INDEX, arguments)`
+///     * The field will be converted to a shader-compatible type using the [`ShaderType`] trait, written to a [`Buffer`], and bound as a storage buffer.
+///     * It supports and optional `read_only` parameter. Defaults to false if not present.
+///
+/// | Arguments              | Values                                                                  | Default              |
+/// |------------------------|-------------------------------------------------------------------------|----------------------|
+/// | `visibility(...)`      | `all`, `none`, or a list-combination of `vertex`, `fragment`, `compute` | `vertex`, `fragment` |
+/// | `read_only`            | if present then value is true, otherwise false                          | `false`              |
+///
 /// Note that fields without field-level binding attributes will be ignored.
 /// ```
 /// # use bevy_render::{color::Color, render_resource::AsBindGroup};