Implement traits over scalars with Simd as a self type

Author: calebzulawski

crates/core_simd/src/lib.rs

@@ -1,5 +1,6 @@
 #![no_std]
 #![feature(
+    arbitrary_self_types,
     convert_float_to_int,
     f16,
     core_intrinsics,

crates/core_simd/src/masks.rs

@@ -291,6 +291,80 @@ where
         unsafe { core::intrinsics::simd::simd_reduce_all(self.0) }
     }
 
+    /// Test if each element is equal to the corresponding element in `other`.
+    #[inline]
+    #[must_use = "method returns a new mask and does not mutate the original value"]
+    pub fn simd_eq(self, other: Self) -> Self {
+        // Safety: `self` is a mask vector, and the result of comparison is always a valid mask.
+        unsafe { Self::from_simd_unchecked(core::intrinsics::simd::simd_eq(self.0, other.0)) }
+    }
+
+    /// Test if each element is not equal to the corresponding element in `other`.
+    #[inline]
+    #[must_use = "method returns a new mask and does not mutate the original value"]
+    pub fn simd_ne(self, other: Self) -> Self {
+        // Safety: `self` is a mask vector, and the result of comparison is always a valid mask.
+        unsafe { Self::from_simd_unchecked(core::intrinsics::simd::simd_ne(self.0, other.0)) }
+    }
+
+    /// Test if each element is less than the corresponding element in `other`.
+    #[inline]
+    #[must_use = "method returns a new mask and does not mutate the original value"]
+    pub fn simd_lt(self, other: Self) -> Self {
+        // Safety: `self` is a mask vector, and the result of comparison is always a valid mask.
+        unsafe { Self::from_simd_unchecked(core::intrinsics::simd::simd_lt(self.0, other.0)) }
+    }
+
+    /// Test if each element is less than or equal to the corresponding element in `other`.
+    #[inline]
+    #[must_use = "method returns a new mask and does not mutate the original value"]
+    pub fn simd_le(self, other: Self) -> Self {
+        // Safety: `self` is a mask vector, and the result of comparison is always a valid mask.
+        unsafe { Self::from_simd_unchecked(core::intrinsics::simd::simd_le(self.0, other.0)) }
+    }
+
+    /// Test if each element is greater than the corresponding element in `other`.
+    #[inline]
+    #[must_use = "method returns a new mask and does not mutate the original value"]
+    pub fn simd_gt(self, other: Self) -> Self {
+        // Safety: `self` is a mask vector, and the result of comparison is always a valid mask.
+        unsafe { Self::from_simd_unchecked(core::intrinsics::simd::simd_gt(self.0, other.0)) }
+    }
+
+    /// Test if each element is greater than or equal to the corresponding element in `other`.
+    #[inline]
+    #[must_use = "method returns a new mask and does not mutate the original value"]
+    pub fn simd_ge(self, other: Self) -> Self {
+        // Safety: `self` is a mask vector, and the result of comparison is always a valid mask.
+        unsafe { Self::from_simd_unchecked(core::intrinsics::simd::simd_ge(self.0, other.0)) }
+    }
+
+    /// Returns the element-wise maximum with `other`.
+    #[inline]
+    #[must_use = "method returns a new mask and does not mutate the original value"]
+    pub fn simd_max(self, other: Self) -> Self {
+        self.simd_gt(other).select(other, self)
+    }
+
+    /// Returns the element-wise minimum with `other`.
+    #[inline]
+    #[must_use = "method returns a new mask and does not mutate the original value"]
+    pub fn simd_min(self, other: Self) -> Self {
+        self.simd_lt(other).select(other, self)
+    }
+
+    /// Restrict each element to a certain interval.
+    #[inline]
+    #[must_use = "method returns a new mask and does not mutate the original value"]
+    #[track_caller]
+    pub fn simd_clamp(self, min: Self, max: Self) -> Self {
+        assert!(
+            min.simd_le(max).all(),
+            "each element in `min` must be less than or equal to the corresponding element in `max`",
+        );
+        self.simd_max(min).simd_min(max)
+    }
+
     /// Creates a bitmask from a mask.
     ///
     /// Each bit is set if the corresponding element in the mask is `true`.
@@ -667,3 +741,10 @@ impl_from! { i16 => i32, i64, isize, i8 }
 impl_from! { i32 => i64, isize, i8, i16 }
 impl_from! { i64 => isize, i8, i16, i32 }
 impl_from! { isize => i8, i16, i32, i64 }
+
+impl<T, const N: usize> core::ops::Receiver for Mask<T, N>
+where
+    T: MaskElement,
+{
+    type Target = T;
+}

crates/core_simd/src/simd/cmp/eq.rs

@@ -4,32 +4,41 @@ use crate::simd::{
 };
 
 /// Parallel `PartialEq`.
