scalar or vector: derive enums and improve enum docs

Author: Firestar99

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

@@ -1,15 +1,26 @@
 use proc_macro2::TokenStream;
 use quote::{ToTokens, quote};
 use syn::punctuated::Punctuated;
-use syn::{Fields, FieldsNamed, FieldsUnnamed, GenericParam, Token};
+use syn::{
+    Data, DataStruct, DataUnion, DeriveInput, 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 item = syn::parse2::<DeriveInput>(item)?;
+    let content = match &item.data {
+        Data::Enum(_) => derive_enum(&spirv_std, &item),
+        Data::Struct(data) => derive_struct(&spirv_std, data),
+        Data::Union(DataUnion { union_token, .. }) => {
+            Err(syn::Error::new_spanned(union_token, "Union not supported"))
+        }
+    }?;
+
+    let ident = &item.ident;
     let gens = &item.generics.params;
     let gen_refs = &item
         .generics
@@ -23,33 +34,57 @@ pub fn derive(item: TokenStream) -> syn::Result<TokenStream> {
         .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 {
+        impl<#gens> #spirv_std::ScalarOrVectorComposite for #ident<#gen_refs> #where_clause {
             #[inline]
             fn transform<F: #spirv_std::ScalarOrVectorTransform>(self, f: &mut F) -> Self {
                 #content
             }
         }
     })
 }
+
+pub fn derive_struct(spirv_std: &TokenStream, data: &DataStruct) -> syn::Result<TokenStream> {
+    Ok(match &data.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),
+    })
+}
+
+pub fn derive_enum(spirv_std: &TokenStream, item: &DeriveInput) -> syn::Result<TokenStream> {
+    let mut attributes = item.attrs.iter().filter(|a| a.path().is_ident("repr"));
+    let repr = match (attributes.next(), attributes.next()) {
+        (None, _) => Err(syn::Error::new_spanned(
+            item,
+            "Missing #[repr(...)] attribute",
+        )),
+        (Some(repr), None) => Ok(repr),
+        (Some(_), Some(_)) => Err(syn::Error::new_spanned(
+            item,
+            "Multiple #[repr(...)] attributes found",
+        )),
+    }?;
+    let prim = &repr.meta.require_list()?.tokens;
+    Ok(quote! {
+        #spirv_std::assert_is_integer::<#prim>();
+        <Self as From<#prim>>::from(#spirv_std::ScalarOrVectorComposite::transform(<Self as Into<#prim>>::into(self), f))
+    })
+}

crates/spirv-std/src/scalar.rs

@@ -117,3 +117,7 @@ impl_scalar! {
     impl Float for f64;
     impl Scalar for bool;
 }
+
+/// used by `ScalarOrVector` derive when working with enums
+#[inline]
+pub fn assert_is_integer<T: Integer>() {}

crates/spirv-std/src/scalar_or_vector.rs

