Feature/optimization

iKeySort/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "i_key_sort"
-version = "0.11.0"
+version = "0.12.0"
 authors = ["Nail Sharipov <nailxsharipov@gmail.com>"]
 edition = "2024"
 description = "Counting sort algorithm."
@@ -15,7 +15,7 @@ lto = true
 codegen-units = 1
 
 [dependencies]
-rayon = { optional = true, version = "^1.11" }
+rayon = { optional = true, version = "^1.12" }
 
 [features]
 

iKeySort/src/sort/mod.rs

@@ -1,13 +1,13 @@
-#[cfg(feature = "allow_multithreading")]
-mod parallel;
-mod serial;
 mod bin_layout;
-mod mapper;
-mod min_max;
-mod spread;
 mod buffer;
 pub mod key;
+mod mapper;
+mod min_max;
 pub mod one_key;
 pub mod one_key_cmp;
+#[cfg(feature = "allow_multithreading")]
+mod parallel;
+mod serial;
+mod spread;
 pub mod two_keys;
-pub mod two_keys_cmp;
+pub mod two_keys_cmp;

iKeySort/src/sort/one_key.rs

@@ -34,12 +34,8 @@ pub trait OneKeySort<T> {
         K: SortKey,
         F: KeyFn<T, K>;
 
-    fn sort_by_one_key_and_buffer<K, F>(
-        &mut self,
-        parallel: bool,
-        buffer: &mut Vec<T>,
-        key: F,
-    ) where
+    fn sort_by_one_key_and_buffer<K, F>(&mut self, parallel: bool, buffer: &mut Vec<T>, key: F)
+    where
         K: SortKey,
         F: KeyFn<T, K>;
 }
@@ -117,12 +113,8 @@ where
     }
 
     #[inline]
-    fn sort_by_one_key_and_buffer<K, F>(
-        &mut self,
-        parallel: bool,
-        buffer: &mut Vec<T>,
-        key: F,
-    ) where
+    fn sort_by_one_key_and_buffer<K, F>(&mut self, parallel: bool, buffer: &mut Vec<T>, key: F)
+    where
         K: SortKey + Send + Sync,
         F: KeyFn<T, K> + Send + Sync,
     {
@@ -140,11 +132,10 @@ where
     }
 }
 