-pub trait SimdPartialEq<T, const N: usize>
-where
-    T: SimdElement,
-{
+pub trait SimdPartialEq: SimdElement {
     /// Test if each element is equal to the corresponding element in `other`.
     #[must_use = "method returns a new mask and does not mutate the original value"]
-    fn simd_eq(self, other: Self) -> Mask<<T as SimdElement>::Mask, N>;
+    fn simd_eq<const N: usize>(
+        self: Simd<Self, N>,
+        other: Simd<Self, N>,
+    ) -> Mask<<Self as SimdElement>::Mask, N>;
 
     /// Test if each element is not equal to the corresponding element in `other`.
     #[must_use = "method returns a new mask and does not mutate the original value"]
-    fn simd_ne(self, other: Self) -> Mask<<T as SimdElement>::Mask, N>;
+    fn simd_ne<const N: usize>(
+        self: Simd<Self, N>,
+        other: Simd<Self, N>,
+    ) -> Mask<<Self as SimdElement>::Mask, N>;
 }
 
 macro_rules! impl_number {
     { $($number:ty),* } => {
         $(
-        impl<const N: usize> SimdPartialEq<$number, N> for Simd<$number, N> {
+        impl SimdPartialEq for $number {
             #[inline]
-            fn simd_eq(self, other: Self) -> Mask<<$number as SimdElement>::Mask, N> {
+            fn simd_eq<const N: usize>(
+                self: Simd<Self, N>,
+                other: Simd<Self, N>,
+            ) -> Mask<<Self as SimdElement>::Mask, N> {
                 // Safety: `self` is a vector, and the result of the comparison
                 // is always a valid mask.
                 unsafe { Mask::from_simd_unchecked(core::intrinsics::simd::simd_eq(self, other)) }
             }
 
             #[inline]
-            fn simd_ne(self, other: Self) -> Mask<<$number as SimdElement>::Mask, N> {
+            fn simd_ne<const N: usize>(
+                self: Simd<Self, N>,
+                other: Simd<Self, N>,
+            ) -> Mask<<Self as SimdElement>::Mask, N> {
                 // Safety: `self` is a vector, and the result of the comparison
                 // is always a valid mask.
                 unsafe { Mask::from_simd_unchecked(core::intrinsics::simd::simd_ne(self, other)) }
@@ -41,50 +50,26 @@ macro_rules! impl_number {
 
 impl_number! { f16, f32, f64, u8, u16, u32, u64, usize, i8, i16, i32, i64, isize }
 
-macro_rules! impl_mask {
-    { $($integer:ty),* } => {
-        $(
-        impl<const N: usize> SimdPartialEq<$integer, N> for Mask<$integer, N> {
-            #[inline]
-            fn simd_eq(self, other: Self) -> Mask<<$integer as SimdElement>::Mask, N> {
-                // Safety: `self` is a vector, and the result of the comparison
-                // is always a valid mask.
-                unsafe { Self::from_simd_unchecked(core::intrinsics::simd::simd_eq(self.to_simd(), other.to_simd())) }
-            }
-
-            #[inline]
-            fn simd_ne(self, other: Self) -> Mask<<$integer as SimdElement>::Mask, N> {
-                // Safety: `self` is a vector, and the result of the comparison
-                // is always a valid mask.
-                unsafe { Self::from_simd_unchecked(core::intrinsics::simd::simd_ne(self.to_simd(), other.to_simd())) }
-            }
-        }
-        )*
-    }
-}
-
-impl_mask! { i8, i16, i32, i64, isize }
-
-impl<T, const N: usize> SimdPartialEq<*const T, N> for Simd<*const T, N> {
+impl<T> SimdPartialEq for *const T {
     #[inline]
-    fn simd_eq(self, other: Self) -> Mask<isize, N> {
+    fn simd_eq<const N: usize>(self: Simd<Self, N>, other: Simd<Self, N>) -> Mask<isize, N> {
         self.addr().simd_eq(other.addr())
     }
 
     #[inline]
-    fn simd_ne(self, other: Self) -> Mask<isize, N> {
+    fn simd_ne<const N: usize>(self: Simd<Self, N>, other: Simd<Self, N>) -> Mask<isize, N> {
         self.addr().simd_ne(other.addr())
     }
 }
 
-impl<T, const N: usize> SimdPartialEq<*mut T, N> for Simd<*mut T, N> {
+impl<T> SimdPartialEq for *mut T {
     #[inline]
-    fn simd_eq(self, other: Self) -> Mask<isize, N> {
+    fn simd_eq<const N: usize>(self: Simd<Self, N>, other: Simd<Self, N>) -> Mask<isize, N> {
         self.addr().simd_eq(other.addr())
     }
 
     #[inline]
-    fn simd_ne(self, other: Self) -> Mask<isize, N> {
+    fn simd_ne<const N: usize>(self: Simd<Self, N>, other: Simd<Self, N>) -> Mask<isize, N> {
         self.addr().simd_ne(other.addr())
     }
 }