Fix/wasm usize overflow

iKeySort/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "i_key_sort"
-version = "0.10.3"
+version = "0.11.0"
 authors = ["Nail Sharipov <nailxsharipov@gmail.com>"]
 edition = "2024"
 description = "Counting sort algorithm."
@@ -21,4 +21,4 @@ rayon = { optional = true, version = "^1.11" }
 
 default = []
 std = []
-allow_multithreading = ["std", "dep:rayon"]
+allow_multithreading = ["std", "dep:rayon"]

iKeySort/src/sort/bin_layout.rs

@@ -36,34 +36,35 @@ impl<K: SortKey> BinLayout<K> {
     #[inline(always)]
     pub fn index(&self, value: K) -> usize {
         debug_assert!(value >= self.min_key, "value must be >= min_key");
-        let offset = value.difference(self.min_key);
-        offset >> self.power
+        value.shifted_distance(self.min_key, self.power)
     }
 
     #[inline(always)]
     pub fn count(&self) -> usize {
         self.index(self.max_key) + 1
     }
 
-    pub(crate) fn with_constraints(min_key: K, max_key: K, constraints: LayoutConstraints) -> BinLayout<K> {
-        let length = max_key.difference(min_key);
-        if length < constraints.max_split_count {
-            return Self {
-                min_key,
-                max_key,
-                power: 0,
-                bin_width_is_one: true,
-            };
-        }
-
-        let scale = length.saturating_add(1).ilog2_ceil();
-        let power = scale.saturating_sub(constraints.max_split_count.ilog2()) as usize;
+    pub(crate) fn with_constraints(
+        min_key: K,
+        max_key: K,
+        constraints: LayoutConstraints,
+    ) -> BinLayout<K> {
+        // `max_split_count` is a sizing hint. Cap it at half of `usize::MAX`
+        let max_split_count = constraints.max_split_count.clamp(1, usize::MAX / 2);
+        let distance = max_key.shifted_distance(min_key, 0);
+        let power = if distance < max_split_count {
+            0
+        } else {
+            max_key
+                .distance_bits(min_key)
+                .saturating_sub(max_split_count.ilog2() as usize)
+        };
 
         Self {
             min_key,
             max_key,
             power,
-            bin_width_is_one: false,
+            bin_width_is_one: power == 0,
         }
     }
 
@@ -83,25 +84,9 @@ impl<K: SortKey> BinLayout<K> {
     }
 }
 
-trait Log2 {
-    fn ilog2_ceil(&self) -> u32;
-}
-
-impl Log2 for usize {
-    #[inline(always)]
-    fn ilog2_ceil(&self) -> u32 {
-        let floor = self.ilog2();
-        if self.is_power_of_two() {
-            floor
-        } else {
-            floor + 1
-        }
-    }
-}
-
 #[cfg(test)]
 mod tests {
-    use crate::sort::bin_layout::{BinLayout, Log2};
+    use crate::sort::bin_layout::{BinLayout, LayoutConstraints, MAX_BINS_COUNT};
 
     #[test]
     fn test_0() {
@@ -124,11 +109,52 @@ mod tests {
     }
 
     #[test]
-    fn test_log2_0() {
-        assert_eq!(1usize.ilog2_ceil(), 0);
-        assert_eq!(2usize.ilog2_ceil(), 1);
-        assert_eq!(3usize.ilog2_ceil(), 2);
-        assert_eq!(4usize.ilog2_ceil(), 2);
-        assert_eq!(5usize.ilog2_ceil(), 3);
+    fn test_i64_full_range() {
+        let layout = BinLayout::<i64>::with_constraints(i64::MIN, i64::MAX, Default::default());
+
+        assert_eq!(layout.power, 56);
+        assert_eq!(layout.index(i64::MIN), 0);
+        assert_eq!(layout.index(i64::MAX), MAX_BINS_COUNT - 1);
+        assert_eq!(layout.count(), MAX_BINS_COUNT);
+    }
+
+    #[test]
+    fn test_max_split_count_is_limited_to_usize() {
+        let layout = BinLayout::<usize>::with_constraints(
+            usize::MIN,
+            usize::MAX,
+            LayoutConstraints {
+                max_split_count: usize::MAX,
+            },
+        );
+
+        assert_eq!(layout.power, 2);
+        assert_eq!(layout.count(), (usize::MAX >> 2) + 1);
+    }
+
+    #[test]
+    fn test_non_power_of_two_max_split_count() {
+        let constraints = LayoutConstraints {
+            max_split_count: 160,
+        };
+
+        let layout = BinLayout::<i64>::with_constraints(0, 159, constraints);
+        assert_eq!(layout.power, 0);
+        assert_eq!(layout.count(), 160);
+
+        let layout = BinLayout::<i64>::with_constraints(0, 40_000, constraints);
+        assert_eq!(layout.power, 9);
+        assert_eq!(layout.count(), 79);
+    }
+
+    #[test]
+    fn test_zero_max_split_count() {
+        let layout = BinLayout::<i64>::with_constraints(
+            i64::MIN,
+            i64::MAX,
+            LayoutConstraints { max_split_count: 0 },
+        );
+
+        assert_eq!(layout.count(), 1);
     }
 }

iKeySort/src/sort/key.rs

@@ -7,77 +7,66 @@ pub trait CmpFn<T>: Fn(&T, &T) -> Ordering + Copy {}
 impl<T, F: Fn(&T, &T) -> Ordering + Copy> CmpFn<T> for F {}
 
 pub trait SortKey: Copy + Ord {
-    fn difference(self, other: Self) -> usize;
-}
-
-impl SortKey for u8 {
-    #[inline(always)]
-    fn difference(self, other: Self) -> usize {
-        debug_assert!(self >= other, "difference() requires self >= other");
-        (self - other) as usize
-    }
-}
-
-impl SortKey for i8 {
-    #[inline(always)]
-    fn difference(self, other: Self) -> usize {
-        debug_assert!(self >= other, "difference() requires self >= other");
-        (self - other) as usize
-    }
-}
+    /// Returns the number of bits required to represent `self - other`.
+    fn distance_bits(self, other: Self) -> usize;
 
-impl SortKey for u16 {
-    #[inline(always)]
-    fn difference(self, other: Self) -> usize {
-        debug_assert!(self >= other, "difference() requires self >= other");
-        (self - other) as usize
-    }
+    /// Returns `(self - other) >> shift`, saturated to the `usize` range.
+    ///
+    /// The distance must be shifted before it is converted to `usize`.
+    fn shifted_distance(self, other: Self, shift: usize) -> usize;
 }
 
-impl SortKey for i16 {
-    #[inline(always)]
-    fn difference(self, other: Self) -> usize {
-        debug_assert!(self >= other, "difference() requires self >= other");
-        (self - other) as usize
-    }
-}
+macro_rules! impl_unsigned_sort_key {
+    ($($ty:ty),+ $(,)?) => {
+        $(
+            impl SortKey for $ty {
+                #[inline(always)]
+                fn distance_bits(self, other: Self) -> usize {
+                    debug_assert!(self >= other, "distance_bits() requires self >= other");
+                    let distance = self - other;
+                    (<$ty>::BITS - distance.leading_zeros()) as usize
+                }
 
-impl SortKey for u32 {
-    #[inline(always)]
-    fn difference(self, other: Self) -> usize {
-        debug_assert!(self >= other, "difference() requires self >= other");
-        (self - other) as usize
-    }
+                #[inline(always)]
+                fn shifted_distance(self, other: Self, shift: usize) -> usize {
+                    debug_assert!(self >= other, "shifted_distance() requires self >= other");
+                    let distance = self - other;
+                    let shifted = u32::try_from(shift)
+                        .ok()
+                        .and_then(|shift| distance.checked_shr(shift))
+                        .unwrap_or(0);
+                    usize::try_from(shifted).unwrap_or(usize::MAX)
+                }
+            }
+        )+
+    };
 }
 
-impl SortKey for i32 {
-    #[inline(always)]
-    fn difference(self, other: Self) -> usize {
-        debug_assert!(self >= other, "difference() requires self >= other");
-        (self - other) as usize
-    }
-}
+macro_rules! impl_signed_sort_key {
+    ($($ty:ty),+ $(,)?) => {
+        $(
+            impl SortKey for $ty {
+                #[inline(always)]
+                fn distance_bits(self, other: Self) -> usize {
+                    debug_assert!(self >= other, "distance_bits() requires self >= other");
+                    let distance = self.abs_diff(other);
+                    (<$ty>::BITS - distance.leading_zeros()) as usize
+                }
 
-impl SortKey for u64 {
-    #[inline(always)]
-    fn difference(self, other: Self) -> usize {
-        debug_assert!(self >= other, "difference() requires self >= other");
-        (self - other) as usize
-    }
+                #[inline(always)]
+                fn shifted_distance(self, other: Self, shift: usize) -> usize {
+                    debug_assert!(self >= other, "shifted_distance() requires self >= other");
+                    let distance = self.abs_diff(other);
+                    let shifted = u32::try_from(shift)
+                        .ok()
+                        .and_then(|shift| distance.checked_shr(shift))
+                        .unwrap_or(0);
+                    usize::try_from(shifted).unwrap_or(usize::MAX)
+                }
+            }
+        )+
+    };
 }
 
-impl SortKey for i64 {
-    #[inline(always)]
-    fn difference(self, other: Self) -> usize {
-        debug_assert!(self >= other, "difference() requires self >= other");
-        (self - other) as usize
-    }
-}
-
-impl SortKey for usize {
-    #[inline(always)]
-    fn difference(self, other: Self) -> usize {
-        debug_assert!(self >= other, "difference() requires self >= other");
-        self - other
-    }
-}
+impl_unsigned_sort_key!(u8, u16, u32, u64, usize);
+impl_signed_sort_key!(i8, i16, i32, i64);

iKeySort/src/sort/one_key.rs

@@ -173,7 +173,7 @@ mod tests {
 
     #[test]
     fn test_5() {
-        test(1000_000);
+        test(1_000_000);
     }
 
     #[test]
@@ -193,4 +193,4 @@ mod tests {
         assert!(arr1 == org);
         assert!(arr2 == org);
     }
-}
+}

iKeySort/src/sort/parallel/slice_one_key.rs

@@ -91,7 +91,7 @@ mod tests {
 
     #[test]
     fn test_5() {
-        test(1000_000);
+        test(1_000_000);
     }
 
     fn test(count: usize) {

iKeySort/src/sort/serial/slice_one_key.rs

@@ -83,7 +83,7 @@ mod tests {
 
     #[test]
     fn test_5() {
-        test(1000_000);
+        test(1_000_000);
     }
 
     fn test(count: usize) {

iKeySort/tests/dynamic.rs

@@ -110,7 +110,7 @@ mod tests {
             })
             .collect();
 
-        segments.sort_unstable_by(|x0, x1| x0.cmp(&x1));
+        segments.sort_unstable();
 
         for arr in res {
             assert_eq!(arr, segments);
@@ -234,7 +234,7 @@ mod tests {
             a += da;
         }
 
-        let mut p0 = points.last().unwrap().clone();
+        let mut p0 = *points.last().unwrap();
         for &pi in points.iter() {
             result.push(Segment::new(p0, pi));
             p0 = pi;

performance/remap/Cargo.toml

@@ -4,4 +4,4 @@ version = "0.1.0"
 edition = "2024"
 
 [dependencies]
-i_key_sort = "0.9.1"
+i_key_sort = { path = "../../iKeySort" }