crates/core_simd/src/masks/bitmask.rs

   1 use crate::LanesAtMost64;
   2
   3 /// A mask where each lane is represented by a single bit.
   4 #[derive(Copy, Clone, Debug)]
   5 #[repr(transparent)]
   6 pub struct BitMask<const LANES: usize>(u64)
   7 where
   8     BitMask<LANES>: LanesAtMost64;
   9
  10 impl<const LANES: usize> BitMask<LANES>
  11 where
  12     Self: LanesAtMost64,
  13 {
  14     /// Construct a mask by setting all lanes to the given value.
  15     pub fn splat(value: bool) -> Self {
  16         if value {
  17             Self(u64::MAX)
  18         } else {
  19             Self(u64::MIN)
  20         }
  21     }
  22
  23     /// Tests the value of the specified lane.
  24     ///
  25     /// # Panics
  26     /// Panics if `lane` is greater than or equal to the number of lanes in the vector.
  27     #[inline]
  28     pub fn test(&self, lane: usize) -> bool {
  29         assert!(lane < LANES, "lane index out of range");
  30         (self.0 >> lane) & 0x1 > 0
  31     }
  32
  33     /// Sets the value of the specified lane.
  34     ///
  35     /// # Panics
  36     /// Panics if `lane` is greater than or equal to the number of lanes in the vector.
  37     #[inline]
  38     pub fn set(&mut self, lane: usize, value: bool) {
  39         assert!(lane < LANES, "lane index out of range");
  40         self.0 ^= ((value ^ self.test(lane)) as u64) << lane
  41     }
  42 }
  43
  44 impl<const LANES: usize> core::ops::BitAnd for BitMask<LANES>
  45 where
  46     Self: LanesAtMost64,
  47 {
  48     type Output = Self;
  49     #[inline]
  50     fn bitand(self, rhs: Self) -> Self {
  51         Self(self.0 & rhs.0)
  52     }
  53 }
  54
  55 impl<const LANES: usize> core::ops::BitAnd<bool> for BitMask<LANES>
  56 where
  57     Self: LanesAtMost64,
  58 {
  59     type Output = Self;
  60     #[inline]
  61     fn bitand(self, rhs: bool) -> Self {
  62         self & Self::splat(rhs)
  63     }
  64 }
  65
  66 impl<const LANES: usize> core::ops::BitAnd<BitMask<LANES>> for bool
  67 where
  68     BitMask<LANES>: LanesAtMost64,
  69 {
  70     type Output = BitMask<LANES>;
  71     #[inline]
  72     fn bitand(self, rhs: BitMask<LANES>) -> BitMask<LANES> {
  73         BitMask::<LANES>::splat(self) & rhs
  74     }
  75 }
  76
  77 impl<const LANES: usize> core::ops::BitOr for BitMask<LANES>
  78 where
  79     Self: LanesAtMost64,
  80 {
  81     type Output = Self;
  82     #[inline]
  83     fn bitor(self, rhs: Self) -> Self {
  84         Self(self.0 | rhs.0)
  85     }
  86 }
  87
  88 impl<const LANES: usize> core::ops::BitOr<bool> for BitMask<LANES>
  89 where
  90     Self: LanesAtMost64,
  91 {
  92     type Output = Self;
  93     #[inline]
  94     fn bitor(self, rhs: bool) -> Self {
  95         self | Self::splat(rhs)
  96     }
  97 }
  98
  99 impl<const LANES: usize> core::ops::BitOr<BitMask<LANES>> for bool
 100 where
 101     BitMask<LANES>: LanesAtMost64,
 102 {
 103     type Output = BitMask<LANES>;
 104     #[inline]
 105     fn bitor(self, rhs: BitMask<LANES>) -> BitMask<LANES> {
 106         BitMask::<LANES>::splat(self) | rhs
 107     }
 108 }
 109
 110 impl<const LANES: usize> core::ops::BitXor for BitMask<LANES>
 111 where
 112     Self: LanesAtMost64,
 113 {
 114     type Output = Self;
 115     #[inline]
 116     fn bitxor(self, rhs: Self) -> Self::Output {
 117         Self(self.0 ^ rhs.0)
 118     }
 119 }
 120
 121 impl<const LANES: usize> core::ops::BitXor<bool> for BitMask<LANES>
 122 where
 123     Self: LanesAtMost64,
 124 {
 125     type Output = Self;
 126     #[inline]
 127     fn bitxor(self, rhs: bool) -> Self::Output {
 128         self ^ Self::splat(rhs)
 129     }
 130 }
 131
 132 impl<const LANES: usize> core::ops::BitXor<BitMask<LANES>> for bool
 133 where
 134     BitMask<LANES>: LanesAtMost64,
 135 {
 136     type Output = BitMask<LANES>;
 137     #[inline]
 138     fn bitxor(self, rhs: BitMask<LANES>) -> Self::Output {
 139         BitMask::<LANES>::splat(self) ^ rhs
 140     }
 141 }
 142
 143 impl<const LANES: usize> core::ops::Not for BitMask<LANES>
 144 where
 145     Self: LanesAtMost64,
 146 {
 147     type Output = BitMask<LANES>;
 148     #[inline]
 149     fn not(self) -> Self::Output {
 150         Self(!self.0)
 151     }
 152 }
 153
 154 impl<const LANES: usize> core::ops::BitAndAssign for BitMask<LANES>
 155 where
 156     Self: LanesAtMost64,
 157 {
 158     #[inline]
 159     fn bitand_assign(&mut self, rhs: Self) {
 160         self.0 &= rhs.0;
 161     }
 162 }
 163
 164 impl<const LANES: usize> core::ops::BitAndAssign<bool> for BitMask<LANES>
 165 where
 166     Self: LanesAtMost64,
 167 {
 168     #[inline]
 169     fn bitand_assign(&mut self, rhs: bool) {
 170         *self &= Self::splat(rhs);
 171     }
 172 }
 173
 174 impl<const LANES: usize> core::ops::BitOrAssign for BitMask<LANES>
 175 where
 176     Self: LanesAtMost64,
 177 {
 178     #[inline]
 179     fn bitor_assign(&mut self, rhs: Self) {
 180         self.0 |= rhs.0;
 181     }
 182 }
 183
 184 impl<const LANES: usize> core::ops::BitOrAssign<bool> for BitMask<LANES>
 185 where
 186     Self: LanesAtMost64,
 187 {
 188     #[inline]
 189     fn bitor_assign(&mut self, rhs: bool) {
 190         *self |= Self::splat(rhs);
 191     }
 192 }
 193
 194 impl<const LANES: usize> core::ops::BitXorAssign for BitMask<LANES>
 195 where
 196     Self: LanesAtMost64,
 197 {
 198     #[inline]
 199     fn bitxor_assign(&mut self, rhs: Self) {
 200         self.0 ^= rhs.0;
 201     }
 202 }
 203
 204 impl<const LANES: usize> core::ops::BitXorAssign<bool> for BitMask<LANES>
 205 where
 206     Self: LanesAtMost64,
 207 {
 208     #[inline]
 209     fn bitxor_assign(&mut self, rhs: bool) {
 210         *self ^= Self::splat(rhs);
 211     }
 212 }