Implement partial_sort_unstable for slice

Signed-off-by: tison <wander4096@gmail.com>
Co-Authored-By: Orson Peters <orsonpeters@gmail.com>
Signed-off-by: tison <wander4096@gmail.com>

library/alloctests/tests/lib.rs

@@ -20,6 +20,8 @@
 #![feature(binary_heap_into_iter_sorted)]
 #![feature(binary_heap_drain_sorted)]
 #![feature(slice_ptr_get)]
+#![feature(slice_range)]
+#![feature(slice_partial_sort_unstable)]
 #![feature(inplace_iteration)]
 #![feature(iter_advance_by)]
 #![feature(iter_next_chunk)]

library/alloctests/tests/sort/mod.rs

@@ -12,6 +12,7 @@ fn sort_by<T, F>(v: &mut [T], compare: F)
 
 mod ffi_types;
 mod known_good_stable_sort;
+mod partial;
 mod patterns;
 mod tests;
 mod zipf;

library/alloctests/tests/sort/partial.rs

@@ -0,0 +1,84 @@
+use std::fmt::Debug;
+use std::ops::{Range, RangeBounds};
+use std::slice;
+
+use super::patterns;
+
+fn check_is_partial_sorted<T: Ord + Clone + Debug, R: RangeBounds<usize>>(v: &mut [T], range: R) {
+    let Range { start, end } = slice::range(range, ..v.len());
+    v.partial_sort_unstable(start..end);
+
+    let max_before = v[..start].iter().max().into_iter();
+    let sorted_range = v[start..end].into_iter();
+    let min_after = v[end..].iter().min().into_iter();
+    let seq = max_before.chain(sorted_range).chain(min_after);
+    assert!(seq.is_sorted());
+}
+
+fn check_is_partial_sorted_ranges<T: Ord + Clone + Debug>(v: &[T]) {
+    let len = v.len();
+
+    check_is_partial_sorted::<T, _>(&mut v.to_vec(), ..);
+    check_is_partial_sorted::<T, _>(&mut v.to_vec(), 0..0);
+    check_is_partial_sorted::<T, _>(&mut v.to_vec(), len..len);
+
+    if len > 0 {
+        check_is_partial_sorted::<T, _>(&mut v.to_vec(), len - 1..len - 1);
+        check_is_partial_sorted::<T, _>(&mut v.to_vec(), 0..1);
+        check_is_partial_sorted::<T, _>(&mut v.to_vec(), len - 1..len);
+
+        for mid in 1..len {
+            check_is_partial_sorted::<T, _>(&mut v.to_vec(), 0..mid);
+            check_is_partial_sorted::<T, _>(&mut v.to_vec(), mid..len);
+            check_is_partial_sorted::<T, _>(&mut v.to_vec(), mid..mid);
+            check_is_partial_sorted::<T, _>(&mut v.to_vec(), mid - 1..mid + 1);
+            check_is_partial_sorted::<T, _>(&mut v.to_vec(), mid - 1..mid);
+            check_is_partial_sorted::<T, _>(&mut v.to_vec(), mid..mid + 1);
+        }
+
+        let quarters = [0, len / 4, len / 2, (3 * len) / 4, len];
+        for &start in &quarters {
+            for &end in &quarters {
+                if start < end {
+                    check_is_partial_sorted::<T, _>(&mut v.to_vec(), start..end);
+                }
+            }
+        }
+    }
+}
+
+#[test]
+fn basic_impl() {
+    check_is_partial_sorted::<i32, _>(&mut [], ..);
+    check_is_partial_sorted::<(), _>(&mut [], ..);
+    check_is_partial_sorted::<(), _>(&mut [()], ..);
+    check_is_partial_sorted::<(), _>(&mut [(), ()], ..);
+    check_is_partial_sorted::<(), _>(&mut [(), (), ()], ..);
+    check_is_partial_sorted::<i32, _>(&mut [], ..);
+
+    check_is_partial_sorted::<i32, _>(&mut [77], ..);
+    check_is_partial_sorted::<i32, _>(&mut [2, 3], ..);
+    check_is_partial_sorted::<i32, _>(&mut [2, 3, 6], ..);
+    check_is_partial_sorted::<i32, _>(&mut [2, 3, 99, 6], ..);
+    check_is_partial_sorted::<i32, _>(&mut [2, 7709, 400, 90932], ..);
+    check_is_partial_sorted::<i32, _>(&mut [15, -1, 3, -1, -3, -1, 7], ..);
+
+    check_is_partial_sorted::<i32, _>(&mut [15, -1, 3, -1, -3, -1, 7], 0..0);
+    check_is_partial_sorted::<i32, _>(&mut [15, -1, 3, -1, -3, -1, 7], 0..1);
+    check_is_partial_sorted::<i32, _>(&mut [15, -1, 3, -1, -3, -1, 7], 0..5);
+    check_is_partial_sorted::<i32, _>(&mut [15, -1, 3, -1, -3, -1, 7], 0..7);
+    check_is_partial_sorted::<i32, _>(&mut [15, -1, 3, -1, -3, -1, 7], 7..7);
+    check_is_partial_sorted::<i32, _>(&mut [15, -1, 3, -1, -3, -1, 7], 6..7);
+    check_is_partial_sorted::<i32, _>(&mut [15, -1, 3, -1, -3, -1, 7], 5..7);
+    check_is_partial_sorted::<i32, _>(&mut [15, -1, 3, -1, -3, -1, 7], 5..5);
+    check_is_partial_sorted::<i32, _>(&mut [15, -1, 3, -1, -3, -1, 7], 4..5);
+    check_is_partial_sorted::<i32, _>(&mut [15, -1, 3, -1, -3, -1, 7], 4..6);
+}
+
+#[test]
+fn random_patterns() {
+    check_is_partial_sorted_ranges(&patterns::random(10));
+    check_is_partial_sorted_ranges(&patterns::random(50));
+    check_is_partial_sorted_ranges(&patterns::random(100));
+    check_is_partial_sorted_ranges(&patterns::random(1000));
+}

