Add comparison ops

5 files changed, 414 insertions, 83 deletions

diff --git a/crates/core_simd/src/intrinsics.rs b/crates/core_simd/src/intrinsics.rs
index b2edc3747ef..3dfc77136f2 100644
--- a/crates/core_simd/src/intrinsics.rs
+++ b/crates/core_simd/src/intrinsics.rs
@@ -45,4 +45,11 @@ extern "platform-intrinsic" {
 
     // ceil
     pub(crate) fn simd_ceil<T>(x: T) -> T;
+
+    pub(crate) fn simd_eq<T, U>(x: T, y: T) -> U;
+    pub(crate) fn simd_ne<T, U>(x: T, y: T) -> U;
+    pub(crate) fn simd_lt<T, U>(x: T, y: T) -> U;
+    pub(crate) fn simd_le<T, U>(x: T, y: T) -> U;
+    pub(crate) fn simd_gt<T, U>(x: T, y: T) -> U;
+    pub(crate) fn simd_ge<T, U>(x: T, y: T) -> U;
 }
diff --git a/crates/core_simd/src/lib.rs b/crates/core_simd/src/lib.rs
index fd4f9dd16fd..3c581ad659b 100644
--- a/crates/core_simd/src/lib.rs
+++ b/crates/core_simd/src/lib.rs
@@ -11,6 +11,7 @@ mod intrinsics;
 mod ops;
 
 pub mod masks;
+pub use masks::opaque::*;
 
 mod vectors_u8;
 pub use vectors_u8::*;
diff --git a/crates/core_simd/src/macros.rs b/crates/core_simd/src/macros.rs
index b8324ffdb92..75584f58b78 100644
--- a/crates/core_simd/src/macros.rs
+++ b/crates/core_simd/src/macros.rs
@@ -336,6 +336,24 @@ macro_rules! define_mask_vector {
             call_repeat! { $lanes => define_mask_vector [$impl_type] splat $type | }
             call_counting_args! { $lanes => define_mask_vector => new $type | }
             call_counting_args! { $lanes => define_mask_vector => new_from_bool $type | }
+
+            /// Tests the value of the specified lane.
+            ///
+            /// # Panics
+            /// Panics if `lane` is greater than or equal to the number of lanes in the vector.
+            #[inline]
+            pub fn test(&self, lane: usize) -> bool {
+                self[lane].test()
+            }
+
+            /// Sets the value of the specified lane.
+            ///
+            /// # Panics
+            /// Panics if `lane` is greater than or equal to the number of lanes in the vector.
+            #[inline]
+            pub fn set(&mut self, lane: usize, value: bool) {
+                self[lane] = value.into();
+            }
         }
 
         base_vector_traits! { $name => [$type; $lanes] }
diff --git a/crates/core_simd/src/masks/mod.rs b/crates/core_simd/src/masks/mod.rs
index e138a1b4dd8..6688db290e2 100644
--- a/crates/core_simd/src/masks/mod.rs
+++ b/crates/core_simd/src/masks/mod.rs
@@ -2,7 +2,10 @@
 
 pub mod wide;
 
