1#![allow(non_camel_case_types)]
4
5#[cfg_attr(
6 not(all(target_arch = "x86_64", target_feature = "avx512f")),
7 path = "masks/full_masks.rs"
8)]
9#[cfg_attr(
10 all(target_arch = "x86_64", target_feature = "avx512f"),
11 path = "masks/bitmask.rs"
12)]
13mod mask_impl;
14
15use crate::simd::{LaneCount, Simd, SimdCast, SimdElement, SupportedLaneCount};
16use core::cmp::Ordering;
17use core::{fmt, mem};
18
19mod sealed {
20 use super::*;
21
22 pub trait Sealed {
29 fn valid<const N: usize>(values: Simd<Self, N>) -> bool
30 where
31 LaneCount<N>: SupportedLaneCount,
32 Self: SimdElement;
33
34 fn eq(self, other: Self) -> bool;
35
36 fn to_usize(self) -> usize;
37 fn max_unsigned() -> u64;
38
39 type Unsigned: SimdElement;
40
41 const TRUE: Self;
42
43 const FALSE: Self;
44 }
45}
46use sealed::Sealed;
47
48pub unsafe trait MaskElement: SimdElement<Mask = Self> + SimdCast + Sealed {}
53
54macro_rules! impl_element {
55 { $ty:ty, $unsigned:ty } => {
56 impl Sealed for $ty {
57 #[inline]
58 fn valid<const N: usize>(value: Simd<Self, N>) -> bool
59 where
60 LaneCount<N>: SupportedLaneCount,
61 {
62 unsafe {
66 use core::intrinsics::simd;
67 let falses: Simd<Self, N> = simd::simd_eq(value, Simd::splat(0 as _));
68 let trues: Simd<Self, N> = simd::simd_eq(value, Simd::splat(-1 as _));
69 let valid: Simd<Self, N> = simd::simd_or(falses, trues);
70 simd::simd_reduce_all(valid)
71 }
72 }
73
74 #[inline]
75 fn eq(self, other: Self) -> bool { self == other }
76
77 #[inline]
78 fn to_usize(self) -> usize {
79 self as usize
80 }
81
82 #[inline]
83 fn max_unsigned() -> u64 {
84 <$unsigned>::MAX as u64
85 }
86
87 type Unsigned = $unsigned;
88
89 const TRUE: Self = -1;
90 const FALSE: Self = 0;
91 }
92
93 unsafe impl MaskElement for $ty {}
95 }
96}
97
98impl_element! { i8, u8 }
99impl_element! { i16, u16 }
100impl_element! { i32, u32 }
101impl_element! { i64, u64 }
102impl_element! { isize, usize }
103
104#[repr(transparent)]
112pub struct Mask<T, const N: usize>(mask_impl::Mask<T, N>)
113where
114 T: MaskElement,
115 LaneCount<N>: SupportedLaneCount;
116
117impl<T, const N: usize> Copy for Mask<T, N>
118where
119 T: MaskElement,
120 LaneCount<N>: SupportedLaneCount,
121{
122}
123
124impl<T, const N: usize> Clone for Mask<T, N>
125where
126 T: MaskElement,
127 LaneCount<N>: SupportedLaneCount,
128{
129 #[inline]
130 fn clone(&self) -> Self {
131 *self
132 }
133}
134
135impl<T, const N: usize> Mask<T, N>
136where
137 T: MaskElement,
138 LaneCount<N>: SupportedLaneCount,
139{
140 #[inline]
142 pub fn splat(value: bool) -> Self {
143 Self(mask_impl::Mask::splat(value))
144 }
145
146 #[inline]
148 pub fn from_array(array: [bool; N]) -> Self {
149 unsafe {
156 let bytes: [u8; N] = mem::transmute_copy(&array);
157 let bools: Simd<i8, N> =
158 core::intrinsics::simd::simd_ne(Simd::from_array(bytes), Simd::splat(0u8));
159 Mask::from_int_unchecked(core::intrinsics::simd::simd_cast(bools))
160 }
161 }
162
163 #[inline]
165 pub fn to_array(self) -> [bool; N] {
166 unsafe {
177 let mut bytes: Simd<i8, N> = core::intrinsics::simd::simd_cast(self.to_int());
178 bytes &= Simd::splat(1i8);
179 mem::transmute_copy(&bytes)
180 }
181 }
182
183 #[inline]
189 #[must_use = "method returns a new mask and does not mutate the original value"]
190 pub unsafe fn from_int_unchecked(value: Simd<T, N>) -> Self {
191 unsafe {
193 core::intrinsics::assume(<T as Sealed>::valid(value));
194 Self(mask_impl::Mask::from_int_unchecked(value))
195 }
196 }
197
198 #[inline]
204 #[must_use = "method returns a new mask and does not mutate the original value"]
205 #[track_caller]
206 pub fn from_int(value: Simd<T, N>) -> Self {
207 assert!(T::valid(value), "all values must be either 0 or -1",);
208 unsafe { Self::from_int_unchecked(value) }
210 }
211
212 #[inline]
215 #[must_use = "method returns a new vector and does not mutate the original value"]
216 pub fn to_int(self) -> Simd<T, N> {
217 self.0.to_int()
218 }
219
220 #[inline]
222 #[must_use = "method returns a new mask and does not mutate the original value"]
223 pub fn cast<U: MaskElement>(self) -> Mask<U, N> {
224 Mask(self.0.convert())
225 }
226
227 #[inline]
232 #[must_use = "method returns a new bool and does not mutate the original value"]
233 pub unsafe fn test_unchecked(&self, index: usize) -> bool {
234 unsafe { self.0.test_unchecked(index) }
236 }
237
238 #[inline]
243 #[must_use = "method returns a new bool and does not mutate the original value"]
244 #[track_caller]
245 pub fn test(&self, index: usize) -> bool {
246 assert!(index < N, "element index out of range");
247 unsafe { self.test_unchecked(index) }
249 }
250
251 #[inline]
256 pub unsafe fn set_unchecked(&mut self, index: usize, value: bool) {
257 unsafe {
259 self.0.set_unchecked(index, value);
260 }
261 }
262
263 #[inline]
268 #[track_caller]
269 pub fn set(&mut self, index: usize, value: bool) {
270 assert!(index < N, "element index out of range");
271 unsafe {
273 self.set_unchecked(index, value);
274 }
275 }
276
277 #[inline]
279 #[must_use = "method returns a new bool and does not mutate the original value"]
280 pub fn any(self) -> bool {
281 self.0.any()
282 }
283
284 #[inline]
286 #[must_use = "method returns a new bool and does not mutate the original value"]
287 pub fn all(self) -> bool {
288 self.0.all()
289 }
290
291 #[inline]
296 #[must_use = "method returns a new integer and does not mutate the original value"]
297 pub fn to_bitmask(self) -> u64 {
298 self.0.to_bitmask_integer()
299 }
300
301 #[inline]
306 #[must_use = "method returns a new mask and does not mutate the original value"]
307 pub fn from_bitmask(bitmask: u64) -> Self {
308 Self(mask_impl::Mask::from_bitmask_integer(bitmask))
309 }
310
311 #[inline]
325 #[must_use = "method returns the index and does not mutate the original value"]
326 pub fn first_set(self) -> Option<usize> {
327 if cfg!(target_feature = "sse") && N <= 64 {
329 let tz = self.to_bitmask().trailing_zeros();
330 return if tz == 64 { None } else { Some(tz as usize) };
331 }
332
333 let index = Simd::from_array(
340 const {
341 let mut index = [0; N];
342 let mut i = 0;
343 while i < N {
344 index[i] = i;
345 i += 1;
346 }
347 index
348 },
349 );
350
351 let index: Simd<T, N> = unsafe { core::intrinsics::simd::simd_cast(index) };
353
354 let masked_index = self.select(index, Self::splat(true).to_int());
355
356 let masked_index: Simd<T::Unsigned, N> =
358 unsafe { core::intrinsics::simd::simd_cast(masked_index) };
359
360 let min_index: T::Unsigned =
362 unsafe { core::intrinsics::simd::simd_reduce_min(masked_index) };
363
364 let min_index: T = unsafe { core::mem::transmute_copy(&min_index) };
366
367 if min_index.eq(T::TRUE) {
368 None
369 } else {
370 Some(min_index.to_usize())
371 }
372 }
373}
374
375impl<T, const N: usize> From<[bool; N]> for Mask<T, N>
377where
378 T: MaskElement,
379 LaneCount<N>: SupportedLaneCount,
380{
381 #[inline]
382 fn from(array: [bool; N]) -> Self {
383 Self::from_array(array)
384 }
385}
386
387impl<T, const N: usize> From<Mask<T, N>> for [bool; N]
388where
389 T: MaskElement,
390 LaneCount<N>: SupportedLaneCount,
391{
392 #[inline]
393 fn from(vector: Mask<T, N>) -> Self {
394 vector.to_array()
395 }
396}
397
398impl<T, const N: usize> Default for Mask<T, N>
399where
400 T: MaskElement,
401 LaneCount<N>: SupportedLaneCount,
402{
403 #[inline]
404 #[must_use = "method returns a defaulted mask with all elements set to false (0)"]
405 fn default() -> Self {
406 Self::splat(false)
407 }
408}
409
410impl<T, const N: usize> PartialEq for Mask<T, N>
411where
412 T: MaskElement + PartialEq,
413 LaneCount<N>: SupportedLaneCount,
414{
415 #[inline]
416 #[must_use = "method returns a new bool and does not mutate the original value"]
417 fn eq(&self, other: &Self) -> bool {
418 self.0 == other.0
419 }
420}
421
422impl<T, const N: usize> PartialOrd for Mask<T, N>
423where
424 T: MaskElement + PartialOrd,
425 LaneCount<N>: SupportedLaneCount,
426{
427 #[inline]
428 #[must_use = "method returns a new Ordering and does not mutate the original value"]
429 fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
430 self.0.partial_cmp(&other.0)
431 }
432}
433
434impl<T, const N: usize> fmt::Debug for Mask<T, N>
435where
436 T: MaskElement + fmt::Debug,
437 LaneCount<N>: SupportedLaneCount,
438{
439 #[inline]
440 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
441 f.debug_list()
442 .entries((0..N).map(|i| self.test(i)))
443 .finish()
444 }
445}
446
447impl<T, const N: usize> core::ops::BitAnd for Mask<T, N>
448where
449 T: MaskElement,
450 LaneCount<N>: SupportedLaneCount,
451{
452 type Output = Self;
453 #[inline]
454 #[must_use = "method returns a new mask and does not mutate the original value"]
455 fn bitand(self, rhs: Self) -> Self {
456 Self(self.0 & rhs.0)
457 }
458}
459
460impl<T, const N: usize> core::ops::BitAnd<bool> for Mask<T, N>
461where
462 T: MaskElement,
463 LaneCount<N>: SupportedLaneCount,
464{
465 type Output = Self;
466 #[inline]
467 #[must_use = "method returns a new mask and does not mutate the original value"]
468 fn bitand(self, rhs: bool) -> Self {
469 self & Self::splat(rhs)
470 }
471}
472
473impl<T, const N: usize> core::ops::BitAnd<Mask<T, N>> for bool
474where
475 T: MaskElement,
476 LaneCount<N>: SupportedLaneCount,
477{
478 type Output = Mask<T, N>;
479 #[inline]
480 #[must_use = "method returns a new mask and does not mutate the original value"]
481 fn bitand(self, rhs: Mask<T, N>) -> Mask<T, N> {
482 Mask::splat(self) & rhs
483 }
484}
485
486impl<T, const N: usize> core::ops::BitOr for Mask<T, N>
487where
488 T: MaskElement,
489 LaneCount<N>: SupportedLaneCount,
490{
491 type Output = Self;
492 #[inline]
493 #[must_use = "method returns a new mask and does not mutate the original value"]
494 fn bitor(self, rhs: Self) -> Self {
495 Self(self.0 | rhs.0)
496 }
497}
498
499impl<T, const N: usize> core::ops::BitOr<bool> for Mask<T, N>
500where
501 T: MaskElement,
502 LaneCount<N>: SupportedLaneCount,
503{
504 type Output = Self;
505 #[inline]
506 #[must_use = "method returns a new mask and does not mutate the original value"]
507 fn bitor(self, rhs: bool) -> Self {
508 self | Self::splat(rhs)
509 }
510}
511
512impl<T, const N: usize> core::ops::BitOr<Mask<T, N>> for bool
513where
514 T: MaskElement,
515 LaneCount<N>: SupportedLaneCount,
516{
517 type Output = Mask<T, N>;
518 #[inline]
519 #[must_use = "method returns a new mask and does not mutate the original value"]
520 fn bitor(self, rhs: Mask<T, N>) -> Mask<T, N> {
521 Mask::splat(self) | rhs
522 }
523}
524
525impl<T, const N: usize> core::ops::BitXor for Mask<T, N>
526where
527 T: MaskElement,
528 LaneCount<N>: SupportedLaneCount,
529{
530 type Output = Self;
531 #[inline]
532 #[must_use = "method returns a new mask and does not mutate the original value"]
533 fn bitxor(self, rhs: Self) -> Self::Output {
534 Self(self.0 ^ rhs.0)
535 }
536}
537
538impl<T, const N: usize> core::ops::BitXor<bool> for Mask<T, N>
539where
540 T: MaskElement,
541 LaneCount<N>: SupportedLaneCount,
542{
543 type Output = Self;
544 #[inline]
545 #[must_use = "method returns a new mask and does not mutate the original value"]
546 fn bitxor(self, rhs: bool) -> Self::Output {
547 self ^ Self::splat(rhs)
548 }
549}
550
551impl<T, const N: usize> core::ops::BitXor<Mask<T, N>> for bool
552where
553 T: MaskElement,
554 LaneCount<N>: SupportedLaneCount,
555{
556 type Output = Mask<T, N>;
557 #[inline]
558 #[must_use = "method returns a new mask and does not mutate the original value"]
559 fn bitxor(self, rhs: Mask<T, N>) -> Self::Output {
560 Mask::splat(self) ^ rhs
561 }
562}
563
564impl<T, const N: usize> core::ops::Not for Mask<T, N>
565where
566 T: MaskElement,
567 LaneCount<N>: SupportedLaneCount,
568{
569 type Output = Mask<T, N>;
570 #[inline]
571 #[must_use = "method returns a new mask and does not mutate the original value"]
572 fn not(self) -> Self::Output {
573 Self(!self.0)
574 }
575}
576
577impl<T, const N: usize> core::ops::BitAndAssign for Mask<T, N>
578where
579 T: MaskElement,
580 LaneCount<N>: SupportedLaneCount,
581{
582 #[inline]
583 fn bitand_assign(&mut self, rhs: Self) {
584 self.0 = self.0 & rhs.0;
585 }
586}
587
588impl<T, const N: usize> core::ops::BitAndAssign<bool> for Mask<T, N>
589where
590 T: MaskElement,
591 LaneCount<N>: SupportedLaneCount,
592{
593 #[inline]
594 fn bitand_assign(&mut self, rhs: bool) {
595 *self &= Self::splat(rhs);
596 }
597}
598
599impl<T, const N: usize> core::ops::BitOrAssign for Mask<T, N>
600where
601 T: MaskElement,
602 LaneCount<N>: SupportedLaneCount,
603{
604 #[inline]
605 fn bitor_assign(&mut self, rhs: Self) {
606 self.0 = self.0 | rhs.0;
607 }
608}
609
610impl<T, const N: usize> core::ops::BitOrAssign<bool> for Mask<T, N>
611where
612 T: MaskElement,
613 LaneCount<N>: SupportedLaneCount,
614{
615 #[inline]
616 fn bitor_assign(&mut self, rhs: bool) {
617 *self |= Self::splat(rhs);
618 }
619}
620
621impl<T, const N: usize> core::ops::BitXorAssign for Mask<T, N>
622where
623 T: MaskElement,
624 LaneCount<N>: SupportedLaneCount,
625{
626 #[inline]
627 fn bitxor_assign(&mut self, rhs: Self) {
628 self.0 = self.0 ^ rhs.0;
629 }
630}
631
632impl<T, const N: usize> core::ops::BitXorAssign<bool> for Mask<T, N>
633where
634 T: MaskElement,
635 LaneCount<N>: SupportedLaneCount,
636{
637 #[inline]
638 fn bitxor_assign(&mut self, rhs: bool) {
639 *self ^= Self::splat(rhs);
640 }
641}
642
643macro_rules! impl_from {
644 { $from:ty => $($to:ty),* } => {
645 $(
646 impl<const N: usize> From<Mask<$from, N>> for Mask<$to, N>
647 where
648 LaneCount<N>: SupportedLaneCount,
649 {
650 #[inline]
651 fn from(value: Mask<$from, N>) -> Self {
652 value.cast()
653 }
654 }
655 )*
656 }
657}
658impl_from! { i8 => i16, i32, i64, isize }
659impl_from! { i16 => i32, i64, isize, i8 }
660impl_from! { i32 => i64, isize, i8, i16 }
661impl_from! { i64 => isize, i8, i16, i32 }
662impl_from! { isize => i8, i16, i32, i64 }