-
 #[cfg(test)]
 mod tests {
-    use alloc::vec::Vec;
     use crate::sort::one_key::OneKeySort;
+    use alloc::vec::Vec;
 
     #[test]
     fn test_0() {

iKeySort/src/sort/one_key_cmp.rs

@@ -155,7 +155,7 @@ where
     F1: KeyFn<T, K>,
     F2: CmpFn<T>,
 {
-    slice.sort_unstable_by(|a, b| key(a).cmp(&key(b)).then(compare(a, b)));
+    slice.sort_unstable_by(|a, b| key(a).cmp(&key(b)).then_with(|| compare(a, b)));
 }
 
 #[cfg(test)]

iKeySort/src/sort/parallel/bin_layout.rs

@@ -1,10 +1,10 @@
 use crate::sort::bin_layout::{BinLayout, LayoutConstraints, MAX_BINS_COUNT};
 use crate::sort::key::{KeyFn, SortKey};
-use crate::sort::parallel::pre_sort::{PreSortFragment, IdRange, FragmentationByCount};
-use rayon::iter::ParallelIterator;
-use rayon::iter::{IntoParallelRefIterator, IntoParallelRefMutIterator};
+use crate::sort::parallel::pre_sort::{ScratchRanges, spread_fragment};
+use crate::sort::parallel::sub_sort::BinMarks;
 use core::ops::Range;
-use core::ptr;
+use rayon::iter::ParallelIterator;
+use rayon::prelude::*;
 
 impl<K: SortKey + Send + Sync> BinLayout<K> {
     #[inline]
@@ -22,7 +22,13 @@ impl<K: SortKey + Send + Sync> BinLayout<K> {
 
         let max_bins_count = cpu.saturating_mul(4).min(MAX_BINS_COUNT);
 
-        Some(BinLayout::with_constraints(min_key, max_key, LayoutConstraints { max_split_count: max_bins_count }))
+        Some(BinLayout::with_constraints(
+            min_key,
+            max_key,
+            LayoutConstraints {
+                max_split_count: max_bins_count,
+            },
+        ))
     }
 }
 
@@ -65,7 +71,7 @@ impl<K: SortKey + Send + Sync> BinLayout<K> {
         src: &mut [T],
         buf: &mut Vec<T>,
         key: F,
-    ) -> Vec<usize>
+    ) -> BinMarks
     where
         T: Copy + Send + Sync,
         F: KeyFn<T, K> + Send + Sync,
@@ -75,99 +81,151 @@ impl<K: SortKey + Send + Sync> BinLayout<K> {
             buf.reserve(src.len());
         }
 
-        let mut fragments = src.fragment_by_count(buf, cpu);
-
-        let groups = self.par_spread(&mut fragments, key);
-
-        // at this time buffer is fully initialized
-        #[allow(clippy::uninit_vec)]
-        unsafe { buf.set_len(src.len()); }
-
-        copy_groups(src, buf, groups)
-    }
-
-    #[inline]
-    fn par_spread<T, F>(&self, fragments: &mut [PreSortFragment<T>], key: F) -> Vec<Vec<Range<usize>>>
-    where
-        T: Copy + Send + Sync,
-        F: KeyFn<T, K> + Send + Sync,
-    {
         let bins_count = self.count();
-        let frags_count = fragments.len();
-
-        let id_ranges = fragments
-            .par_iter_mut()
-            .map(|f| f.spread(self.clone(), key))
-            .reduce(
-                || Vec::<IdRange>::with_capacity(frags_count * bins_count),
-                |mut a, mut b| {
-                    a.append(&mut b);
-                    a
-                },
-            );
-
-        let mut counter = vec![0usize; bins_count];
-        for range in id_ranges.iter() {
-            counter[range.index] += 1;
-        }
-
-        let mut groups: Vec<Vec<Range<usize>>> = (0..bins_count).map(|_| Vec::new()).collect();
-        for (group, &count) in groups.iter_mut().zip(counter.iter()) {
-            if count > 0 {
-                *group = Vec::with_capacity(count);
-            }
-        }
-
-        for id_range in id_ranges {
-            unsafe {
-                groups
-                    .get_unchecked_mut(id_range.index)
-                    .push(id_range.range);
-            }
+        let step_len = src.len().div_ceil(cpu);
+        let fragments_count = src.len().div_ceil(step_len);
+        let ranges_count = fragments_count
+            .checked_mul(bins_count)
+            .expect("too many bin ranges");
+        let mut scratch = ScratchRanges::new(ranges_count);
+
+        src.par_chunks(step_len)
+            .zip(buf.spare_capacity_mut()[..src.len()].par_chunks_mut(step_len))
+            .zip(scratch.ranges.par_chunks_mut(bins_count))
+            .enumerate()
+            .for_each(|(index, ((src, buf), ranges))| {
+                spread_fragment(src, buf, ranges, index * step_len, self, key);
+            });
+
+        // SAFETY: every fragment has initialized its disjoint part of the buffer.
+        // If a key function panics, execution never reaches set_len.
+        unsafe {
+            buf.set_len(src.len());
         }
 
-        groups.retain(|g| !g.is_empty());
-        groups
+        copy_groups(src, buf, &scratch.ranges, bins_count)
     }
 }
 
 #[inline]
