scalar or vector: add #[derive(ScalarOrVectorComposite)]

Author: Firestar99

crates/spirv-std/macros/src/lib.rs

@@ -74,6 +74,7 @@
 mod debug_printf;
 mod image;
 mod sample_param_permutations;
+mod scalar_or_vector_composite;
 
 use crate::debug_printf::{DebugPrintfInput, debug_printf_inner};
 use proc_macro::TokenStream;
@@ -309,3 +310,10 @@ pub fn debug_printfln(input: TokenStream) -> TokenStream {
 pub fn gen_sample_param_permutations(_attr: TokenStream, item: TokenStream) -> TokenStream {
     sample_param_permutations::gen_sample_param_permutations(item)
 }
+
+#[proc_macro_derive(ScalarOrVectorComposite)]
+pub fn derive_scalar_or_vector_composite(item: TokenStream) -> TokenStream {
+    scalar_or_vector_composite::derive(item.into())
+        .unwrap_or_else(syn::Error::into_compile_error)
+        .into()
+}

crates/spirv-std/macros/src/scalar_or_vector_composite.rs

@@ -0,0 +1,54 @@
+use proc_macro2::TokenStream;
+use quote::{ToTokens, quote};
+use syn::punctuated::Punctuated;
+use syn::{Fields, FieldsNamed, FieldsUnnamed, GenericParam, Token};
+
+pub fn derive(item: TokenStream) -> syn::Result<TokenStream> {
+    // Whenever we'll properly resolve the crate symbol, replace this.
+    let spirv_std = quote!(spirv_std);
+
+    // Defer all validation to our codegen backend. Rather than erroring here, emit garbage.
+    let item = syn::parse2::<syn::ItemStruct>(item)?;
+    let struct_ident = &item.ident;
+    let gens = &item.generics.params;
+    let gen_refs = &item
+        .generics
+        .params
+        .iter()
+        .map(|p| match p {
+            GenericParam::Lifetime(p) => p.lifetime.to_token_stream(),
+            GenericParam::Type(p) => p.ident.to_token_stream(),
+            GenericParam::Const(p) => p.ident.to_token_stream(),
+        })
+        .collect::<Punctuated<_, Token![,]>>();
+    let where_clause = &item.generics.where_clause;
+
+    let content =
+        match item.fields {
+            Fields::Named(FieldsNamed { named, .. }) => {
+                let content = named.iter().map(|f| {
+                let ident = &f.ident;
+                quote!(#ident: #spirv_std::ScalarOrVectorComposite::transform(self.#ident, f))
+            }).collect::<Punctuated<_, Token![,]>>();
+                quote!(Self { #content })
+            }
+            Fields::Unnamed(FieldsUnnamed { unnamed, .. }) => {
+                let content = (0..unnamed.len())
+                    .map(|i| {
+                        let i = syn::Index::from(i);
+                        quote!(#spirv_std::ScalarOrVectorComposite::transform(self.#i, f))
+                    })
+                    .collect::<Punctuated<_, Token![,]>>();
+                quote!(Self(#content))
+            }
+            Fields::Unit => quote!(Self),
+        };
+
+    Ok(quote! {
+        impl<#gens> #spirv_std::ScalarOrVectorComposite for #struct_ident<#gen_refs> #where_clause {
+            fn transform<F: #spirv_std::ScalarOrVectorTransform>(self, f: &mut F) -> Self {
+                #content
+            }
+        }
+    })
+}

crates/spirv-std/src/lib.rs

@@ -87,6 +87,7 @@
 /// Public re-export of the `spirv-std-macros` crate.
 #[macro_use]
 pub extern crate spirv_std_macros as macros;
+pub use macros::ScalarOrVectorComposite;
 pub use macros::spirv;
 
 pub mod arch;

tests/compiletests/ui/arch/subgroup/subgroup_composite.rs

@@ -0,0 +1,46 @@
+// build-pass
+// compile-flags: -C target-feature=+GroupNonUniform,+GroupNonUniformBallot,+GroupNonUniformShuffle,+GroupNonUniformShuffleRelative,+ext:SPV_KHR_vulkan_memory_model
+// normalize-stderr-test "OpLine .*\n" -> ""
+
+use glam::*;
+use spirv_std::ScalarOrVectorComposite;
+use spirv_std::arch::*;
+use spirv_std::spirv;
+
+#[derive(Copy, Clone, ScalarOrVectorComposite)]
+pub struct MyStruct {
+    a: f32,
+    b: UVec3,
+    c: Nested,
+    d: Zst,
+}
+
+#[derive(Copy, Clone, ScalarOrVectorComposite)]
+pub struct Nested(i32);
+
+#[derive(Copy, Clone, ScalarOrVectorComposite)]
+pub struct Zst;
+
+#[spirv(compute(threads(32)))]
+pub fn main(
+    #[spirv(local_invocation_index)] inv_id: UVec3,
+    #[spirv(descriptor_set = 0, binding = 0, storage_buffer)] output: &mut UVec3,
+) {
+    unsafe {
+        let my_struct = MyStruct {
+            a: 1.,
+            b: inv_id,
+            c: Nested(-42),
+            d: Zst,
+        };
+
+        let mut out = UVec3::ZERO;
+        out += subgroup_broadcast(my_struct, 19).b;
+        out += subgroup_broadcast_first(my_struct).b;
+        out += subgroup_shuffle(my_struct, 2).b;
+        out += subgroup_shuffle_xor(my_struct, 4).b;
+        out += subgroup_shuffle_up(my_struct, 5).b;
+        out += subgroup_shuffle_down(my_struct, 7).b;
+        *output = out;
+    }
+}