From 0a1e7453204383a109f4854cc588fd3b35c990c6 Mon Sep 17 00:00:00 2001
From: Caleb Zulawski <caleb.zulawski@gmail.com>
Date: Sun, 10 Sep 2023 18:20:47 -0400
Subject: [PATCH] Return unsigned integers from some signed integer functions

---
 crates/core_simd/src/elements/int.rs | 37 +++++++++++++++-------------
 crates/core_simd/tests/ops_macros.rs |  8 +++---
 2 files changed, 24 insertions(+), 21 deletions(-)

diff --git a/crates/core_simd/src/elements/int.rs b/crates/core_simd/src/elements/int.rs
index e81ed7bf601..c341c59545c 100644
--- a/crates/core_simd/src/elements/int.rs
+++ b/crates/core_simd/src/elements/int.rs
@@ -1,6 +1,7 @@
 use super::sealed::Sealed;
 use crate::simd::{
-    intrinsics, LaneCount, Mask, Simd, SimdCast, SimdElement, SimdPartialOrd, SupportedLaneCount,
+    intrinsics, LaneCount, Mask, Simd, SimdCast, SimdElement, SimdPartialOrd, SimdUint,
+    SupportedLaneCount,
 };
 
 /// Operations on SIMD vectors of signed integers.
@@ -11,6 +12,9 @@ pub trait SimdInt: Copy + Sealed {
     /// Scalar type contained by this SIMD vector type.
     type Scalar;
 
+    /// A SIMD vector of unsigned integers with the same element size.
+    type Unsigned;
+
     /// A SIMD vector with a different element type.
     type Cast<T: SimdElement>;
 
@@ -200,20 +204,20 @@ pub trait SimdInt: Copy + Sealed {
     fn reverse_bits(self) -> Self;
 
     /// Returns the number of leading zeros in the binary representation of each element.
-    fn leading_zeros(self) -> Self;
+    fn leading_zeros(self) -> Self::Unsigned;
 
     /// Returns the number of trailing zeros in the binary representation of each element.
-    fn trailing_zeros(self) -> Self;
+    fn trailing_zeros(self) -> Self::Unsigned;
 
     /// Returns the number of leading ones in the binary representation of each element.
-    fn leading_ones(self) -> Self;
+    fn leading_ones(self) -> Self::Unsigned;
 
     /// Returns the number of trailing ones in the binary representation of each element.
-    fn trailing_ones(self) -> Self;
+    fn trailing_ones(self) -> Self::Unsigned;
 }
 
 macro_rules! impl_trait {
-    { $($ty:ty),* } => {
+    { $($ty:ident ($unsigned:ident)),* } => {
         $(
         impl<const LANES: usize> Sealed for Simd<$ty, LANES>
         where
@@ -227,6 +231,7 @@ impl<const LANES: usize> SimdInt for Simd<$ty, LANES>
         {
             type Mask = Mask<<$ty as SimdElement>::Mask, LANES>;
             type Scalar = $ty;
+            type Unsigned = Simd<$unsigned, LANES>;
             type Cast<T: SimdElement> = Simd<T, LANES>;
 
             #[inline]
@@ -340,29 +345,27 @@ fn reverse_bits(self) -> Self {
             }
 
             #[inline]
-            fn leading_zeros(self) -> Self {
-                // Safety: `self` is an integer vector
-                unsafe { intrinsics::simd_ctlz(self) }
+            fn leading_zeros(self) -> Self::Unsigned {
+                self.cast::<$unsigned>().leading_zeros()
             }
 
             #[inline]
-            fn trailing_zeros(self) -> Self {
-                // Safety: `self` is an integer vector
-                unsafe { intrinsics::simd_cttz(self) }
+            fn trailing_zeros(self) -> Self::Unsigned {
+                self.cast::<$unsigned>().trailing_zeros()
             }
 
             #[inline]
-            fn leading_ones(self) -> Self {
-                (!self).leading_zeros()
+            fn leading_ones(self) -> Self::Unsigned {
+                self.cast::<$unsigned>().leading_ones()
             }
 
             #[inline]
-            fn trailing_ones(self) -> Self {
-                (!self).trailing_zeros()
+            fn trailing_ones(self) -> Self::Unsigned {
+                self.cast::<$unsigned>().trailing_ones()
             }
         }
         )*
     }
 }
 
-impl_trait! { i8, i16, i32, i64, isize }
+impl_trait! { i8 (u8), i16 (u16), i32 (u32), i64 (u64), isize (usize) }
diff --git a/crates/core_simd/tests/ops_macros.rs b/crates/core_simd/tests/ops_macros.rs
index 23e914e64b5..135f3ecf7b2 100644
--- a/crates/core_simd/tests/ops_macros.rs
+++ b/crates/core_simd/tests/ops_macros.rs
@@ -213,7 +213,7 @@ fn reverse_bits<const LANES: usize>() {
             fn leading_zeros<const LANES: usize>() {
                 test_helpers::test_unary_elementwise(
                     &$vector::<LANES>::leading_zeros,
-                    &|x| x.leading_zeros() as $scalar,
+                    &|x| x.leading_zeros() as _,
                     &|_| true,
                 )
             }
@@ -221,7 +221,7 @@ fn leading_zeros<const LANES: usize>() {
             fn trailing_zeros<const LANES: usize>() {
                 test_helpers::test_unary_elementwise(
                     &$vector::<LANES>::trailing_zeros,
-                    &|x| x.trailing_zeros() as $scalar,
+                    &|x| x.trailing_zeros() as _,
                     &|_| true,
                 )
             }
@@ -229,7 +229,7 @@ fn trailing_zeros<const LANES: usize>() {
             fn leading_ones<const LANES: usize>() {
                 test_helpers::test_unary_elementwise(
                     &$vector::<LANES>::leading_ones,
-                    &|x| x.leading_ones() as $scalar,
+                    &|x| x.leading_ones() as _,
                     &|_| true,
                 )
             }
@@ -237,7 +237,7 @@ fn leading_ones<const LANES: usize>() {
             fn trailing_ones<const LANES: usize>() {
                 test_helpers::test_unary_elementwise(
                     &$vector::<LANES>::trailing_ones,
-                    &|x| x.trailing_ones() as $scalar,
+                    &|x| x.trailing_ones() as _,
                     &|_| true,
                 )
             }