From: Caleb Zulawski <caleb.zulawski@gmail.com>
Date: Sat, 8 May 2021 00:07:07 +0000 (+0000)
Subject: Add select for masks
X-Git-Url: https://git.lizzy.rs/?a=commitdiff_plain;ds=sidebyside;h=0bf5eb5f72f354489ce897a2bd1f6c1e9102fdd2;p=rust.git

Add select for masks
---

diff --git a/crates/core_simd/src/intrinsics.rs b/crates/core_simd/src/intrinsics.rs
index 798c4461f79..3779d96a40e 100644
--- a/crates/core_simd/src/intrinsics.rs
+++ b/crates/core_simd/src/intrinsics.rs
@@ -82,6 +82,7 @@
 
     // select
     pub(crate) fn simd_select<T, U>(m: T, a: U, b: U) -> U;
+    #[allow(unused)]
     pub(crate) fn simd_select_bitmask<T, U>(m: T, a: U, b: U) -> U;
 }
 
diff --git a/crates/core_simd/src/select.rs b/crates/core_simd/src/select.rs
index 66b0839cf24..a00af73ef4d 100644
--- a/crates/core_simd/src/select.rs
+++ b/crates/core_simd/src/select.rs
@@ -1,10 +1,13 @@
 mod sealed {
-pub trait Sealed {}
+    pub trait Sealed {}
 }
 use sealed::Sealed;
 
 /// Supporting trait for vector `select` function
-pub trait Select<Mask>: Sealed {}
+pub trait Select<Mask>: Sealed {
+    #[doc(hidden)]
+    fn select(mask: Mask, true_values: Self, false_values: Self) -> Self;
+}
 
 macro_rules! impl_select {
     {
@@ -17,9 +20,32 @@ impl<const LANES: usize> Select<crate::$mask<LANES>> for crate::$type<LANES>
             crate::$mask<LANES>: crate::Mask,
             crate::$bits_ty<LANES>: crate::LanesAtMost32,
             Self: crate::LanesAtMost32,
-        {}
+        {
+            #[doc(hidden)]
+            #[inline]
+            fn select(mask: crate::$mask<LANES>, true_values: Self, false_values: Self) -> Self {
+                unsafe { crate::intrinsics::simd_select(mask.to_int(), true_values, false_values) }
+            }
+        }
         )*
 
+        impl<const LANES: usize> Sealed for crate::$mask<LANES>
+        where
+            Self: crate::Mask,
+            crate::$bits_ty<LANES>: crate::LanesAtMost32,
+        {}
+        impl<const LANES: usize> Select<Self> for crate::$mask<LANES>
+        where
+            Self: crate::Mask,
+            crate::$bits_ty<LANES>: crate::LanesAtMost32,
+        {
+            #[doc(hidden)]
+            #[inline]
+            fn select(mask: Self, true_values: Self, false_values: Self) -> Self {
+                mask & true_values | !mask & false_values
+            }
+        }
+
         impl<const LANES: usize> crate::$mask<LANES>
         where
             Self: crate::Mask,
@@ -38,8 +64,19 @@ impl<const LANES: usize> crate::$mask<LANES>
             /// let c = mask.select(a, b);
             /// assert_eq!(c.to_array(), [0, 5, 6, 3]);
             /// ```
+            ///
+            /// `select` can also be used with masks:
+            /// ```
+            /// # use core_simd::{Mask32};
+            /// let a = Mask32::from_array([true, true, false, false]);
+            /// let b = Mask32::from_array([false, false, true, true]);
+            /// let mask = Mask32::from_array([true, false, false, true]);
+            /// let c = mask.select(a, b);
+            /// assert_eq!(c.to_array(), [true, false, true, false]);
+            /// ```
+            #[inline]
             pub fn select<S: Select<Self>>(self, true_values: S, false_values: S) -> S {
-                unsafe { crate::intrinsics::simd_select(self.to_int(), true_values, false_values) }
+                S::select(self, true_values, false_values)
             }
         }
     }
diff --git a/crates/core_simd/tests/mask_ops_impl/mod.rs b/crates/core_simd/tests/mask_ops_impl/mod.rs
index ff36af95651..b414167866e 100644
--- a/crates/core_simd/tests/mask_ops_impl/mod.rs
+++ b/crates/core_simd/tests/mask_ops_impl/mod.rs
@@ -1,8 +1,8 @@
 #[macro_use]
 mod mask_macros;
 
-mod mask8;
 mod mask16;
 mod mask32;
 mod mask64;
+mod mask8;
 mod masksize;