Implement mask_up_to properly for vectors 128 elements and larger

Author: farnoy

crates/core_simd/src/masks.rs

@@ -34,6 +34,7 @@ mod sealed {
         fn eq(self, other: Self) -> bool;
 
         fn to_usize(self) -> usize;
+        fn max_unsigned() -> u64;
 
         type Unsigned: SimdElement;
 
@@ -78,6 +79,11 @@ macro_rules! impl_element {
                 self as usize
             }
 
+            #[inline]
+            fn max_unsigned() -> u64 {
+                <$unsigned>::MAX as u64
+            }
+
             type Unsigned = $unsigned;
 
             const TRUE: Self = -1;

crates/core_simd/src/vector.rs

@@ -1,6 +1,7 @@
 use crate::simd::{
     cmp::SimdPartialOrd,
     ptr::{SimdConstPtr, SimdMutPtr},
+    num::SimdUint,
     LaneCount, Mask, MaskElement, SupportedLaneCount, Swizzle,
 };
 
@@ -1190,14 +1191,13 @@ where
 }
 
 #[inline]
-fn lane_indices<T, const N: usize>() -> Simd<T, N>
+fn lane_indices<const N: usize>() -> Simd<usize, N>
 where
-    T: MaskElement + Default + core::convert::From<i8> + core::ops::Add<T, Output = T>,
     LaneCount<N>: SupportedLaneCount,
 {
-    let mut index = [T::default(); N];
-    for i in 1..N {
-        index[i] = index[i - 1] + T::from(1);
+    let mut index = [0; N];
+    for i in 0..N {
+        index[i] = i;
     }
     Simd::from_array(index)
 }
@@ -1208,7 +1208,18 @@ where
     LaneCount<N>: SupportedLaneCount,
     M: MaskElement,
 {
-    let index = lane_indices::<i8, N>();
-    let lt = index.simd_lt(Simd::splat(i8::try_from(len).unwrap_or(i8::MAX)));
-    lt.cast()
+    let index = lane_indices::<N>();
+    let max_value: u64 = M::max_unsigned();
+    macro_rules! case {
+        ($ty:ty) => {
+            if N < <$ty>::MAX as usize && max_value as $ty as u64 == max_value {
+                return index.cast().simd_lt(Simd::splat(len.min(N) as $ty)).cast();
+            } 
+        };
+    }
+    case!(u8);
+    case!(u16);
+    case!(u32);
+    case!(u64);
+    index.simd_lt(Simd::splat(len)).cast()
 }