From 12aadfd4a61971d0fa373c728617b949bd87386c Mon Sep 17 00:00:00 2001
From: Andrii Borziak <username@users.noreply.github.com>
Date: Wed, 22 Feb 2023 22:43:29 +0000
Subject: [PATCH] Support raw buffers in AsBindGroup macro (#7701)

# Objective

There was PR that introduced support for storage buffer is `AsBindGroup` macro [#6129](https://github.com/bevyengine/bevy/pull/6129), but it does not give more granular control over storage buffer, it will always copy all the data no matter which part of it was updated. There is also currently another open PR #6669 that tries to achieve exactly that, it is just not up to date and seems abandoned (Sorry if that is not right). In this PR I'm proposing a solution for both of these approaches to co-exist using `#[storage(n, buffer)]` and `#[storage(n)]` to distinguish between the cases.

We could also discuss in this PR if there is a need to extend this support to DynamicBuffers as well.
---
 .../bevy_render/macros/src/as_bind_group.rs   | 48 ++++++++++++++-----
 .../src/render_resource/bind_group.rs         |  2 +
 2 files changed, 37 insertions(+), 13 deletions(-)

diff --git a/crates/bevy_render/macros/src/as_bind_group.rs b/crates/bevy_render/macros/src/as_bind_group.rs
index 8af04edd68..acf2fed037 100644
--- a/crates/bevy_render/macros/src/as_bind_group.rs
+++ b/crates/bevy_render/macros/src/as_bind_group.rs
@@ -200,6 +200,7 @@ pub fn derive_as_bind_group(ast: syn::DeriveInput) -> Result<TokenStream> {
                     let StorageAttrs {
                         visibility,
                         read_only,
+                        buffer,
                     } = get_storage_binding_attr(nested_meta_items)?;
                     let visibility =
                         visibility.hygenic_quote(&quote! { #render_path::render_resource });
@@ -207,18 +208,32 @@ pub fn derive_as_bind_group(ast: syn::DeriveInput) -> Result<TokenStream> {
                     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(),
-                            },
-                        ))
-                    }});
+                    let min_binding_size = if buffer {
+                        quote! {None}
+                    } else {
+                        quote! {Some(<#field_ty as #render_path::render_resource::ShaderType>::min_size())}
+                    };
+
+                    if buffer {
+                        binding_impls.push(quote! {
+                            #render_path::render_resource::OwnedBindingResource::Buffer({
+                                self.#field_name.clone()
+                            })
+                        });
+                    } else {
+                        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 {
@@ -227,7 +242,7 @@ pub fn derive_as_bind_group(ast: syn::DeriveInput) -> Result<TokenStream> {
                             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()),
+                                min_binding_size: #min_binding_size,
                             },
                             count: None,
                         }
@@ -907,13 +922,16 @@ fn get_sampler_binding_type_value(lit_str: &LitStr) -> Result<SamplerBindingType
 struct StorageAttrs {
     visibility: ShaderStageVisibility,
     read_only: bool,
+    buffer: bool,
 }
 
 const READ_ONLY: Symbol = Symbol("read_only");
+const BUFFER: Symbol = Symbol("buffer");
 
 fn get_storage_binding_attr(metas: Vec<NestedMeta>) -> Result<StorageAttrs> {
     let mut visibility = ShaderStageVisibility::vertex_fragment();
     let mut read_only = false;
+    let mut buffer = false;
 
     for meta in metas {
         use syn::{Meta::List, Meta::Path, NestedMeta::Meta};
@@ -925,6 +943,9 @@ fn get_storage_binding_attr(metas: Vec<NestedMeta>) -> Result<StorageAttrs> {
             Meta(Path(path)) if path == READ_ONLY => {
                 read_only = true;
             }
+            Meta(Path(path)) if path == BUFFER => {
+                buffer = true;
+            }
             _ => {
                 return Err(Error::new_spanned(
                     meta,
@@ -937,5 +958,6 @@ fn get_storage_binding_attr(metas: Vec<NestedMeta>) -> Result<StorageAttrs> {
     Ok(StorageAttrs {
         visibility,
         read_only,
+        buffer,
     })
 }
diff --git a/crates/bevy_render/src/render_resource/bind_group.rs b/crates/bevy_render/src/render_resource/bind_group.rs
index 46e30771b0..15d87a969f 100644
--- a/crates/bevy_render/src/render_resource/bind_group.rs
+++ b/crates/bevy_render/src/render_resource/bind_group.rs
@@ -84,6 +84,8 @@ impl Deref for BindGroup {
 ///     color_texture: Handle<Image>,
 ///     #[storage(3, read_only)]
 ///     values: Vec<f32>,
+///     #[storage(4, read_only, buffer)]
+///     buffer: Buffer,
 /// }
 /// ```
 ///