Use bitmask trait

Author: calebzulawski

crates/core_simd/src/masks.rs

@@ -12,8 +12,10 @@
 )]
 mod mask_impl;
 
-use crate::simd::intrinsics;
-use crate::simd::{LaneCount, Simd, SimdElement, SupportedLaneCount};
+mod to_bitmask;
+pub use to_bitmask::ToBitMask;
+
+use crate::simd::{intrinsics, LaneCount, Simd, SimdElement, SupportedLaneCount};
 use core::cmp::Ordering;
 use core::{fmt, mem};
 
@@ -216,22 +218,6 @@ where
         }
     }
 
-    /// Convert this mask to a bitmask, with one bit set per lane.
-    #[cfg(feature = "generic_const_exprs")]
-    #[inline]
-    #[must_use = "method returns a new array and does not mutate the original value"]
-    pub fn to_bitmask(self) -> [u8; LaneCount::<LANES>::BITMASK_LEN] {
-        self.0.to_bitmask()
-    }
-
-    /// Convert a bitmask to a mask.
-    #[cfg(feature = "generic_const_exprs")]
-    #[inline]
-    #[must_use = "method returns a new mask and does not mutate the original value"]
-    pub fn from_bitmask(bitmask: [u8; LaneCount::<LANES>::BITMASK_LEN]) -> Self {
-        Self(mask_impl::Mask::from_bitmask(bitmask))
-    }
-
     /// Returns true if any lane is set, or false otherwise.
     #[inline]
     #[must_use = "method returns a new bool and does not mutate the original value"]

crates/core_simd/src/masks/bitmask.rs

@@ -115,20 +115,14 @@ where
         unsafe { Self(intrinsics::simd_bitmask(value), PhantomData) }
     }
 
-    #[cfg(feature = "generic_const_exprs")]
     #[inline]
