spirv_std/
scalar_or_vector.rs

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