library/core/src/slice/mod.rs

@@ -3244,6 +3244,219 @@ pub fn sort_unstable_by_key<K, F>(&mut self, mut f: F)
         sort::unstable::sort(self, &mut |a, b| f(a).lt(&f(b)));
     }
 
+    /// Partially sorts the slice in ascending order **without** preserving the initial order of equal elements.
+    ///
+    /// Upon completion, for the specified range `start..end`, it's guaranteed that:
+    ///
+    /// 1. Every element in `self[..start]` is smaller than or equal to
+    /// 2. Every element in `self[start..end]`, which is sorted, and smaller than or equal to
+    /// 3. Every element in `self[end..]`.
+    ///
+    /// This partial sort is unstable, meaning it may reorder equal elements in the specified range.
+    /// It may reorder elements outside the specified range as well, but the guarantees above still hold.
+    ///
+    /// This partial sort is in-place (i.e., does not allocate), and *O*(*n* + *k* \* log(*k*)) worst-case,
+    /// where *n* is the length of the slice and *k* is the length of the specified range.
+    ///
+    /// See the documentation of [`sort_unstable`] for implementation notes.
+    ///
+    /// # Panics
+    ///
+    /// May panic if the implementation of [`Ord`] for `T` does not implement a total order, or if
+    /// the [`Ord`] implementation panics, or if the specified range is out of bounds.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// #![feature(slice_partial_sort_unstable)]
+    ///
+    /// let mut v = [4, -5, 1, -3, 2];
+    ///
+    /// // empty range at the beginning, nothing changed
+    /// v.partial_sort_unstable(0..0);
+    /// assert_eq!(v, [4, -5, 1, -3, 2]);
+    ///
+    /// // empty range in the middle, partitioning the slice
+    /// v.partial_sort_unstable(2..2);
+    /// for i in 0..2 {
+    ///    assert!(v[i] <= v[2]);
+    /// }
+    /// for i in 3..v.len() {
+    ///   assert!(v[2] <= v[i]);
+    /// }
+    ///
+    /// // single element range, same as select_nth_unstable
+    /// v.partial_sort_unstable(2..3);
+    /// for i in 0..2 {
+    ///    assert!(v[i] <= v[2]);
+    /// }
+    /// for i in 3..v.len() {
+    ///   assert!(v[2] <= v[i]);
+    /// }
+    ///
+    /// // partial sort a subrange
+    /// v.partial_sort_unstable(1..4);
+    /// assert_eq!(&v[1..4], [-3, 1, 2]);
+    ///
+    /// // partial sort the whole range, same as sort_unstable
+    /// v.partial_sort_unstable(..);
+    /// assert_eq!(v, [-5, -3, 1, 2, 4]);
+    /// ```
+    ///
+    /// [`sort_unstable`]: slice::sort_unstable
+    #[unstable(feature = "slice_partial_sort_unstable", issue = "149046")]
+    #[inline]
+    pub fn partial_sort_unstable<R>(&mut self, range: R)
+    where
+        T: Ord,
+        R: RangeBounds<usize>,
+    {
+        sort::unstable::partial_sort(self, range, T::lt);
+    }
+
+    /// Partially sorts the slice in ascending order with a comparison function, **without**
+    /// preserving the initial order of equal elements.
+    ///
+    /// Upon completion, for the specified range `start..end`, it's guaranteed that:
+    ///
+    /// 1. Every element in `self[..start]` is smaller than or equal to
+    /// 2. Every element in `self[start..end]`, which is sorted, and smaller than or equal to
+    /// 3. Every element in `self[end..]`.
+    ///
+    /// This partial sort is unstable, meaning it may reorder equal elements in the specified range.
+    /// It may reorder elements outside the specified range as well, but the guarantees above still hold.
+    ///
+    /// This partial sort is in-place (i.e., does not allocate), and *O*(*n* + *k* \* log(*k*)) worst-case,
+    /// where *n* is the length of the slice and *k* is the length of the specified range.
+    ///
+    /// See the documentation of [`sort_unstable_by`] for implementation notes.
+    ///
+    /// # Panics
+    ///
+    /// May panic if the `compare` does not implement a total order, or if
+    /// the `compare` itself panics, or if the specified range is out of bounds.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// #![feature(slice_partial_sort_unstable)]
+    ///
+    /// let mut v = [4, -5, 1, -3, 2];
+    ///
+    /// // empty range at the beginning, nothing changed
+    /// v.partial_sort_unstable_by(0..0, |a, b| b.cmp(a));
+    /// assert_eq!(v, [4, -5, 1, -3, 2]);
+    ///
+    /// // empty range in the middle, partitioning the slice
+    /// v.partial_sort_unstable_by(2..2, |a, b| b.cmp(a));
+    /// for i in 0..2 {
+    ///    assert!(v[i] >= v[2]);
+    /// }
+    /// for i in 3..v.len() {
+    ///   assert!(v[2] >= v[i]);
+    /// }
+    ///
+    /// // single element range, same as select_nth_unstable
+    /// v.partial_sort_unstable_by(2..3, |a, b| b.cmp(a));
+    /// for i in 0..2 {
+    ///    assert!(v[i] >= v[2]);
+    /// }
+    /// for i in 3..v.len() {
+    ///   assert!(v[2] >= v[i]);
+    /// }
+    ///
+    /// // partial sort a subrange
+    /// v.partial_sort_unstable_by(1..4, |a, b| b.cmp(a));
+    /// assert_eq!(&v[1..4], [2, 1, -3]);
+    ///
+    /// // partial sort the whole range, same as sort_unstable
+    /// v.partial_sort_unstable_by(.., |a, b| b.cmp(a));
+    /// assert_eq!(v, [4, 2, 1, -3, -5]);
+    /// ```
+    ///
+    /// [`sort_unstable_by`]: slice::sort_unstable_by
+    #[unstable(feature = "slice_partial_sort_unstable", issue = "149046")]
+    #[inline]
+    pub fn partial_sort_unstable_by<F, R>(&mut self, range: R, mut compare: F)
+    where
+        F: FnMut(&T, &T) -> Ordering,
+        R: RangeBounds<usize>,
+    {
+        sort::unstable::partial_sort(self, range, |a, b| compare(a, b) == Less);
+    }
+
+    /// Partially sorts the slice in ascending order with a key extraction function, **without**
+    /// preserving the initial order of equal elements.
+    ///
+    /// Upon completion, for the specified range `start..end`, it's guaranteed that:
+    ///
+    /// 1. Every element in `self[..start]` is smaller than or equal to
+    /// 2. Every element in `self[start..end]`, which is sorted, and smaller than or equal to
+    /// 3. Every element in `self[end..]`.
+    ///
+    /// This partial sort is unstable, meaning it may reorder equal elements in the specified range.
+    /// It may reorder elements outside the specified range as well, but the guarantees above still hold.
+    ///
+    /// This partial sort is in-place (i.e., does not allocate), and *O*(*n* + *k* \* log(*k*)) worst-case,
+    /// where *n* is the length of the slice and *k* is the length of the specified range.
+    ///
+    /// See the documentation of [`sort_unstable_by_key`] for implementation notes.
+    ///
+    /// # Panics
+    ///
+    /// May panic if the implementation of [`Ord`] for `K` does not implement a total order, or if
+    /// the [`Ord`] implementation panics, or if the specified range is out of bounds.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// #![feature(slice_partial_sort_unstable)]
+    ///
+    /// let mut v = [4i32, -5, 1, -3, 2];
+    ///
+    /// // empty range at the beginning, nothing changed
+    /// v.partial_sort_unstable_by_key(0..0, |k| k.abs());
+    /// assert_eq!(v, [4, -5, 1, -3, 2]);
+    ///
+    /// // empty range in the middle, partitioning the slice
+    /// v.partial_sort_unstable_by_key(2..2, |k| k.abs());
+    /// for i in 0..2 {
+    ///    assert!(v[i].abs() <= v[2].abs());
+    /// }
+    /// for i in 3..v.len() {
+    ///   assert!(v[2].abs() <= v[i].abs());
+    /// }
+    ///
+    /// // single element range, same as select_nth_unstable
+    /// v.partial_sort_unstable_by_key(2..3, |k| k.abs());
+    /// for i in 0..2 {
+    ///    assert!(v[i].abs() <= v[2].abs());
+    /// }
+    /// for i in 3..v.len() {
+    ///   assert!(v[2].abs() <= v[i].abs());
+    /// }
+    ///
+    /// // partial sort a subrange
+    /// v.partial_sort_unstable_by_key(1..4, |k| k.abs());
+    /// assert_eq!(&v[1..4], [2, -3, 4]);
+    ///
+    /// // partial sort the whole range, same as sort_unstable
+    /// v.partial_sort_unstable_by_key(.., |k| k.abs());
+    /// assert_eq!(v, [1, 2, -3, 4, -5]);
+    /// ```
+    ///
+    /// [`sort_unstable_by_key`]: slice::sort_unstable_by_key
+    #[unstable(feature = "slice_partial_sort_unstable", issue = "149046")]
+    #[inline]
+    pub fn partial_sort_unstable_by_key<K, F, R>(&mut self, range: R, mut f: F)
+    where
+        F: FnMut(&T) -> K,
+        K: Ord,
+        R: RangeBounds<usize>,
+    {
+        sort::unstable::partial_sort(self, range, |a, b| f(a).lt(&f(b)));
+    }
+
     /// Reorders the slice such that the element at `index` is at a sort-order position. All
     /// elements before `index` will be `<=` to this value, and all elements after will be `>=` to
     /// it.

