From 563d2a2cfce14b5807bb6cacce336e0e87950289 Mon Sep 17 00:00:00 2001
From: Caleb Zulawski <caleb.zulawski@gmail.com>
Date: Thu, 22 Apr 2021 04:09:56 +0000
Subject: [PATCH] Add select function

---
 crates/core_simd/src/intrinsics.rs |  1 +
 crates/core_simd/src/lib.rs        |  3 ++
 crates/core_simd/src/select.rs     | 52 ++++++++++++++++++++++++++++++
 3 files changed, 56 insertions(+)
 create mode 100644 crates/core_simd/src/select.rs

diff --git a/crates/core_simd/src/intrinsics.rs b/crates/core_simd/src/intrinsics.rs
index 8cbb0cbccf7..798c4461f79 100644
--- a/crates/core_simd/src/intrinsics.rs
+++ b/crates/core_simd/src/intrinsics.rs
@@ -81,6 +81,7 @@
     pub(crate) fn simd_bitmask<T, U>(x: T) -> U;
 
     // select
+    pub(crate) fn simd_select<T, U>(m: T, a: U, b: U) -> U;
     pub(crate) fn simd_select_bitmask<T, U>(m: T, a: U, b: U) -> U;
 }
 
diff --git a/crates/core_simd/src/lib.rs b/crates/core_simd/src/lib.rs
index 2d4176ce342..3489dfb5f3e 100644
--- a/crates/core_simd/src/lib.rs
+++ b/crates/core_simd/src/lib.rs
@@ -14,6 +14,9 @@
 #[macro_use]
 mod reduction;
 
+mod select;
+pub use select::Select;
+
 mod comparisons;
 mod fmt;
 mod intrinsics;
diff --git a/crates/core_simd/src/select.rs b/crates/core_simd/src/select.rs
new file mode 100644
index 00000000000..66b0839cf24
--- /dev/null
+++ b/crates/core_simd/src/select.rs
@@ -0,0 +1,52 @@
+mod sealed {
+pub trait Sealed {}
+}
+use sealed::Sealed;
+
+/// Supporting trait for vector `select` function
+pub trait Select<Mask>: Sealed {}
+
+macro_rules! impl_select {
+    {
+        $mask:ident ($bits_ty:ident): $($type:ident),*
+    } => {
+        $(
+        impl<const LANES: usize> Sealed for crate::$type<LANES> where Self: crate::LanesAtMost32 {}
+        impl<const LANES: usize> Select<crate::$mask<LANES>> for crate::$type<LANES>
+        where
+            crate::$mask<LANES>: crate::Mask,
+            crate::$bits_ty<LANES>: crate::LanesAtMost32,
+            Self: crate::LanesAtMost32,
+        {}
+        )*
+
+        impl<const LANES: usize> crate::$mask<LANES>
+        where
+            Self: crate::Mask,
+            crate::$bits_ty<LANES>: crate::LanesAtMost32,
+        {
+            /// Choose lanes from two vectors.
+            ///
+            /// For each lane in the mask, choose the corresponding lane from `true_values` if
+            /// that lane mask is true, and `false_values` if that lane mask is false.
+            ///
+            /// ```
+            /// # use core_simd::{Mask32, SimdI32};
+            /// let a = SimdI32::from_array([0, 1, 2, 3]);
+            /// let b = SimdI32::from_array([4, 5, 6, 7]);
+            /// let mask = Mask32::from_array([true, false, false, true]);
+            /// let c = mask.select(a, b);
+            /// assert_eq!(c.to_array(), [0, 5, 6, 3]);
+            /// ```
+            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) }
+            }
+        }
+    }
+}
+
+impl_select! { Mask8 (SimdI8): SimdU8, SimdI8 }
+impl_select! { Mask16 (SimdI16): SimdU16, SimdI16 }
+impl_select! { Mask32 (SimdI32): SimdU32, SimdI32, SimdF32}
+impl_select! { Mask64 (SimdI64): SimdU64, SimdI64, SimdF64}
+impl_select! { MaskSize (SimdIsize): SimdUsize, SimdIsize }