moved permutations to complete ImageWithMethods trait impls

Author: oisyn

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

@@ -631,30 +631,42 @@ const SAMPLE_PARAM_TYPES: [&str; SAMPLE_PARAM_COUNT] = ["B", "L", "S"];
 const SAMPLE_PARAM_OPERANDS: [&str; SAMPLE_PARAM_COUNT] = ["Bias", "Lod", "Sample"];
 const SAMPLE_PARAM_NAMES: [&str; SAMPLE_PARAM_COUNT] = ["bias", "lod", "sample_index"];
 
-struct SampleFnRewriter(usize);
+struct SampleImplRewriter(usize, syn::Type);
 
-impl SampleFnRewriter {
-    pub fn rewrite(mask: usize, f: &syn::ItemFn) -> syn::ItemFn {
-        let mut new_f = f.clone();
-        let mut ty = String::from("SampleParams<");
+impl SampleImplRewriter {
+    pub fn rewrite(mask: usize, f: &syn::ItemImpl) -> syn::ItemImpl {
+        let mut new_impl = f.clone();
+        let mut ty_str = String::from("SampleParams<");
 
+        // based on the mask, form a `SampleParams` type string and add the generic parameters to the `impl<>` generics
+        // example type string: `"SampleParams<SomeTy<B>, NoneTy, NoneTy>"`
         for i in 0..SAMPLE_PARAM_COUNT {
             if mask & (1 << i) != 0 {
-                new_f.sig.generics.params.push(syn::GenericParam::Type(
+                new_impl.generics.params.push(syn::GenericParam::Type(
                     syn::Ident::new(SAMPLE_PARAM_TYPES[i], Span::call_site()).into(),
                 ));
-                ty.push_str(SAMPLE_PARAM_TYPES[i]);
+                ty_str.push_str("SomeTy<");
+                ty_str.push_str(SAMPLE_PARAM_TYPES[i]);
+                ty_str.push('>');
             } else {
-                ty.push_str("()");
+                ty_str.push_str("NoneTy");
             }
-            ty.push(',');
+            ty_str.push(',');
         }
-        ty.push('>');
-        if let Some(syn::FnArg::Typed(p)) = new_f.sig.inputs.last_mut() {
-            *p.ty.as_mut() = syn::parse(ty.parse().unwrap()).unwrap();
+        ty_str.push_str(">");
+        let ty: syn::Type = syn::parse(ty_str.parse().unwrap()).unwrap();
+
+        if let Some(t) = &mut new_impl.trait_ {
+            if let syn::PathArguments::AngleBracketed(a) =
+                &mut t.1.segments.last_mut().unwrap().arguments
+            {
+                if let Some(syn::GenericArgument::Type(t)) = a.args.last_mut() {
+                    *t = ty.clone();
+                }
+            }
         }
-        SampleFnRewriter(mask).visit_item_fn_mut(&mut new_f);
-        new_f
+        SampleImplRewriter(mask, ty).visit_item_impl_mut(&mut new_impl);
+        new_impl
     }
 
     fn get_operands(&self) -> String {
@@ -686,18 +698,22 @@ impl SampleFnRewriter {
     fn add_regs(&self, t: &mut Vec<TokenTree>) {
         for i in 0..SAMPLE_PARAM_COUNT {
             if self.0 & (1 << i) != 0 {
-                let s = format!("{0} = in(reg) &param.{0}", SAMPLE_PARAM_NAMES[i]);
-                t.push(TokenTree::Literal(proc_macro2::Literal::string(s.as_str())));
-                t.push(TokenTree::Punct(proc_macro2::Punct::new(
-                    ',',
-                    proc_macro2::Spacing::Alone,
-                )))
+                let s = format!("{0} = in(reg) &param.{0},", SAMPLE_PARAM_NAMES[i]);
+                let ts: proc_macro2::TokenStream = s.parse().unwrap();
+                t.extend(ts);
             }
         }
     }
 }
 
-impl syn::visit_mut::VisitMut for SampleFnRewriter {
+impl VisitMut for SampleImplRewriter {
+    fn visit_impl_item_method_mut(&mut self, item: &mut syn::ImplItemMethod) {
+        if let Some(syn::FnArg::Typed(p)) = item.sig.inputs.last_mut() {
+            *p.ty.as_mut() = self.1.clone();
+        }
+        syn::visit_mut::visit_impl_item_method_mut(self, item);
+    }
+
     fn visit_macro_mut(&mut self, m: &mut syn::Macro) {
         if m.path.is_ident("asm") {
             let t = m.tokens.clone();
@@ -745,12 +761,13 @@ impl syn::visit_mut::VisitMut for SampleFnRewriter {
 #[proc_macro_attribute]
 #[doc(hidden)]
 pub fn gen_sample_param_permutations(_attr: TokenStream, item: TokenStream) -> TokenStream {
-    let item_fn = syn::parse_macro_input!(item as syn::ItemFn);
+    let item_impl = syn::parse_macro_input!(item as syn::ItemImpl);
     let mut fns = Vec::new();
 
-    for m in 1..((1 << SAMPLE_PARAM_COUNT) - 1) {
-        fns.push(SampleFnRewriter::rewrite(m, &item_fn));
+    for m in 1..(1 << SAMPLE_PARAM_COUNT) {
+        fns.push(SampleImplRewriter::rewrite(m, &item_impl));
     }
 
+    println!("{}", quote! { #(#fns)* }.to_string());
     quote! { #(#fns)* }.into()
 }

crates/spirv-std/src/image.rs

@@ -10,12 +10,17 @@ use core::arch::asm;
 #[rustfmt::skip]
 mod params;
 
-pub use self::params::{ImageCoordinate, ImageCoordinateSubpassData, SampleParams, SampleType};
+/// Contains extra image operands
+pub mod sample;
+
+pub use self::params::{ImageCoordinate, ImageCoordinateSubpassData, SampleType};
 pub use crate::macros::Image;
 pub use spirv_std_types::image_params::{
     AccessQualifier, Arrayed, Dimensionality, ImageDepth, ImageFormat, Multisampled, Sampled,
 };
 
+use sample::{NoneTy, SampleParams, SomeTy};
+
 use crate::{float::Float, integer::Integer, vector::Vector, Sampler};
 
 /// Re-export of primitive types to ensure the `Image` proc macro always points
@@ -158,33 +163,6 @@ impl<
         }
         result.truncate_into()
     }
-
-    #[crate::macros::gen_sample_param_permutations]
-    #[crate::macros::gpu_only]
-    #[doc(alias = "OpImageFetch")]
-    pub fn fetch_with<I>(
-        &self,
-        coordinate: impl ImageCoordinate<I, DIM, ARRAYED>,
-        params: SampleParams,
-    ) -> SampledType::SampleResult
-    where
-        I: Integer,
-        B: Integer,
-    {
-        let mut result = SampledType::Vec4::default();
-        unsafe {
-            asm! {
-                "%image = OpLoad _ {this}",
-                "%coordinate = OpLoad _ {coordinate}",
-                "%result = OpImageFetch typeof*{result} %image %coordinate $PARAMS",
-                "OpStore {result} %result",
-                result = in(reg) &mut result,
-                this = in(reg) self,
-                coordinate = in(reg) &coordinate,
-            }
-        }
-        result.truncate_into()
-    }
 }
 
 impl<
@@ -1142,6 +1120,79 @@ impl<
     }
 }
 
+/// Helper trait that defines all `*_with` methods on an `Image` that use the extra image operands,
+/// such as bias or lod, defined by the `SampleParams` struct.
+pub trait ImageWithMethods<
+    SampledType: SampleType<FORMAT, COMPONENTS>,
+    const DIM: u32,
+    const DEPTH: u32,
+    const ARRAYED: u32,
+    const MULTISAMPLED: u32,
+    const SAMPLED: u32,
+    const FORMAT: u32,
+    const COMPONENTS: u32,
+    Params,
+>
+{
+    /// Fetch a single texel with a sampler set at compile time
+    fn fetch_with<I>(
+        &self,
+        coordinate: impl ImageCoordinate<I, DIM, ARRAYED>,
+        params: Params,
+    ) -> SampledType::SampleResult
+    where
+        I: Integer;
+}
+
+#[crate::macros::gen_sample_param_permutations]
+impl<
+        SampledType: SampleType<FORMAT, COMPONENTS>,
+        const DIM: u32,
+        const DEPTH: u32,
+        const ARRAYED: u32,
+        const MULTISAMPLED: u32,
+        const SAMPLED: u32,
+        const FORMAT: u32,
+        const COMPONENTS: u32,
+    >
+    ImageWithMethods<
+        SampledType,
+        DIM,
+        DEPTH,
+        ARRAYED,
+        MULTISAMPLED,
+        SAMPLED,
+        FORMAT,
+        COMPONENTS,
+        SampleParams,
+    > for Image<SampledType, DIM, DEPTH, ARRAYED, MULTISAMPLED, SAMPLED, FORMAT, COMPONENTS>
+{
+    #[crate::macros::gpu_only]
+    #[doc(alias = "OpImageFetch")]
+    fn fetch_with<I>(
+        &self,
+        coordinate: impl ImageCoordinate<I, DIM, ARRAYED>,
+        params: SampleParams,
+    ) -> SampledType::SampleResult
+    where
+        I: Integer,
+    {
+        let mut result = SampledType::Vec4::default();
+        unsafe {
+            asm! {
+                "%image = OpLoad _ {this}",
+                "%coordinate = OpLoad _ {coordinate}",
+                "%result = OpImageFetch typeof*{result} %image %coordinate $PARAMS",
+                "OpStore {result} %result",
+                result = in(reg) &mut result,
+                this = in(reg) self,
+                coordinate = in(reg) &coordinate,
+            }
+        }
+        result.truncate_into()
+    }
+}
+
 /// This is a marker trait to represent the constraints on `OpImageGather` too complex to be
 /// represented by const generics. Specifically:
 ///

crates/spirv-std/src/image/params.rs

@@ -194,66 +194,3 @@ impl<V: Vector<S, 4>, S: Scalar>
 pub trait ImageCoordinateSubpassData<T, const ARRAYED: u32> {}
 impl<V: Vector<I, 2>, I: Integer> ImageCoordinateSubpassData<I, { Arrayed::False as u32 }> for V {}
 impl<V: Vector<I, 3>, I: Integer> ImageCoordinateSubpassData<I, { Arrayed::True as u32 }> for V {}
-
-/// Helper struct that allows building image operands. Start with a global function that returns this
-/// struct, and then chain additional calls.
-/// Example: `image.sample_with(coords, params::bias(3.0).sample_index(1))`
-pub struct SampleParams<B, L, S> {
-    bias: B,
-    lod: L,
-    sample_index: S,
-}
-
-pub fn bias<B>(bias: B) -> SampleParams<B, (), ()> {
-    SampleParams {
-        bias,
-        lod: (),
-        sample_index: (),
-    }
-}
-
-pub fn lod<L>(lod: L) -> SampleParams<(), L, ()> {
-    SampleParams {
-        bias: (),
-        lod,
-        sample_index: (),
-    }
-}
-
-pub fn sample_index<S>(sample_index: S) -> SampleParams<(), (), S> {
-    SampleParams {
-        bias: (),
-        lod: (),
-        sample_index,
-    }
-}
-
-impl<L, S> SampleParams<(), L, S> {
-    pub fn bias<B>(self, bias: B) -> SampleParams<B, L, S> {
-        SampleParams {
-            bias,
-            lod: self.lod,
-            sample_index: self.sample_index,
-        }
-    }
-}
-
-impl<B, S> SampleParams<B, (), S> {
-    pub fn lod<L>(self, lod: L) -> SampleParams<B, L, S> {
-        SampleParams {
-            bias: self.bias,
-            lod: lod,
-            sample_index: self.sample_index,
-        }
-    }
-}
-
-impl<B, L> SampleParams<B, L, ()> {
-    pub fn sample_index<S>(self, sample_index: S) -> SampleParams<B, L, S> {
-        SampleParams {
-            bias: self.bias,
-            lod: self.lod,
-            sample_index: sample_index,
-        }
-    }
-}

crates/spirv-std/src/image/sample.rs

@@ -0,0 +1,87 @@
+/// Helper trait to mimic `Option<T>`, but where the variant are types
+pub trait OptionTy {
+    /// Whether this is a `NoneTy` (when false) or a `SomeTy<T>` (when true)
+    const EXISTS: bool;
+}
+
+impl OptionTy for NoneTy {
+    const EXISTS: bool = false;
+}
+
+impl<T> OptionTy for SomeTy<T> {
+    const EXISTS: bool = true;
+}
+/// Helper struct that denotes that the type doesn't exist, analog to `Option::None`
+pub struct NoneTy;
+
+/// Helper struct that denotes that the type does exist and is of type T, analog to `Option::Some(T)`
+pub struct SomeTy<T>(T);
+
+/// Helper struct that allows building image operands. Start with a global function that returns this
+/// struct, and then chain additional calls.
+/// Example: `image.sample_with(coords, params::bias(3.0).sample_index(1))`
+pub struct SampleParams<B: OptionTy, L: OptionTy, S: OptionTy> {
+    bias: B,
+    lod: L,
+    sample_index: S,
+}
+
+/// Sets the 'Bias' image operand
+pub fn bias<B>(bias: B) -> SampleParams<SomeTy<B>, NoneTy, NoneTy> {
+    SampleParams {
+        bias: SomeTy(bias),
+        lod: NoneTy,
+        sample_index: NoneTy,
+    }
+}
+
+/// Sets the 'Lod' image operand
+pub fn lod<L>(lod: L) -> SampleParams<NoneTy, SomeTy<L>, NoneTy> {
+    SampleParams {
+        bias: NoneTy,
+        lod: SomeTy(lod),
+        sample_index: NoneTy,
+    }
+}
+
+/// Sets the 'Sample' image operand
+pub fn sample_index<S>(sample_index: S) -> SampleParams<NoneTy, NoneTy, SomeTy<S>> {
+    SampleParams {
+        bias: NoneTy,
+        lod: NoneTy,
+        sample_index: SomeTy(sample_index),
+    }
+}
+
+impl<L: OptionTy, S: OptionTy> SampleParams<NoneTy, L, S> {
+    /// Sets the 'Bias' image operand
+    pub fn bias<B>(self, bias: B) -> SampleParams<SomeTy<B>, L, S> {
+        SampleParams {
+            bias: SomeTy(bias),
+            lod: self.lod,
+            sample_index: self.sample_index,
+        }
+    }
+}
+
+impl<B: OptionTy, S: OptionTy> SampleParams<B, NoneTy, S> {
+    /// Sets the 'Lod' image operand
+    pub fn lod<L>(self, lod: L) -> SampleParams<B, SomeTy<L>, S> {
+        SampleParams {
+            bias: self.bias,
+            lod: SomeTy(lod),
+            sample_index: self.sample_index,
+        }
+    }
+}
+
+impl<B: OptionTy, L: OptionTy> SampleParams<B, L, NoneTy> {
+    /// Sets the 'Sample' image operand
+    pub fn sample_index<S>(self, sample_index: S) -> SampleParams<B, L, SomeTy<S>> {
+        SampleParams {
+            bias: self.bias,
+            lod: self.lod,
+            sample_index: SomeTy(sample_index),
+        }
+    }
+}