library/core/src/slice/sort/unstable/mod.rs

@@ -1,11 +1,13 @@
 //! This module contains the entry points for `slice::sort_unstable`.
 
 use crate::mem::SizedTypeProperties;
+use crate::ops::{Range, RangeBounds};
+use crate::slice::sort::select::partition_at_index;
 #[cfg(not(any(feature = "optimize_for_size", target_pointer_width = "16")))]
 use crate::slice::sort::shared::find_existing_run;
 #[cfg(not(any(feature = "optimize_for_size", target_pointer_width = "16")))]
 use crate::slice::sort::shared::smallsort::insertion_sort_shift_left;
-use crate::{cfg_select, intrinsics};
+use crate::{cfg_select, intrinsics, slice};
 
 pub(crate) mod heapsort;
 pub(crate) mod quicksort;
@@ -17,7 +19,10 @@
 /// Upholds all safety properties outlined here:
 /// <https://github.com/Voultapher/sort-research-rs/blob/main/writeup/sort_safety/text.md>
 #[inline(always)]
-pub fn sort<T, F: FnMut(&T, &T) -> bool>(v: &mut [T], is_less: &mut F) {
+pub fn sort<T, F>(v: &mut [T], is_less: &mut F)
+where
+    F: FnMut(&T, &T) -> bool,
+{
     // Arrays of zero-sized types are always all-equal, and thus sorted.
     if T::IS_ZST {
         return;
@@ -52,6 +57,54 @@ pub fn sort<T, F: FnMut(&T, &T) -> bool>(v: &mut [T], is_less: &mut F) {
     }
 }
 
+/// Unstable partial sort the range `start..end`, after which it's guaranteed that:
+///
+/// 1. Every element in `v[..start]` is smaller than or equal to
+/// 2. Every element in `v[start..end]`, which is sorted, and smaller than or equal to
+/// 3. Every element in `v[end..]`.
+#[inline]
+pub fn partial_sort<T, F, R>(v: &mut [T], range: R, mut is_less: F)
+where
+    F: FnMut(&T, &T) -> bool,
+    R: RangeBounds<usize>,
+{
+    // Arrays of zero-sized types are always all-equal, and thus sorted.
+    if T::IS_ZST {
+        return;
+    }
+
+    let len = v.len();
+    let Range { start, end } = slice::range(range, ..len);
+
+    if end - start <= 1 {
+        // Empty range or single element. This case can be resolved in at most
+        // single partition_at_index call, without further sorting.
+
+        if end == 0 || start == len {
+            // Do nothing if it is an empty range at start or end: all guarantees
+            // are already upheld.
+            return;
+        }
+
+        partition_at_index(v, start, &mut is_less);
+        return;
+    }
+
+    // A heuristic factor to decide whether to partition the slice or not.
+    // If the range bound is close to the edges of the slice, it's not worth
+    // partitioning first.
+    const PARTITION_THRESHOLD: usize = 8;
+    let mut v = v;
+    if end + PARTITION_THRESHOLD <= len {
+        v = partition_at_index(v, end - 1, &mut is_less).0;
+    }
+    if start >= PARTITION_THRESHOLD {
+        v = partition_at_index(v, start, &mut is_less).2;
+    }
+
+    sort(v, &mut is_less);
+}
+
 /// See [`sort`]
 ///
 /// Deliberately don't inline the main sorting routine entrypoint to ensure the