spirv_std/
scalar_or_vector.rs

1use crate::glam::Vec3A;
2use crate::{Scalar, Vector};
3use core::num::NonZeroUsize;
4
5pub(crate) mod sealed {
6    /// A marker trait used to prevent other traits from being implemented outside
7    /// of `spirv-std`.
8    pub trait Sealed {}
9}
10
11/// Abstract trait representing either a [`Scalar`] or [`Vector`] type.
12///
13/// # Safety
14/// Your type must also implement [`Scalar`] or [`Vector`], see their safety sections as well.
15pub unsafe trait ScalarOrVector: ScalarComposite + Default {
16    /// Either the scalar component type of the vector or the scalar itself.
17    type Scalar: Scalar;
18
19    /// The dimension of the vector, or 1 if it is a scalar
20    const N: NonZeroUsize;
21}
22
23/// A `ScalarComposite` is a type that is either
24/// * a [`Scalar`]
25/// * a [`Vector`] (since vectors are made from scalars)
26/// * an array of `ScalarComposite`
27/// * a struct where all members are `ScalarComposite`
28/// * an enum with a `repr` that is a [`Scalar`]
29///
30/// By calling [`Self::transform`] you can visit all the individual [`Scalar`] and [`Vector`] values this composite is
31/// build out of and transform them into some other value. This is particularly useful for subgroup intrinsics sending
32/// data to other threads.
33///
34/// To derive `ScalarComposite` on a struct, all members must also implement `ScalarComposite`.
35///
36/// To derive `ScalarComposite` on an enum, the enum must implement `From<N>` and `Into<N>` where `N` is defined by the
37/// `#[repr(N)]` attribute on the enum and must be an [`Integer`], like `u32`.
38/// Note that some [safe subgroup operations] may return an "undefined result", so your `From<N>` must gracefully handle
39/// arbitrary bit patterns being passed to it. While panicking is legal, it is discouraged as it may result in
40/// unexpected control flow.
41/// To implement these conversion traits, we recommend [`FromPrimitive`] and [`IntoPrimitive`] from the [`num_enum`]
42/// crate. [`FromPrimitive`] requires the enum to either be exhaustive or have a variant to default to, by either
43/// implementing [`Default`] or marking a variant with `#[num_enum(default)]`. Note to disable default
44/// features on the [`num_enum`] crate, or it won't compile on SPIR-V.
45///
46/// [`Integer`]: crate::Integer
47/// [safe subgroup operations]: crate::arch::subgroup_shuffle
48/// [`FromPrimitive`]: https://docs.rs/num_enum/latest/num_enum/derive.FromPrimitive.html
49/// [`IntoPrimitive`]: https://docs.rs/num_enum/latest/num_enum/derive.IntoPrimitive.html
50/// [`num_enum`]: https://crates.io/crates/num_enum
51pub trait ScalarComposite: Copy + Send + Sync + 'static {
52    /// Transform the individual [`Scalar`] and [`Vector`] values of this type to a different value.
53    ///
54    /// See [`Self`] for more detail.
55    fn transform<F: ScalarOrVectorTransform>(self, f: &mut F) -> Self;
56}
57
58/// A transform operation for [`ScalarComposite::transform`]
59pub trait ScalarOrVectorTransform {
60    /// transform a [`ScalarOrVector`]
61    fn transform<T: ScalarOrVector>(&mut self, value: T) -> T;
62
63    /// transform a [`Scalar`], defaults to [`Self::transform`]
64    #[inline]
65    fn transform_scalar<T: Scalar>(&mut self, value: T) -> T {
66        self.transform(value)
67    }
68
69    /// transform a [`Vector`], defaults to [`Self::transform`]
70    #[inline]
71    fn transform_vector<V: Vector<S, N>, S: Scalar, const N: usize>(&mut self, value: V) -> V {
72        self.transform(value)
73    }
74}
75
76/// `Default` is unfortunately necessary until rust-gpu improves
77impl<T: ScalarComposite + Default, const N: usize> ScalarComposite for [T; N] {
78    #[inline]
79    fn transform<F: ScalarOrVectorTransform>(self, f: &mut F) -> Self {
80        let mut out = [T::default(); N];
81        for i in 0..N {
82            out[i] = self[i].transform(f);
83        }
84        out
85    }
86}
87
88/// `glam::Vec3A` is not considered a SPIR-V vector and thus needs a manual implementation
89impl ScalarComposite for Vec3A {
90    #[inline]
91    fn transform<F: ScalarOrVectorTransform>(self, f: &mut F) -> Self {
92        Vec3A::new(
93            f.transform_scalar(self.x),
94            f.transform_scalar(self.y),
95            f.transform_scalar(self.z),
96        )
97    }
98}
spirv_std/scalar_or_vector.rs

spirv_std/
scalar_or_vector.rs