-trait MaskImpl {
+mod ops;
+pub use ops::*;
+
+pub(crate) trait MaskImpl {
     type Mask;
 }
 
@@ -93,15 +96,67 @@ macro_rules! define_opaque_mask {
     } => {
         $(#[$attr])*
         #[allow(non_camel_case_types)]
-        pub struct $name(<[$width; $lanes] as MaskImpl>::Mask);
+        pub struct $name(<[$width; $lanes] as crate::masks::MaskImpl>::Mask);
 
         impl $name {
+            pub(crate) fn new_from_inner(inner: <[$width; $lanes] as crate::masks::MaskImpl>::Mask) -> Self {
+                Self(inner)
+            }
+
             /// Construct a mask by setting all lanes to the given value.
             pub fn splat(value: bool) -> Self {
-                Self(<[$width; $lanes] as MaskImpl>::Mask::splat(value.into()))
+                Self(<[$width; $lanes] as crate::masks::MaskImpl>::Mask::splat(value.into()))
             }
 
             call_counting_args! { $lanes => define_opaque_mask => new [$width; $lanes] }
+
+            /// Tests the value of the specified lane.
+            ///
+            /// # Panics
+            /// Panics if `lane` is greater than or equal to the number of lanes in the vector.
+            #[inline]
+            pub fn test(&self, lane: usize) -> bool {
+                self.0.test(lane)
+            }
+
+            /// Sets the value of the specified lane.
+            ///
+            /// # Panics
+            /// Panics if `lane` is greater than or equal to the number of lanes in the vector.
+            #[inline]
+            pub fn set(&mut self, lane: usize, value: bool) {
+                self.0.set(lane, value);
+            }
+        }
+
+        impl Copy for $name {}
+
+        impl Clone for $name {
+            #[inline]
+            fn clone(&self) -> Self {
+                *self
+            }
+        }
+
+        impl Default for $name {
+            #[inline]
+            fn default() -> Self {
+                Self::splat(false)
+            }
+        }
+
+        impl PartialEq for $name {
+            #[inline]
+            fn eq(&self, other: &Self) -> bool {
+                self.0 == other.0
+            }
+        }
+
+        impl PartialOrd for $name {
+            #[inline]
+            fn partial_cmp(&self, other: &Self) -> Option<core::cmp::Ordering> {
+                self.0.partial_cmp(&other.0)
+            }
         }
     };
     { new [$width:ty; $lanes:tt] $($var:ident)* } => {
@@ -109,107 +164,149 @@ macro_rules! define_opaque_mask {
         #[allow(clippy::too_many_arguments)]
         #[inline]
         pub const fn new($($var: bool),*) -> Self {
-            Self(<[$width; $lanes] as MaskImpl>::Mask::new_from_bool($($var),*))
+            Self(<[$width; $lanes] as crate::masks::MaskImpl>::Mask::new_from_bool($($var),*))
         }
     }
 }
 
-define_opaque_mask! {
-    /// Mask for 8 8-bit lanes
-    struct mask8x8([u8; 8]);
-}
+pub(crate) mod opaque {
+    define_opaque_mask! {
+        /// Mask for 8 8-bit lanes.
+        ///
+        /// The layout of this type is unspecified.
+        struct mask8x8([u8; 8]);
+    }
 
-define_opaque_mask! {
-    /// Mask for 16 8-bit lanes
-    struct mask8x16([u8; 16]);
-}
+    define_opaque_mask! {
+        /// Mask for 16 8-bit lanes.
+        ///
+        /// The layout of this type is unspecified.
+        struct mask8x16([u8; 16]);
+    }
 
-define_opaque_mask! {
-    /// Mask for 32 8-bit lanes
-    struct mask8x32([u8; 32]);
-}
+    define_opaque_mask! {
+        /// Mask for 32 8-bit lanes.
+        ///
+        /// The layout of this type is unspecified.
+        struct mask8x32([u8; 32]);
+    }
 
-define_opaque_mask! {
-    /// Mask for 64 8-bit lanes
-    struct mask8x64([u8; 64]);
-}
+    define_opaque_mask! {
+        /// Mask for 64 8-bit lanes.
+        ///
+        /// The layout of this type is unspecified.
+        struct mask8x64([u8; 64]);
+    }
 
-define_opaque_mask! {
-    /// Mask for 4 16-bit lanes
-    struct mask16x4([u16; 4]);
-}
+    define_opaque_mask! {
+        /// Mask for 4 16-bit lanes.
+        ///
+        /// The layout of this type is unspecified.
+        struct mask16x4([u16; 4]);
+    }
 
-define_opaque_mask! {
-    /// Mask for 8 16-bit lanes
-    struct mask16x8([u16; 8]);
-}
+    define_opaque_mask! {
+        /// Mask for 8 16-bit lanes.
+        ///
+        /// The layout of this type is unspecified.
+        struct mask16x8([u16; 8]);
+    }
 
-define_opaque_mask! {
-    /// Mask for 16 16-bit lanes
-    struct mask16x16([u16; 16]);
-}
+    define_opaque_mask! {
+        /// Mask for 16 16-bit lanes.
+        ///
+        /// The layout of this type is unspecified.
+        struct mask16x16([u16; 16]);
+    }
 
-define_opaque_mask! {
-    /// Mask for 32 16-bit lanes
-    struct mask16x32([u16; 32]);
-}
+    define_opaque_mask! {
+        /// Mask for 32 16-bit lanes.
+        ///
+        /// The layout of this type is unspecified.
+        struct mask16x32([u16; 32]);
+    }
 
-define_opaque_mask! {
-    /// Mask for 2 32-bit lanes
-    struct mask32x2([u32; 2]);
-}
+    define_opaque_mask! {
+        /// Mask for 2 32-bit lanes.
+        ///
+        /// The layout of this type is unspecified.
+        struct mask32x2([u32; 2]);
+    }
 
-define_opaque_mask! {
-    /// Mask for 4 32-bit lanes
-    struct mask32x4([u32; 4]);
-}
+    define_opaque_mask! {
+        /// Mask for 4 32-bit lanes.
+        ///
+        /// The layout of this type is unspecified.
+        struct mask32x4([u32; 4]);
+    }
 
-define_opaque_mask! {
-    /// Mask for 8 32-bit lanes
-    struct mask32x8([u32; 8]);
-}
+    define_opaque_mask! {
+        /// Mask for 8 32-bit lanes.
+        ///
+        /// The layout of this type is unspecified.
+        struct mask32x8([u32; 8]);
+    }
 
-define_opaque_mask! {
-    /// Mask for 16 32-bit lanes
-    struct mask32x16([u32; 16]);
-}
+    define_opaque_mask! {
+        /// Mask for 16 32-bit lanes.
+        ///
+        /// The layout of this type is unspecified.
+        struct mask32x16([u32; 16]);
+    }
 
-define_opaque_mask! {
-    /// Mask for 2 64-bit lanes
-    struct mask64x2([u64; 2]);
-}
+    define_opaque_mask! {
+        /// Mask for 2 64-bit lanes.
+        ///
+        /// The layout of this type is unspecified.
+        struct mask64x2([u64; 2]);
+    }
 
-define_opaque_mask! {
-    /// Mask for 4 64-bit lanes
-    struct mask64x4([u64; 4]);
-}
+    define_opaque_mask! {
+        /// Mask for 4 64-bit lanes.
+        ///
+        /// The layout of this type is unspecified.
+        struct mask64x4([u64; 4]);
+    }
 
-define_opaque_mask! {
-    /// Mask for 8 64-bit lanes
-    struct mask64x8([u64; 8]);
-}
+    define_opaque_mask! {
+        /// Mask for 8 64-bit lanes.
+        ///
+        /// The layout of this type is unspecified.
+        struct mask64x8([u64; 8]);
+    }
 
-define_opaque_mask! {
-    /// Mask for 2 128-bit lanes
-    struct mask128x2([u128; 2]);
-}
+    define_opaque_mask! {
+        /// Mask for 2 128-bit lanes.
+        ///
+        /// The layout of this type is unspecified.
+        struct mask128x2([u128; 2]);
+    }
 
-define_opaque_mask! {
-    /// Mask for 4 128-bit lanes
-    struct mask128x4([u128; 4]);
-}
+    define_opaque_mask! {
+        /// Mask for 4 128-bit lanes.
+        ///
+        /// The layout of this type is unspecified.
+        struct mask128x4([u128; 4]);
+    }
 
-define_opaque_mask! {
-    /// Mask for 2 `isize`-wide lanes
-    struct masksizex2([usize; 2]);
-}
+    define_opaque_mask! {
+        /// Mask for 2 `isize`-wide lanes.
+        ///
+        /// The layout of this type is unspecified.
+        struct masksizex2([usize; 2]);
+    }
 
-define_opaque_mask! {
-    /// Mask for 4 `isize`-wide lanes
-    struct masksizex4([usize; 4]);
-}
+    define_opaque_mask! {
+        /// Mask for 4 `isize`-wide lanes.
+        ///
+        /// The layout of this type is unspecified.
+        struct masksizex4([usize; 4]);
+    }
 
-define_opaque_mask! {
-    /// Mask for 8 `isize`-wide lanes
-    struct masksizex8([usize; 8]);
+    define_opaque_mask! {
+        /// Mask for 8 `isize`-wide lanes.
+        ///
+        /// The layout of this type is unspecified.
+        struct masksizex8([usize; 8]);
+    }
 }
diff --git a/crates/core_simd/src/masks/ops.rs b/crates/core_simd/src/masks/ops.rs
new file mode 100644
index 00000000000..85ce955459a
--- /dev/null
+++ b/crates/core_simd/src/masks/ops.rs
@@ -0,0 +1,208 @@
+/// Mask-related operations using a particular mask layout.
+pub trait MaskExt<Mask> {
+    /// Test if each lane is equal to the corresponding lane in `other`.
+    fn lanes_eq(self, other: Self) -> Mask;
+
+    /// Test if each lane is not equal to the corresponding lane in `other`.
+    fn lanes_ne(self, other: Self) -> Mask;
+
+    /// Test if each lane is less than the corresponding lane in `other`.
+    fn lanes_lt(self, other: Self) -> Mask;
+
+    /// Test if each lane is greater than the corresponding lane in `other`.
+    fn lanes_gt(self, other: Self) -> Mask;
+
+    /// Test if each lane is less than or equal to the corresponding lane in `other`.
+    fn lanes_le(self, other: Self) -> Mask;
+
+    /// Test if each lane is greater than or equal to the corresponding lane in `other`.
+    fn lanes_ge(self, other: Self) -> Mask;
+}
+
+macro_rules! implement_mask_ext {
+    { $($vector:ty => $($mask:ty),*;)* } => {
+        $( // vector
+            $( // mask
+                impl MaskExt<$mask> for $vector {
+                    #[inline]
+                    fn lanes_eq(self, other: Self) -> $mask {
+                        unsafe { crate::intrinsics::simd_eq(self, other) }
+                    }
+
+                    #[inline]
+                    fn lanes_ne(self, other: Self) -> $mask {
+                        unsafe { crate::intrinsics::simd_ne(self, other) }
+                    }
+
+                    #[inline]
+                    fn lanes_lt(self, other: Self) -> $mask {
+                        unsafe { crate::intrinsics::simd_lt(self, other) }
+                    }
+
+                    #[inline]
+                    fn lanes_gt(self, other: Self) -> $mask {
+                        unsafe { crate::intrinsics::simd_gt(self, other) }
+                    }
+
+                    #[inline]
+                    fn lanes_le(self, other: Self) -> $mask {
+                        unsafe { crate::intrinsics::simd_le(self, other) }
+                    }
+
+                    #[inline]
+                    fn lanes_ge(self, other: Self) -> $mask {
+                        unsafe { crate::intrinsics::simd_ge(self, other) }
+                    }
+                }
+            )*
+        )*
+    }
+}
+
+implement_mask_ext! {
+    crate::u8x8 => crate::masks::wide::m8x8;
+    crate::u8x16 => crate::masks::wide::m8x16;
+    crate::u8x32 => crate::masks::wide::m8x32;
+    crate::u8x64 => crate::masks::wide::m8x64;
+    crate::u16x4 => crate::masks::wide::m16x4;
+    crate::u16x8 => crate::masks::wide::m16x8;
+    crate::u16x16 => crate::masks::wide::m16x16;
+    crate::u16x32 => crate::masks::wide::m16x32;
+    crate::u32x2 => crate::masks::wide::m32x2;
+    crate::u32x4 => crate::masks::wide::m32x4;
+    crate::u32x8 => crate::masks::wide::m32x8;
+    crate::u32x16 => crate::masks::wide::m32x16;
+    crate::u64x2 => crate::masks::wide::m64x2;
+    crate::u64x4 => crate::masks::wide::m64x4;
+    crate::u64x8 => crate::masks::wide::m64x8;
+    crate::u128x2 => crate::masks::wide::m128x2;
+    crate::u128x4 => crate::masks::wide::m128x4;
+    crate::usizex2 => crate::masks::wide::msizex2;
+    crate::usizex4 => crate::masks::wide::msizex4;
+    crate::usizex8 => crate::masks::wide::msizex8;
+
+    crate::i8x8 => crate::masks::wide::m8x8;
+    crate::i8x16 => crate::masks::wide::m8x16;
+    crate::i8x32 => crate::masks::wide::m8x32;
+    crate::i8x64 => crate::masks::wide::m8x64;
+    crate::i16x4 => crate::masks::wide::m16x4;
+    crate::i16x8 => crate::masks::wide::m16x8;
+    crate::i16x16 => crate::masks::wide::m16x16;
+    crate::i16x32 => crate::masks::wide::m16x32;
+    crate::i32x2 => crate::masks::wide::m32x2;
+    crate::i32x4 => crate::masks::wide::m32x4;
+    crate::i32x8 => crate::masks::wide::m32x8;
+    crate::i32x16 => crate::masks::wide::m32x16;
+    crate::i64x2 => crate::masks::wide::m64x2;
+    crate::i64x4 => crate::masks::wide::m64x4;
+    crate::i64x8 => crate::masks::wide::m64x8;
+    crate::i128x2 => crate::masks::wide::m128x2;
+    crate::i128x4 => crate::masks::wide::m128x4;
+    crate::isizex2 => crate::masks::wide::msizex2;
+    crate::isizex4 => crate::masks::wide::msizex4;
+    crate::isizex8 => crate::masks::wide::msizex8;
+
+    crate::f32x2 => crate::masks::wide::m32x2;
+    crate::f32x4 => crate::masks::wide::m32x4;
+    crate::f32x8 => crate::masks::wide::m32x8;
+    crate::f32x16 => crate::masks::wide::m32x16;
+    crate::f64x2 => crate::masks::wide::m64x2;
+    crate::f64x4 => crate::masks::wide::m64x4;
+    crate::f64x8 => crate::masks::wide::m64x8;
+}
+
+macro_rules! implement_mask_ops {
+    { $($vector:ty => $mask:ty,)* } => {
+        $( // vector
+            impl $vector {
+                /// Test if each lane is equal to the corresponding lane in `other`.
+                #[inline]
+                pub fn lanes_eq(self, other: Self) -> $mask {
+                    <$mask>::new_from_inner(MaskExt::lanes_eq(self, other))
+                }
+
+                /// Test if each lane is not equal to the corresponding lane in `other`.
+                #[inline]
+                pub fn lanes_ne(self, other: Self) -> $mask {
+                    <$mask>::new_from_inner(MaskExt::lanes_ne(self, other))
+                }
+
+                /// Test if each lane is less than the corresponding lane in `other`.
+                #[inline]
+                pub fn lanes_lt(self, other: Self) -> $mask {
+                    <$mask>::new_from_inner(MaskExt::lanes_lt(self, other))
+                }
+
+                /// Test if each lane is greater than the corresponding lane in `other`.
+                #[inline]
+                pub fn lanes_gt(self, other: Self) -> $mask {
+                    <$mask>::new_from_inner(MaskExt::lanes_gt(self, other))
+                }
+
+                /// Test if each lane is less than or equal to the corresponding lane in `other`.
+                #[inline]
+                pub fn lanes_le(self, other: Self) -> $mask {
+                    <$mask>::new_from_inner(MaskExt::lanes_le(self, other))
+                }
+
+                /// Test if each lane is greater than or equal to the corresponding lane in `other`.
+                #[inline]
+                pub fn lanes_ge(self, other: Self) -> $mask {
+                    <$mask>::new_from_inner(MaskExt::lanes_ge(self, other))
+                }
+            }
+        )*
+    }
+}
+
+implement_mask_ops! {
+    crate::u8x8 => crate::mask8x8,
+    crate::u8x16 => crate::mask8x16,
+    crate::u8x32 => crate::mask8x32,
+    crate::u8x64 => crate::mask8x64,
+    crate::u16x4 => crate::mask16x4,
+    crate::u16x8 => crate::mask16x8,
+    crate::u16x16 => crate::mask16x16,
+    crate::u16x32 => crate::mask16x32,
+    crate::u32x2 => crate::mask32x2,
+    crate::u32x4 => crate::mask32x4,
+    crate::u32x8 => crate::mask32x8,
+    crate::u32x16 => crate::mask32x16,
+    crate::u64x2 => crate::mask64x2,
+    crate::u64x4 => crate::mask64x4,
+    crate::u64x8 => crate::mask64x8,
+    crate::u128x2 => crate::mask128x2,
+    crate::u128x4 => crate::mask128x4,
+    crate::usizex2 => crate::masksizex2,
+    crate::usizex4 => crate::masksizex4,
+    crate::usizex8 => crate::masksizex8,
+
+    crate::i8x8 => crate::mask8x8,
+    crate::i8x16 => crate::mask8x16,
+    crate::i8x32 => crate::mask8x32,
+    crate::i8x64 => crate::mask8x64,
+    crate::i16x4 => crate::mask16x4,
+    crate::i16x8 => crate::mask16x8,
+    crate::i16x16 => crate::mask16x16,
+    crate::i16x32 => crate::mask16x32,
+    crate::i32x2 => crate::mask32x2,
+    crate::i32x4 => crate::mask32x4,
+    crate::i32x8 => crate::mask32x8,
+    crate::i32x16 => crate::mask32x16,
+    crate::i64x2 => crate::mask64x2,
+    crate::i64x4 => crate::mask64x4,
+    crate::i64x8 => crate::mask64x8,
+    crate::i128x2 => crate::mask128x2,
+    crate::i128x4 => crate::mask128x4,
+    crate::isizex2 => crate::masksizex2,
+    crate::isizex4 => crate::masksizex4,
+    crate::isizex8 => crate::masksizex8,
+
+    crate::f32x2 => crate::mask32x2,
+    crate::f32x4 => crate::mask32x4,
+    crate::f32x8 => crate::mask32x8,
+    crate::f32x16 => crate::mask32x16,
+    crate::f64x2 => crate::mask64x2,
+    crate::f64x4 => crate::mask64x4,
+    crate::f64x8 => crate::mask64x8,
+}