Make implementation functions safe

calebzulawski · calebzulawski · commit af6956c43c79 · 2022-02-11T02:35:28.000Z
diff --git a/crates/core_simd/src/masks/bitmask.rs b/crates/core_simd/src/masks/bitmask.rs
@@ -1,7 +1,7 @@
 #![allow(unused_imports)]
 use super::MaskElement;
 use crate::simd::intrinsics;
-use crate::simd::{LaneCount, Simd, SupportedLaneCount};
+use crate::simd::{LaneCount, Simd, SupportedLaneCount, ToBitMask, ToBitMaskArray};
 use core::marker::PhantomData;
 
 /// A mask where each lane is represented by a single bit.
@@ -115,31 +115,40 @@ where
         unsafe { Self(intrinsics::simd_bitmask(value), PhantomData) }
     }
 
-    // Safety: N must be the exact number of bytes required to hold the bitmask for this mask
     #[inline]
     #[must_use = "method returns a new array and does not mutate the original value"]
-    pub unsafe fn to_bitmask_array<const N: usize>(self) -> [u8; N] {
-        // Safety: these are the same type and we are laundering the generic
+    pub fn to_bitmask_array<const N: usize>(self) -> [u8; N] {
+        assert!(core::mem::size_of::<Self>() == N);
+
+        // Safety: converting an integer to an array of bytes of the same size is safe
         unsafe { core::mem::transmute_copy(&self.0) }
     }
 
     // Safety: N must be the exact number of bytes required to hold the bitmask for this mask
     #[inline]
     #[must_use = "method returns a new mask and does not mutate the original value"]
-    pub unsafe fn from_bitmask_array<const N: usize>(bitmask: [u8; N]) -> Self {
-        // Safety: these are the same type and we are laundering the generic
+    pub fn from_bitmask_array<const N: usize>(bitmask: [u8; N]) -> Self {
+        assert!(core::mem::size_of::<Self>() == N);
+
+        // Safety: converting an array of bytes to an integer of the same size is safe
         Self(unsafe { core::mem::transmute_copy(&bitmask) }, PhantomData)
     }
 
-    // Safety: U must be the integer with the exact number of bits required to hold the bitmask for
     #[inline]
-    pub unsafe fn to_bitmask_integer<U>(self) -> U {
+    pub fn to_bitmask_integer<U>(self) -> U
+    where
+        super::Mask<T, LANES>: ToBitMask<BitMask = U>,
+    {
+        // Safety: these are the same types
         unsafe { core::mem::transmute_copy(&self.0) }
     }
 
-    // Safety: U must be the integer with the exact number of bits required to hold the bitmask for
     #[inline]
-    pub unsafe fn from_bitmask_integer<U>(bitmask: U) -> Self {
+    pub fn from_bitmask_integer<U>(bitmask: U) -> Self
+    where
+        super::Mask<T, LANES>: ToBitMask<BitMask = U>,
+    {
+        // Safety: these are the same types
         unsafe { Self(core::mem::transmute_copy(&bitmask), PhantomData) }
     }
 
diff --git a/crates/core_simd/src/masks/full_masks.rs b/crates/core_simd/src/masks/full_masks.rs
@@ -2,7 +2,7 @@
 
 use super::MaskElement;
 use crate::simd::intrinsics;
-use crate::simd::{LaneCount, Simd, SupportedLaneCount};
+use crate::simd::{LaneCount, Simd, SupportedLaneCount, ToBitMask, ToBitMaskArray};
 
 #[repr(transparent)]
 pub struct Mask<T, const LANES: usize>(Simd<T, LANES>)
@@ -126,12 +126,26 @@ where
         unsafe { Mask(intrinsics::simd_cast(self.0)) }
     }
 
-    // Safety: N must be the exact number of bytes required to hold the bitmask for this mask
     #[inline]
     #[must_use = "method returns a new array and does not mutate the original value"]
-    pub unsafe fn to_bitmask_array<const N: usize>(self) -> [u8; N] {
+    pub fn to_bitmask_array<const N: usize>(self) -> [u8; N]
+    where
+        super::Mask<T, LANES>: ToBitMaskArray,
+        [(); <super::Mask<T, LANES> as ToBitMaskArray>::BYTES]: Sized,
+    {
+        assert_eq!(<super::Mask<T, LANES> as ToBitMaskArray>::BYTES, N);
+
+        // Safety: N is the correct bitmask size
+        //
+        // The transmute below allows this function to be marked safe, since it will prevent
+        // monomorphization errors in the case of an incorrect size.
         unsafe {
-            let mut bitmask: [u8; N] = intrinsics::simd_bitmask(self.0);
+            // Compute the bitmask
+            let bitmask: [u8; <super::Mask<T, LANES> as ToBitMaskArray>::BYTES] =
+                intrinsics::simd_bitmask(self.0);
+
+            // Transmute to the return type, previously asserted to be the same size
+            let mut bitmask: [u8; N] = core::mem::transmute_copy(&bitmask);
 
             // There is a bug where LLVM appears to implement this operation with the wrong
             // bit order.
@@ -146,10 +160,19 @@ where
         }
     }
 
-    // Safety: N must be the exact number of bytes required to hold the bitmask for this mask
     #[inline]
     #[must_use = "method returns a new mask and does not mutate the original value"]
-    pub unsafe fn from_bitmask_array<const N: usize>(mut bitmask: [u8; N]) -> Self {
+    pub fn from_bitmask_array<const N: usize>(mut bitmask: [u8; N]) -> Self
+    where
+        super::Mask<T, LANES>: ToBitMaskArray,
+        [(); <super::Mask<T, LANES> as ToBitMaskArray>::BYTES]: Sized,
+    {
+        assert_eq!(<super::Mask<T, LANES> as ToBitMaskArray>::BYTES, N);
+
+        // Safety: N is the correct bitmask size
+        //
+        // The transmute below allows this function to be marked safe, since it will prevent
+        // monomorphization errors in the case of an incorrect size.
         unsafe {
             // There is a bug where LLVM appears to implement this operation with the wrong
             // bit order.
@@ -160,6 +183,11 @@ where
                 }
             }
 
+            // Transmute to the bitmask type, previously asserted to be the same size
+            let bitmask: [u8; <super::Mask<T, LANES> as ToBitMaskArray>::BYTES] =
+                core::mem::transmute_copy(&bitmask);
+
+            // Compute the regular mask
             Self::from_int_unchecked(intrinsics::simd_select_bitmask(
                 bitmask,
                 Self::splat(true).to_int(),
@@ -168,11 +196,12 @@ where
         }
     }
 
-    // Safety: U must be the integer with the exact number of bits required to hold the bitmask for
-    // this mask
     #[inline]
-    pub(crate) unsafe fn to_bitmask_integer<U: ReverseBits>(self) -> U {
-        // Safety: caller must only return bitmask types
+    pub(crate) fn to_bitmask_integer<U: ReverseBits>(self) -> U
+    where
+        super::Mask<T, LANES>: ToBitMask<BitMask = U>,
+    {
+        // Safety: U is required to be the appropriate bitmask type
         let bitmask: U = unsafe { intrinsics::simd_bitmask(self.0) };
 
         // There is a bug where LLVM appears to implement this operation with the wrong
@@ -188,7 +217,10 @@ where
     // Safety: U must be the integer with the exact number of bits required to hold the bitmask for
     // this mask
     #[inline]
-    pub(crate) unsafe fn from_bitmask_integer<U: ReverseBits>(bitmask: U) -> Self {
+    pub(crate) fn from_bitmask_integer<U: ReverseBits>(bitmask: U) -> Self
+    where
+        super::Mask<T, LANES>: ToBitMask<BitMask = U>,
+    {
         // There is a bug where LLVM appears to implement this operation with the wrong
         // bit order.
         // TODO fix this in a better way
@@ -198,7 +230,7 @@ where
             bitmask
         };
 
-        // Safety: caller must only pass bitmask types
+        // Safety: U is required to be the appropriate bitmask type
         unsafe {
             Self::from_int_unchecked(intrinsics::simd_select_bitmask(
                 bitmask,
diff --git a/crates/core_simd/src/masks/to_bitmask.rs b/crates/core_simd/src/masks/to_bitmask.rs
@@ -1,9 +1,26 @@
 use super::{mask_impl, Mask, MaskElement};
+use crate::simd::{LaneCount, SupportedLaneCount};
+
+mod sealed {
+    pub trait Sealed {}
+}
+pub use sealed::Sealed;
+
+impl<T, const LANES: usize> Sealed for Mask<T, LANES>
+where
+    T: MaskElement,
+    LaneCount<LANES>: SupportedLaneCount,
+{
+}
 
 /// Converts masks to and from integer bitmasks.
 ///
 /// Each bit of the bitmask corresponds to a mask lane, starting with the LSB.
-pub trait ToBitMask {
+///
+/// # Safety
+/// This trait is `unsafe` and sealed, since the `BitMask` type must match the number of lanes in
+/// the mask.
+pub unsafe trait ToBitMask: Sealed {
     /// The integer bitmask type.
     type BitMask;
 
@@ -17,7 +34,11 @@ pub trait ToBitMask {
 /// Converts masks to and from byte array bitmasks.
 ///
 /// Each bit of the bitmask corresponds to a mask lane, starting with the LSB of the first byte.
-pub trait ToBitMaskArray {
+///
+/// # Safety
+/// This trait is `unsafe` and sealed, since the `BYTES` value must match the number of lanes in
+/// the mask.
+pub unsafe trait ToBitMaskArray: Sealed {
     /// The length of the bitmask array.
     const BYTES: usize;
 
@@ -31,15 +52,15 @@ pub trait ToBitMaskArray {
 macro_rules! impl_integer_intrinsic {
     { $(unsafe impl ToBitMask<BitMask=$int:ty> for Mask<_, $lanes:literal>)* } => {
         $(
-        impl<T: MaskElement> ToBitMask for Mask<T, $lanes> {
+        unsafe impl<T: MaskElement> ToBitMask for Mask<T, $lanes> {
             type BitMask = $int;
 
             fn to_bitmask(self) -> $int {
-                unsafe { self.0.to_bitmask_integer() }
+                self.0.to_bitmask_integer()
             }
 
             fn from_bitmask(bitmask: $int) -> Self {
-                unsafe { Self(mask_impl::Mask::from_bitmask_integer(bitmask)) }
+                Self(mask_impl::Mask::from_bitmask_integer(bitmask))
             }
         }
         )*
@@ -59,20 +80,18 @@ pub const fn bitmask_len(lanes: usize) -> usize {
 }
 
 macro_rules! impl_array_bitmask {
-    { $(impl ToBitMask<[u8; _]> for Mask<_, $lanes:literal>)* } => {
+    { $(impl ToBitMaskArray<[u8; _]> for Mask<_, $lanes:literal>)* } => {
         $(
-        impl<T: MaskElement> ToBitMaskArray for Mask<T, $lanes>
+        unsafe impl<T: MaskElement> ToBitMaskArray for Mask<T, $lanes>
         {
              const BYTES: usize = bitmask_len($lanes);
 
              fn to_bitmask_array(self) -> [u8; Self::BYTES] {
-                 // Safety: BYTES is the exact size required
-                 unsafe { self.0.to_bitmask_array() }
+                 self.0.to_bitmask_array()
              }
 
              fn from_bitmask_array(bitmask: [u8; Self::BYTES]) -> Self {
-                 // Safety: BYTES is the exact size required
-                 unsafe { Mask(mask_impl::Mask::from_bitmask_array(bitmask)) }
+                 Mask(mask_impl::Mask::from_bitmask_array(bitmask))
              }
         }
         )*
@@ -81,11 +100,11 @@ macro_rules! impl_array_bitmask {
 
 // FIXME this should be specified generically, but it doesn't seem to work with rustc, yet
 impl_array_bitmask! {
-    impl ToBitMask<[u8; _]> for Mask<_, 1>
-    impl ToBitMask<[u8; _]> for Mask<_, 2>
-    impl ToBitMask<[u8; _]> for Mask<_, 4>
-    impl ToBitMask<[u8; _]> for Mask<_, 8>
-    impl ToBitMask<[u8; _]> for Mask<_, 16>
-    impl ToBitMask<[u8; _]> for Mask<_, 32>
-    impl ToBitMask<[u8; _]> for Mask<_, 64>
+    impl ToBitMaskArray<[u8; _]> for Mask<_, 1>
+    impl ToBitMaskArray<[u8; _]> for Mask<_, 2>
+    impl ToBitMaskArray<[u8; _]> for Mask<_, 4>
+    impl ToBitMaskArray<[u8; _]> for Mask<_, 8>
+    impl ToBitMaskArray<[u8; _]> for Mask<_, 16>
+    impl ToBitMaskArray<[u8; _]> for Mask<_, 32>
+    impl ToBitMaskArray<[u8; _]> for Mask<_, 64>
 }