-    #[must_use = "method returns a new array and does not mutate the original value"]
-    pub fn to_bitmask(self) -> [u8; LaneCount::<LANES>::BITMASK_LEN] {
-        // Safety: these are the same type and we are laundering the generic
+    pub unsafe fn to_bitmask_intrinsic<U>(self) -> U {
         unsafe { core::mem::transmute_copy(&self.0) }
     }
 
-    #[cfg(feature = "generic_const_exprs")]
     #[inline]
-    #[must_use = "method returns a new mask and does not mutate the original value"]
-    pub fn from_bitmask(bitmask: [u8; LaneCount::<LANES>::BITMASK_LEN]) -> Self {
-        // Safety: these are the same type and we are laundering the generic
-        Self(unsafe { core::mem::transmute_copy(&bitmask) }, PhantomData)
+    pub unsafe fn from_bitmask_intrinsic<U>(bitmask: U) -> Self {
+        unsafe { Self(core::mem::transmute_copy(&bitmask), PhantomData) }
     }
 
     #[inline]

crates/core_simd/src/masks/full_masks.rs

@@ -109,41 +109,16 @@ where
         unsafe { Mask(intrinsics::simd_cast(self.0)) }
     }
 
-    #[cfg(feature = "generic_const_exprs")]
     #[inline]
-    #[must_use = "method returns a new array and does not mutate the original value"]
-    pub fn to_bitmask(self) -> [u8; LaneCount::<LANES>::BITMASK_LEN] {
-        unsafe {
-            let mut bitmask: [u8; LaneCount::<LANES>::BITMASK_LEN] =
-                intrinsics::simd_bitmask(self.0);
-
-            // There is a bug where LLVM appears to implement this operation with the wrong
-            // bit order.
-            // TODO fix this in a better way
-            if cfg!(target_endian = "big") {
-                for x in bitmask.as_mut() {
-                    *x = x.reverse_bits();
-                }
-            }
-
-            bitmask
-        }
+    pub unsafe fn to_bitmask_intrinsic<U>(self) -> U {
+        // Safety: caller must only return bitmask types
+        unsafe { intrinsics::simd_bitmask(self.0) }
     }
 
-    #[cfg(feature = "generic_const_exprs")]
     #[inline]
-    #[must_use = "method returns a new mask and does not mutate the original value"]
-    pub fn from_bitmask(mut bitmask: [u8; LaneCount::<LANES>::BITMASK_LEN]) -> Self {
+    pub unsafe fn from_bitmask_intrinsic<U>(bitmask: U) -> Self {
+        // Safety: caller must only pass bitmask types
         unsafe {
-            // There is a bug where LLVM appears to implement this operation with the wrong
-            // bit order.
-            // TODO fix this in a better way
-            if cfg!(target_endian = "big") {
-                for x in bitmask.as_mut() {
-                    *x = x.reverse_bits();
-                }
-            }
-
             Self::from_int_unchecked(intrinsics::simd_select_bitmask(
                 bitmask,
                 Self::splat(true).to_int(),

crates/core_simd/src/masks/to_bitmask.rs

@@ -0,0 +1,78 @@
+use super::{mask_impl, Mask, MaskElement};
+
+/// Converts masks to and from bitmasks.
+///
+/// In a bitmask, each bit represents if the corresponding lane in the mask is set.
+pub trait ToBitMask<BitMask> {
+    /// Converts a mask to a bitmask.
+    fn to_bitmask(self) -> BitMask;
+
+    /// Converts a bitmask to a mask.
+    fn from_bitmask(bitmask: BitMask) -> Self;
+}
+
+macro_rules! impl_integer_intrinsic {
+    { $(unsafe impl ToBitMask<$int:ty> for Mask<_, $lanes:literal>)* } => {
+        $(
+        impl<T: MaskElement> ToBitMask<$int> for Mask<T, $lanes> {
+            fn to_bitmask(self) -> $int {
+                unsafe { self.0.to_bitmask_intrinsic() }
+            }
+
+            fn from_bitmask(bitmask: $int) -> Self {
+                unsafe { Self(mask_impl::Mask::from_bitmask_intrinsic(bitmask)) }
+            }
+        }
+        )*
+    }
+}
+
+impl_integer_intrinsic! {
+    unsafe impl ToBitMask<u8> for Mask<_, 8>
+    unsafe impl ToBitMask<u16> for Mask<_, 16>
+    unsafe impl ToBitMask<u32> for Mask<_, 32>
+    unsafe impl ToBitMask<u64> for Mask<_, 64>
+}
+
+macro_rules! impl_integer_via {
+    { $(impl ToBitMask<$int:ty, via $via:ty> for Mask<_, $lanes:literal>)* } => {
+        $(
+        impl<T: MaskElement> ToBitMask<$int> for Mask<T, $lanes> {
+            fn to_bitmask(self) -> $int {
+                let bitmask: $via = self.to_bitmask();
+                bitmask as _
+            }
+
+            fn from_bitmask(bitmask: $int) -> Self {
+                Self::from_bitmask(bitmask as $via)
+            }
+        }
+        )*
+    }
+}
+
+impl_integer_via! {
+    impl ToBitMask<u16, via u8> for Mask<_, 8>
+    impl ToBitMask<u32, via u8> for Mask<_, 8>
+    impl ToBitMask<u64, via u8> for Mask<_, 8>
+
+    impl ToBitMask<u32, via u16> for Mask<_, 16>
+    impl ToBitMask<u64, via u16> for Mask<_, 16>
+
+    impl ToBitMask<u64, via u32> for Mask<_, 32>
+}
+
+#[cfg(target_pointer_width = "32")]
+impl_integer_via! {
+    impl ToBitMask<usize, via u8> for Mask<_, 8>
+    impl ToBitMask<usize, via u16> for Mask<_, 16>
+    impl ToBitMask<usize, via u32> for Mask<_, 32>
+}
+
+#[cfg(target_pointer_width = "64")]
+impl_integer_via! {
+    impl ToBitMask<usize, via u8> for Mask<_, 8>
+    impl ToBitMask<usize, via u16> for Mask<_, 16>
+    impl ToBitMask<usize, via u32> for Mask<_, 32>
+    impl ToBitMask<usize, via u64> for Mask<_, 64>
+}

crates/core_simd/tests/masks.rs

@@ -68,16 +68,16 @@ macro_rules! test_mask_api {
                 assert_eq!(core_simd::Mask::<$type, 8>::from_int(int), mask);
             }
 
-            #[cfg(feature = "generic_const_exprs")]
             #[test]
             fn roundtrip_bitmask_conversion() {
+                use core_simd::ToBitMask;
                 let values = [
                     true, false, false, true, false, false, true, false,
                     true, true, false, false, false, false, false, true,
                 ];
                 let mask = core_simd::Mask::<$type, 16>::from_array(values);
-                let bitmask = mask.to_bitmask();
-                assert_eq!(bitmask, [0b01001001, 0b10000011]);
+                let bitmask: u16 = mask.to_bitmask();
+                assert_eq!(bitmask, 0b1000001101001001);
                 assert_eq!(core_simd::Mask::<$type, 16>::from_bitmask(bitmask), mask);
             }
         }