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, ToBitMask};
   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             intrinsics::simd_select_bitmask(self.0, Simd::splat(T::TRUE), Simd::splat(T::FALSE))
 109         }
 110     }
 111
 112     #[inline]
 113     #[must_use = "method returns a new mask and does not mutate the original value"]
 114     pub unsafe fn from_int_unchecked(value: Simd<T, LANES>) -> Self {
 115         unsafe { Self(intrinsics::simd_bitmask(value), PhantomData) }
 116     }
 117
 118     #[inline]
 119     pub fn to_bitmask_integer<U>(self) -> U
 120     where
 121         super::Mask<T, LANES>: ToBitMask<BitMask = U>,
 122     {
 123         // Safety: these are the same types
 124         unsafe { core::mem::transmute_copy(&self.0) }
 125     }
 126
 127     #[inline]
 128     pub fn from_bitmask_integer<U>(bitmask: U) -> Self
 129     where
 130         super::Mask<T, LANES>: ToBitMask<BitMask = U>,
 131     {
 132         // Safety: these are the same types
 133         unsafe { Self(core::mem::transmute_copy(&bitmask), PhantomData) }
 134     }
 135
 136     #[inline]
 137     #[must_use = "method returns a new mask and does not mutate the original value"]
 138     pub fn convert<U>(self) -> Mask<U, LANES>
 139     where
 140         U: MaskElement,
 141     {
 142         // Safety: bitmask layout does not depend on the element width
 143         unsafe { core::mem::transmute_copy(&self) }
 144     }
 145
 146     #[inline]
 147     #[must_use = "method returns a new bool and does not mutate the original value"]
 148     pub fn any(self) -> bool {
 149         self != Self::splat(false)
 150     }
 151
 152     #[inline]
 153     #[must_use = "method returns a new bool and does not mutate the original value"]
 154     pub fn all(self) -> bool {
 155         self == Self::splat(true)
 156     }
 157 }
 158
 159 impl<T, const LANES: usize> core::ops::BitAnd for Mask<T, LANES>
 160 where
 161     T: MaskElement,
 162     LaneCount<LANES>: SupportedLaneCount,
 163     <LaneCount<LANES> as SupportedLaneCount>::BitMask: AsRef<[u8]> + AsMut<[u8]>,
 164 {
 165     type Output = Self;
 166     #[inline]
 167     #[must_use = "method returns a new mask and does not mutate the original value"]
 168     fn bitand(mut self, rhs: Self) -> Self {
 169         for (l, r) in self.0.as_mut().iter_mut().zip(rhs.0.as_ref().iter()) {
 170             *l &= r;
 171         }
 172         self
 173     }
 174 }
 175
 176 impl<T, const LANES: usize> core::ops::BitOr for Mask<T, LANES>
 177 where
 178     T: MaskElement,
 179     LaneCount<LANES>: SupportedLaneCount,
 180     <LaneCount<LANES> as SupportedLaneCount>::BitMask: AsRef<[u8]> + AsMut<[u8]>,
 181 {
 182     type Output = Self;
 183     #[inline]
 184     #[must_use = "method returns a new mask and does not mutate the original value"]
 185     fn bitor(mut self, rhs: Self) -> Self {
 186         for (l, r) in self.0.as_mut().iter_mut().zip(rhs.0.as_ref().iter()) {
 187             *l |= r;
 188         }
 189         self
 190     }
 191 }
 192
 193 impl<T, const LANES: usize> core::ops::BitXor for Mask<T, LANES>
 194 where
 195     T: MaskElement,
 196     LaneCount<LANES>: SupportedLaneCount,
 197 {
 198     type Output = Self;
 199     #[inline]
 200     #[must_use = "method returns a new mask and does not mutate the original value"]
 201     fn bitxor(mut self, rhs: Self) -> Self::Output {
 202         for (l, r) in self.0.as_mut().iter_mut().zip(rhs.0.as_ref().iter()) {
 203             *l ^= r;
 204         }
 205         self
 206     }
 207 }
 208
 209 impl<T, const LANES: usize> core::ops::Not for Mask<T, LANES>
 210 where
 211     T: MaskElement,
 212     LaneCount<LANES>: SupportedLaneCount,
 213 {
 214     type Output = Self;
 215     #[inline]
 216     #[must_use = "method returns a new mask and does not mutate the original value"]
 217     fn not(mut self) -> Self::Output {
 218         for x in self.0.as_mut() {
 219             *x = !*x;
 220         }
 221         if LANES % 8 > 0 {
 222             *self.0.as_mut().last_mut().unwrap() &= u8::MAX >> (8 - LANES % 8);
 223         }
 224         self
 225     }
 226 }