library/portable-simd/crates/core_simd/src/masks/bitmask.rs

   1 #![allow(unused_imports)]
   2 use super::MaskElement;
   3 use crate::simd::intrinsics;
   4 use crate::simd::{LaneCount, Simd, SupportedLaneCount};
   5 use core::marker::PhantomData;
   6
   7 /// A mask where each lane is represented by a single bit.
   8 #[repr(transparent)]
   9 pub struct Mask<T, const LANES: usize>(
  10     <LaneCount<LANES> as SupportedLaneCount>::BitMask,
  11     PhantomData<T>,
  12 )
  13 where
  14     T: MaskElement,
  15     LaneCount<LANES>: SupportedLaneCount;
  16
  17 impl<T, const LANES: usize> Copy for Mask<T, LANES>
  18 where
  19     T: MaskElement,
  20     LaneCount<LANES>: SupportedLaneCount,
  21 {
  22 }
  23
  24 impl<T, const LANES: usize> Clone for Mask<T, LANES>
  25 where
  26     T: MaskElement,
  27     LaneCount<LANES>: SupportedLaneCount,
  28 {
  29     fn clone(&self) -> Self {
  30         *self
  31     }
  32 }
  33
  34 impl<T, const LANES: usize> PartialEq for Mask<T, LANES>
  35 where
  36     T: MaskElement,
  37     LaneCount<LANES>: SupportedLaneCount,
  38 {
  39     fn eq(&self, other: &Self) -> bool {
  40         self.0.as_ref() == other.0.as_ref()
  41     }
  42 }
  43
  44 impl<T, const LANES: usize> PartialOrd for Mask<T, LANES>
  45 where
  46     T: MaskElement,
  47     LaneCount<LANES>: SupportedLaneCount,
  48 {
  49     fn partial_cmp(&self, other: &Self) -> Option<core::cmp::Ordering> {
  50         self.0.as_ref().partial_cmp(other.0.as_ref())
  51     }
  52 }
  53
  54 impl<T, const LANES: usize> Eq for Mask<T, LANES>
  55 where
  56     T: MaskElement,
  57     LaneCount<LANES>: SupportedLaneCount,
  58 {
  59 }
  60
  61 impl<T, const LANES: usize> Ord for Mask<T, LANES>
  62 where
  63     T: MaskElement,
  64     LaneCount<LANES>: SupportedLaneCount,
  65 {
  66     fn cmp(&self, other: &Self) -> core::cmp::Ordering {
  67         self.0.as_ref().cmp(other.0.as_ref())
  68     }
  69 }
  70
  71 impl<T, const LANES: usize> Mask<T, LANES>
  72 where
  73     T: MaskElement,
  74     LaneCount<LANES>: SupportedLaneCount,
  75 {
  76     #[inline]
  77     #[must_use = "method returns a new mask and does not mutate the original value"]
  78     pub fn splat(value: bool) -> Self {
  79         let mut mask = <LaneCount<LANES> as SupportedLaneCount>::BitMask::default();
  80         if value {
  81             mask.as_mut().fill(u8::MAX)
  82         } else {
  83             mask.as_mut().fill(u8::MIN)
  84         }
  85         if LANES % 8 > 0 {
  86             *mask.as_mut().last_mut().unwrap() &= u8::MAX >> (8 - LANES % 8);
  87         }
  88         Self(mask, PhantomData)
  89     }
  90
  91     #[inline]
  92     #[must_use = "method returns a new bool and does not mutate the original value"]
  93     pub unsafe fn test_unchecked(&self, lane: usize) -> bool {
  94         (self.0.as_ref()[lane / 8] >> (lane % 8)) & 0x1 > 0
  95     }
  96
  97     #[inline]
  98     pub unsafe fn set_unchecked(&mut self, lane: usize, value: bool) {
  99         unsafe {
 100             self.0.as_mut()[lane / 8] ^= ((value ^ self.test_unchecked(lane)) as u8) << (lane % 8)
 101         }
 102     }
 103
 104     #[inline]
 105     #[must_use = "method returns a new vector and does not mutate the original value"]
 106     pub fn to_int(self) -> Simd<T, LANES> {
 107         unsafe {
 108             crate::intrinsics::simd_select_bitmask(
 109                 self.0,
 110                 Simd::splat(T::TRUE),
 111                 Simd::splat(T::FALSE),
 112             )
 113         }
 114     }
 115
 116     #[inline]
 117     #[must_use = "method returns a new mask and does not mutate the original value"]
 118     pub unsafe fn from_int_unchecked(value: Simd<T, LANES>) -> Self {
 119         unsafe { Self(crate::intrinsics::simd_bitmask(value), PhantomData) }
 120     }
 121
 122     #[cfg(feature = "generic_const_exprs")]
 123     #[inline]
 124     #[must_use = "method returns a new array and does not mutate the original value"]
 125     pub fn to_bitmask(self) -> [u8; LaneCount::<LANES>::BITMASK_LEN] {
 126         // Safety: these are the same type and we are laundering the generic
 127         unsafe { core::mem::transmute_copy(&self.0) }
 128     }
 129
 130     #[cfg(feature = "generic_const_exprs")]
 131     #[inline]
 132     #[must_use = "method returns a new mask and does not mutate the original value"]
 133     pub fn from_bitmask(bitmask: [u8; LaneCount::<LANES>::BITMASK_LEN]) -> Self {
 134         // Safety: these are the same type and we are laundering the generic
 135         Self(unsafe { core::mem::transmute_copy(&bitmask) }, PhantomData)
 136     }
 137
 138     #[inline]
 139     #[must_use = "method returns a new mask and does not mutate the original value"]
 140     pub fn convert<U>(self) -> Mask<U, LANES>
 141     where
 142         U: MaskElement,
 143     {
 144         unsafe { core::mem::transmute_copy(&self) }
 145     }
 146
 147     #[inline]
 148     #[must_use = "method returns a new bool and does not mutate the original value"]
 149     pub fn any(self) -> bool {
 150         self != Self::splat(false)
 151     }
 152
 153     #[inline]
 154     #[must_use = "method returns a new bool and does not mutate the original value"]
 155     pub fn all(self) -> bool {
 156         self == Self::splat(true)
 157     }
 158 }
 159
 160 impl<T, const LANES: usize> core::ops::BitAnd for Mask<T, LANES>
 161 where
 162     T: MaskElement,
 163     LaneCount<LANES>: SupportedLaneCount,
 164     <LaneCount<LANES> as SupportedLaneCount>::BitMask: AsRef<[u8]> + AsMut<[u8]>,
 165 {
 166     type Output = Self;
 167     #[inline]
 168     #[must_use = "method returns a new mask and does not mutate the original value"]
 169     fn bitand(mut self, rhs: Self) -> Self {
 170         for (l, r) in self.0.as_mut().iter_mut().zip(rhs.0.as_ref().iter()) {
 171             *l &= r;
 172         }
 173         self
 174     }
 175 }
 176
 177 impl<T, const LANES: usize> core::ops::BitOr for Mask<T, LANES>
 178 where
 179     T: MaskElement,
 180     LaneCount<LANES>: SupportedLaneCount,
 181     <LaneCount<LANES> as SupportedLaneCount>::BitMask: AsRef<[u8]> + AsMut<[u8]>,
 182 {
 183     type Output = Self;
 184     #[inline]
 185     #[must_use = "method returns a new mask and does not mutate the original value"]
 186     fn bitor(mut self, rhs: Self) -> Self {
 187         for (l, r) in self.0.as_mut().iter_mut().zip(rhs.0.as_ref().iter()) {
 188             *l |= r;
 189         }
 190         self
 191     }
 192 }
 193
 194 impl<T, const LANES: usize> core::ops::BitXor for Mask<T, LANES>
 195 where
 196     T: MaskElement,
 197     LaneCount<LANES>: SupportedLaneCount,
 198 {
 199     type Output = Self;
 200     #[inline]
 201     #[must_use = "method returns a new mask and does not mutate the original value"]
 202     fn bitxor(mut self, rhs: Self) -> Self::Output {
 203         for (l, r) in self.0.as_mut().iter_mut().zip(rhs.0.as_ref().iter()) {
 204             *l ^= r;
 205         }
 206         self
 207     }
 208 }
 209
 210 impl<T, const LANES: usize> core::ops::Not for Mask<T, LANES>
 211 where
 212     T: MaskElement,
 213     LaneCount<LANES>: SupportedLaneCount,
 214 {
 215     type Output = Self;
 216     #[inline]
 217     #[must_use = "method returns a new mask and does not mutate the original value"]
 218     fn not(mut self) -> Self::Output {
 219         for x in self.0.as_mut() {
 220             *x = !*x;
 221         }
 222         if LANES % 8 > 0 {
 223             *self.0.as_mut().last_mut().unwrap() &= u8::MAX >> (8 - LANES % 8);
 224         }
 225         self
 226     }
 227 }