@@ -32,11 +32,23 @@ pub unsafe trait ScalarOrVector: ScalarOrVectorComposite + Default {
 /// data to other threads.
 ///
 /// To derive `#[derive(VectorOrScalarComposite)]` on a struct, all members must also implement
-/// `VectorOrScalarComposite`. To derive it on an enum, the enum must have `#[repr(N)]` where `N` is an [`Integer`].
-/// Additionally, you must derive `num_enum::FromPrimitive` and `num_enum::ToPrimitive`, which requires the enum to be
-/// either exhaustive, implement [`Default`] or a variant of the enum to have the `#[num_enum(default)]` attribute.
+/// `VectorOrScalarComposite`.
+///
+/// To derive it on an enum, the enum must implement `From<N>` and `Into<N>` where `N` is defined by the `#[repr(N)]`
+/// attribute on the enum and is an [`Integer`], like `u32`.
+/// Note that some [safe subgroup operations] may return an "undefined result", so your `From<N>` must gracefully handle
+/// arbitrary bit patterns being passed to it. While panicking is legal, it is discouraged as it may result in
+/// unexpected control flow.
+/// To implement these conversion traits, we recommend [`FromPrimitive`] and [`IntoPrimitive`] from the [`num_enum`]
+/// crate. [`FromPrimitive`] requires that either the enum is exhaustive, or you provide it with a variant to default
+/// to, by either implementing [`Default`] or marking a variant with `#[num_enum(default)]`. Note to disable default
+/// features on the [`num_enum`] crate, or it won't compile on SPIR-V.
 ///
 /// [`Integer`]: crate::Integer
+/// [subgroup operations]: crate::arch::subgroup_shuffle
+/// [`FromPrimitive`]: https://docs.rs/num_enum/latest/num_enum/derive.FromPrimitive.html
+/// [`IntoPrimitive`]: https://docs.rs/num_enum/latest/num_enum/derive.IntoPrimitive.html
+/// [`num_enum`]: https://crates.io/crates/num_enum
 pub trait ScalarOrVectorComposite: Copy + Send + Sync + 'static {
     /// Transform the individual [`Scalar`] and [`Vector`] values of this type to a different value.
     ///

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

@@ -0,0 +1,53 @@
+// build-pass
+// compile-flags: -C target-feature=+GroupNonUniform,+GroupNonUniformShuffle,+ext:SPV_KHR_vulkan_memory_model
+// compile-flags: -C llvm-args=--disassemble-fn=subgroup_composite_enum::disassembly
+// normalize-stderr-test "OpLine .*\n" -> ""
+
+use glam::*;
+use spirv_std::ScalarOrVectorComposite;
+use spirv_std::arch::*;
+use spirv_std::spirv;
+
+#[repr(u32)]
+#[derive(Copy, Clone, Default, ScalarOrVectorComposite)]
+pub enum MyEnum {
+    #[default]
+    A,
+    B,
+    C,
+}
+
+impl From<u32> for MyEnum {
+    #[inline]
+    fn from(value: u32) -> Self {
+        match value {
+            0 => Self::A,
+            1 => Self::B,
+            2 => Self::C,
+            _ => Self::default(),
+        }
+    }
+}
+
+impl From<MyEnum> for u32 {
+    #[inline]
+    fn from(value: MyEnum) -> Self {
+        value as u32
+    }
+}
+
+/// this should be 3 `subgroup_shuffle` instructions, with all calls inlined
+fn disassembly(my_struct: MyEnum, id: u32) -> MyEnum {
+    subgroup_shuffle(my_struct, id)
+}
+
+#[spirv(compute(threads(32)))]
+pub fn main(
+    #[spirv(local_invocation_index)] inv_id: UVec3,
+    #[spirv(descriptor_set = 0, binding = 0, storage_buffer)] output: &mut MyEnum,
+) {
+    unsafe {
+        let my_enum = MyEnum::from(inv_id.x % 3);
+        *output = disassembly(my_enum, 5);
+    }
+}

tests/compiletests/ui/arch/subgroup/subgroup_composite_enum.stderr

@@ -0,0 +1,20 @@
+%1 = OpFunction  %2  None %3
+%4 = OpFunctionParameter  %2
+%5 = OpFunctionParameter  %2
+%6 = OpLabel
+%8 = OpGroupNonUniformShuffle  %2  %9 %4 %5
+OpNoLine
+OpSelectionMerge %10 None
+OpSwitch %8 %11 0 %12 1 %13 2 %14
+%11 = OpLabel
+OpBranch %10
+%12 = OpLabel
+OpBranch %10
+%13 = OpLabel
+OpBranch %10
+%14 = OpLabel
+OpBranch %10
+%10 = OpLabel
+%15 = OpPhi  %2  %16 %11 %16 %12 %17 %13 %18 %14
+OpReturnValue %15
+OpFunctionEnd

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

@@ -0,0 +1,89 @@
+// build-fail
+// normalize-stderr-test "\S*/crates/spirv-std/src/" -> "$$SPIRV_STD_SRC/"
+// normalize-stderr-test "\.rs:\d+:\d+" -> ".rs:"
+// normalize-stderr-test "(\n)\d* *([ -])([\|\+\-\=])" -> "$1   $2$3"
+
+use glam::*;
+use spirv_std::ScalarOrVectorComposite;
+use spirv_std::arch::*;
+use spirv_std::spirv;
+
+macro_rules! enum_repr_from {
+    ($ident:ident, $repr:ty) => {
+        impl From<$repr> for $ident {
+            #[inline]
+            fn from(value: $repr) -> Self {
+                match value {
+                    0 => Self::A,
+                    1 => Self::B,
+                    2 => Self::C,
+                    _ => Self::default(),
+                }
+            }
+        }
+
+        impl From<$ident> for $repr {
+            #[inline]
+            fn from(value: $ident) -> Self {
+                value as $repr
+            }
+        }
+    };
+}
+
+#[derive(Copy, Clone, Default, ScalarOrVectorComposite)]
+pub enum NoRepr {
+    #[default]
+    A,
+    B,
+    C,
+}
+
+#[repr(u32)]
+#[repr(u16)]
+#[derive(Copy, Clone, Default, ScalarOrVectorComposite)]
+pub enum TwoRepr {
+    #[default]
+    A,
+    B,
+    C,
+}
+
+#[repr(C)]
+#[derive(Copy, Clone, Default, ScalarOrVectorComposite)]
+pub enum CRepr {
+    #[default]
+    A,
+    B,
+    C,
+}
+
+#[repr(i32)]
+#[derive(Copy, Clone, Default, ScalarOrVectorComposite)]
+pub enum NoFrom {
+    #[default]
+    A,
+    B,
+    C,
+}
+
+#[repr(i32)]
+#[derive(Copy, Clone, Default, ScalarOrVectorComposite)]
+pub enum WrongFrom {
+    #[default]
+    A,
+    B,
+    C,
+}
+
+enum_repr_from!(WrongFrom, u32);
+
+#[repr(i32)]
+#[derive(Copy, Clone, ScalarOrVectorComposite)]
+pub enum NoDefault {
+    A,
+    B,
+    C,
+}
+
+enum_repr_from!(NoDefault, i32);