crates/core_simd/src/select.rs

   1 mod sealed {
   2     pub trait Sealed {}
   3 }
   4 use sealed::Sealed;
   5
   6 /// Supporting trait for vector `select` function
   7 pub trait Select<Mask>: Sealed {
   8     #[doc(hidden)]
   9     fn select(mask: Mask, true_values: Self, false_values: Self) -> Self;
  10 }
  11
  12 macro_rules! impl_select {
  13     {
  14         $mask:ident ($bits_ty:ident): $($type:ident),*
  15     } => {
  16         $(
  17         impl<const LANES: usize> Sealed for crate::$type<LANES> where Self: crate::LanesAtMost32 {}
  18         impl<const LANES: usize> Select<crate::$mask<LANES>> for crate::$type<LANES>
  19         where
  20             crate::$mask<LANES>: crate::Mask,
  21             crate::$bits_ty<LANES>: crate::LanesAtMost32,
  22             Self: crate::LanesAtMost32,
  23         {
  24             #[doc(hidden)]
  25             #[inline]
  26             fn select(mask: crate::$mask<LANES>, true_values: Self, false_values: Self) -> Self {
  27                 unsafe { crate::intrinsics::simd_select(mask.to_int(), true_values, false_values) }
  28             }
  29         }
  30         )*
  31
  32         impl<const LANES: usize> Sealed for crate::$mask<LANES>
  33         where
  34             Self: crate::Mask,
  35             crate::$bits_ty<LANES>: crate::LanesAtMost32,
  36         {}
  37         impl<const LANES: usize> Select<Self> for crate::$mask<LANES>
  38         where
  39             Self: crate::Mask,
  40             crate::$bits_ty<LANES>: crate::LanesAtMost32,
  41         {
  42             #[doc(hidden)]
  43             #[inline]
  44             fn select(mask: Self, true_values: Self, false_values: Self) -> Self {
  45                 mask & true_values | !mask & false_values
  46             }
  47         }
  48
  49         impl<const LANES: usize> crate::$mask<LANES>
  50         where
  51             Self: crate::Mask,
  52             crate::$bits_ty<LANES>: crate::LanesAtMost32,
  53         {
  54             /// Choose lanes from two vectors.
  55             ///
  56             /// For each lane in the mask, choose the corresponding lane from `true_values` if
  57             /// that lane mask is true, and `false_values` if that lane mask is false.
  58             ///
  59             /// ```
  60             /// # use core_simd::{Mask32, SimdI32};
  61             /// let a = SimdI32::from_array([0, 1, 2, 3]);
  62             /// let b = SimdI32::from_array([4, 5, 6, 7]);
  63             /// let mask = Mask32::from_array([true, false, false, true]);
  64             /// let c = mask.select(a, b);
  65             /// assert_eq!(c.to_array(), [0, 5, 6, 3]);
  66             /// ```
  67             ///
  68             /// `select` can also be used with masks:
  69             /// ```
  70             /// # use core_simd::{Mask32};
  71             /// let a = Mask32::from_array([true, true, false, false]);
  72             /// let b = Mask32::from_array([false, false, true, true]);
  73             /// let mask = Mask32::from_array([true, false, false, true]);
  74             /// let c = mask.select(a, b);
  75             /// assert_eq!(c.to_array(), [true, false, true, false]);
  76             /// ```
  77             #[inline]
  78             pub fn select<S: Select<Self>>(self, true_values: S, false_values: S) -> S {
  79                 S::select(self, true_values, false_values)
  80             }
  81         }
  82     }
  83 }
  84
  85 impl_select! { Mask8 (SimdI8): SimdU8, SimdI8 }
  86 impl_select! { Mask16 (SimdI16): SimdU16, SimdI16 }
  87 impl_select! { Mask32 (SimdI32): SimdU32, SimdI32, SimdF32}
  88 impl_select! { Mask64 (SimdI64): SimdU64, SimdI64, SimdF64}
  89 impl_select! { MaskSize (SimdIsize): SimdUsize, SimdIsize }