library/portable-simd/crates/core_simd/src/select.rs

   1 use crate::simd::intrinsics;
   2 use crate::simd::{LaneCount, Mask, MaskElement, Simd, SimdElement, SupportedLaneCount};
   3
   4 mod sealed {
   5     pub trait Sealed<Mask> {
   6         fn select(mask: Mask, true_values: Self, false_values: Self) -> Self;
   7     }
   8 }
   9 use sealed::Sealed;
  10
  11 /// Supporting trait for vector `select` function
  12 pub trait Select<Mask>: Sealed<Mask> {}
  13
  14 impl<T, const LANES: usize> Sealed<Mask<T::Mask, LANES>> for Simd<T, LANES>
  15 where
  16     T: SimdElement,
  17     LaneCount<LANES>: SupportedLaneCount,
  18 {
  19     #[inline]
  20     #[must_use = "method returns a new vector and does not mutate the original inputs"]
  21     fn select(mask: Mask<T::Mask, LANES>, true_values: Self, false_values: Self) -> Self {
  22         unsafe { intrinsics::simd_select(mask.to_int(), true_values, false_values) }
  23     }
  24 }
  25
  26 impl<T, const LANES: usize> Select<Mask<T::Mask, LANES>> for Simd<T, LANES>
  27 where
  28     T: SimdElement,
  29     LaneCount<LANES>: SupportedLaneCount,
  30 {
  31 }
  32
  33 impl<T, const LANES: usize> Sealed<Self> for Mask<T, LANES>
  34 where
  35     T: MaskElement,
  36     LaneCount<LANES>: SupportedLaneCount,
  37 {
  38     #[inline]
  39     #[must_use = "method returns a new vector and does not mutate the original inputs"]
  40     fn select(mask: Self, true_values: Self, false_values: Self) -> Self {
  41         mask & true_values | !mask & false_values
  42     }
  43 }
  44
  45 impl<T, const LANES: usize> Select<Self> for Mask<T, LANES>
  46 where
  47     T: MaskElement,
  48     LaneCount<LANES>: SupportedLaneCount,
  49 {
  50 }
  51
  52 impl<T, const LANES: usize> Mask<T, LANES>
  53 where
  54     T: MaskElement,
  55     LaneCount<LANES>: SupportedLaneCount,
  56 {
  57     /// Choose lanes from two vectors.
  58     ///
  59     /// For each lane in the mask, choose the corresponding lane from `true_values` if
  60     /// that lane mask is true, and `false_values` if that lane mask is false.
  61     ///
  62     /// ```
  63     /// # #![feature(portable_simd)]
  64     /// # #[cfg(feature = "std")] use core_simd::{Simd, Mask};
  65     /// # #[cfg(not(feature = "std"))] use core::simd::{Simd, Mask};
  66     /// let a = Simd::from_array([0, 1, 2, 3]);
  67     /// let b = Simd::from_array([4, 5, 6, 7]);
  68     /// let mask = Mask::from_array([true, false, false, true]);
  69     /// let c = mask.select(a, b);
  70     /// assert_eq!(c.to_array(), [0, 5, 6, 3]);
  71     /// ```
  72     ///
  73     /// `select` can also be used on masks:
  74     /// ```
  75     /// # #![feature(portable_simd)]
  76     /// # #[cfg(feature = "std")] use core_simd::Mask;
  77     /// # #[cfg(not(feature = "std"))] use core::simd::Mask;
  78     /// let a = Mask::<i32, 4>::from_array([true, true, false, false]);
  79     /// let b = Mask::<i32, 4>::from_array([false, false, true, true]);
  80     /// let mask = Mask::<i32, 4>::from_array([true, false, false, true]);
  81     /// let c = mask.select(a, b);
  82     /// assert_eq!(c.to_array(), [true, false, true, false]);
  83     /// ```
  84     #[inline]
  85     #[must_use = "method returns a new vector and does not mutate the original inputs"]
  86     pub fn select<S: Select<Self>>(self, true_values: S, false_values: S) -> S {
  87         S::select(self, true_values, false_values)
  88     }
  89 }