-fn copy_groups<T>(mut dst: &mut [T], src: &[T], mut groups: Vec<Vec<Range<usize>>>) -> Vec<usize> {
-    let mut marks = Vec::with_capacity(groups.len());
-    let last_group = if let Some(last) = groups.pop() {
-        last
-    } else {
-        return marks;
-    };
-
-    let mut base = 0;
-    for ranges in groups.iter() {
-        let length = copy_ranges(dst, src, ranges);
-        (_, dst) = dst.split_at_mut(length);
-
-        base += length;
-        marks.push(base);
+fn copy_groups<T: Copy>(
+    dst: &mut [T],
+    src: &[T],
+    ranges: &[Range<usize>],
+    bins_count: usize,
+) -> BinMarks {
+    let mut marks = BinMarks::new();
+    let mut offset = 0;
+    for bin in 0..bins_count {
+        let start = offset;
+        for row in ranges.chunks_exact(bins_count) {
+            let range = &row[bin];
+            if range.is_empty() {
+                continue;
+            }
+            let end = offset + range.len();
+            dst[offset..end].copy_from_slice(&src[range.clone()]);
+            offset = end;
+        }
+        if offset > start && offset < dst.len() {
+            marks.push(offset);
+        }
     }
-
-    let _ = copy_ranges(dst, src, &last_group);
-
+    debug_assert_eq!(offset, dst.len());
     marks
 }
 
-#[inline]
-fn copy_ranges<T>(dst: &mut [T], src: &[T], ranges: &[Range<usize>]) -> usize {
-    let mut offset = 0;
-    let dst_base = dst.as_mut_ptr();
-    let src_base = src.as_ptr();
-    for range in ranges.iter() {
-        unsafe {
-            let dst_ptr = dst_base.add(offset);
-            let src_ptr = src_base.add(range.start);
-            ptr::copy_nonoverlapping(src_ptr, dst_ptr, range.len());
-        }
+#[cfg(test)]
+mod scratch_tests {
+    use super::BinLayout;
+    use crate::sort::parallel::sub_sort::FragmentationByMarks;
+    use rayon::iter::ParallelIterator;
+    use std::panic::{AssertUnwindSafe, catch_unwind};
+    use std::sync::atomic::{AtomicBool, Ordering};
+
+    #[test]
+    fn sparse_and_full_bins_reuse_scratch_with_uneven_fragments() {
+        let pool = rayon::ThreadPoolBuilder::new()
+            .num_threads(3)
+            .build()
+            .unwrap();
+        pool.install(|| {
+            let mut buffer = Vec::new();
+            for bins in [256, 2, 17, 256] {
+                let mut values: Vec<i64> = (0..1_031).rev().map(|i| i % bins).collect();
+                let mut expected = values.clone();
+                expected.sort_unstable();
+                let layout = BinLayout::with_keys(&values, |x| *x).unwrap();
+                let marks = layout.par_pre_sort(3, &mut values, &mut buffer, |x| *x);
+                assert_eq!(values, expected);
+                assert_eq!(marks.len(), bins as usize - 1);
+                values
+                    .fragment_by_marks(&mut buffer, &marks)
+                    .par_iter_mut()
+                    .for_each(|fragment| {
+                        assert!(fragment.src.iter().all(|x| *x == fragment.src[0]))
+                    });
+            }
+            let mut values = [i64::MIN, 0, i64::MAX].repeat(347);
+            let mut expected = values.clone();
+            expected.sort_unstable();
+            let layout = BinLayout::with_keys(&values, |x| *x).unwrap();
+            let marks = layout.par_pre_sort(3, &mut values, &mut buffer, |x| *x);
+            assert_eq!(values, expected);
+            assert_eq!(marks.len(), 2);
+        });
+    }
 
-        offset += range.len();
+    #[test]
+    fn nested_spread_can_rent_scratch_on_the_same_thread() {
+        let pool = rayon::ThreadPoolBuilder::new()
+            .num_threads(1)
+            .build()
+            .unwrap();
+        pool.install(|| {
+            let mut values: Vec<i64> = (0..64).rev().collect();
+            let layout = BinLayout::with_keys(&values, |x| *x).unwrap();
+            let nested = AtomicBool::new(false);
+            layout.par_pre_sort(1, &mut values, &mut Vec::new(), |x| {
+                if !nested.swap(true, Ordering::Relaxed) {
+                    let mut inner: Vec<i64> = (0..256).rev().collect();
+                    let inner_layout = BinLayout::with_keys(&inner, |x| *x).unwrap();
+                    inner_layout.par_pre_sort(1, &mut inner, &mut Vec::new(), |x| *x);
+                    assert!(inner.is_sorted());
+                }
+                *x
+            });
+            assert!(values.is_sorted());
+            assert!(nested.load(Ordering::Relaxed));
+        });
     }
 
-    offset
+    #[test]
+    fn panic_during_scatter_keeps_buffer_uninitialized_and_reusable() {
+        let pool = rayon::ThreadPoolBuilder::new()
+            .num_threads(2)
+            .build()
+            .unwrap();
+        pool.install(|| {
+            let mut values: Vec<i64> = (0..128).rev().collect();
+            let original = values.clone();
+            let layout = BinLayout::with_keys(&values, |x| *x).unwrap();
+            let seen = AtomicBool::new(false);
+            let mut buffer = vec![0; values.len()];
+            let result = catch_unwind(AssertUnwindSafe(|| {
+                layout.par_pre_sort(2, &mut values, &mut buffer, |x| {
+                    if *x == 0 && seen.swap(true, Ordering::Relaxed) {
+                        panic!("key failed during scatter");
+                    }
+                    *x
+                });
+            }));
+            assert!(result.is_err());
+            assert!(buffer.is_empty());
+            assert_eq!(values, original);
+            layout.par_pre_sort(2, &mut values, &mut buffer, |x| *x);
+            assert!(values.is_sorted());
+            assert_eq!(buffer.len(), values.len());
+        });
+    }
 }

iKeySort/src/sort/parallel/cpu_count.rs

@@ -3,7 +3,7 @@ pub(super) struct CPUCount;
 impl CPUCount {
     #[inline]
     pub(super) fn count() -> usize {
-        std::thread::available_parallelism().map_or(1, |n| n.get())
+        rayon::current_num_threads()
     }
 
     #[cfg(not(debug_assertions))]

iKeySort/src/sort/parallel/mod.rs

@@ -1,8 +1,8 @@
+mod bin_layout;
 mod cpu_count;
 mod pre_sort;
 pub(crate) mod slice_one_key;
 pub(crate) mod slice_one_key_cmp;
 pub(crate) mod slice_two_keys;
 pub(crate) mod slice_two_keys_cmp;
-mod bin_layout;
 mod sub_sort;