Safer sort partition

Author: kornelski

library/core/src/slice/sort/select.rs

@@ -10,7 +10,7 @@ use crate::cfg_select;
 use crate::mem::{self, SizedTypeProperties};
 #[cfg(not(feature = "optimize_for_size"))]
 use crate::slice::sort::shared::pivot::choose_pivot;
-use crate::slice::sort::shared::smallsort::insertion_sort_shift_left;
+use crate::slice::sort::shared::smallsort::{insertion_sort_shift_left, panic_on_ord_violation};
 use crate::slice::sort::unstable::quicksort::partition;
 
 /// Reorders the slice such that the element at `index` is at its final sorted position.
@@ -104,7 +104,9 @@ fn partition_at_index_loop<'a, T, F>(
             let pivot = &v[pivot_pos];
 
             if !is_less(p, pivot) {
-                let num_lt = partition(v, pivot_pos, &mut |a, b| !is_less(b, a));
+                let Some(num_lt) = partition(v, pivot_pos, &mut |a, b| !is_less(b, a)) else {
+                    panic_on_ord_violation();
+                };
 
                 // Continue sorting elements greater than the pivot. We know that `mid` contains
                 // the pivot. So we can continue after `mid`.
@@ -122,14 +124,15 @@ fn partition_at_index_loop<'a, T, F>(
             }
         }
 
-        let mid = partition(v, pivot_pos, is_less);
+        let Some(mid) = partition(v, pivot_pos, is_less) else {
+            panic_on_ord_violation();
+        };
 
         // Split the slice into `left`, `pivot`, and `right`.
         let (left, right) = v.split_at_mut(mid);
-        let (pivot, right) = right.split_at_mut(1);
-        let pivot = &pivot[0];
 
         if mid < index {
+            let (pivot, right) = right.split_first_mut().unwrap();
             v = right;
             index = index - mid - 1;
             ancestor_pivot = Some(pivot);
@@ -198,7 +201,9 @@ fn median_of_medians<T, F: FnMut(&T, &T) -> bool>(mut v: &mut [T], is_less: &mut
             return;
         }
 
-        let p = median_of_ninthers(v, is_less);
+        let Some(p) = median_of_ninthers(v, is_less) else {
+            panic_on_ord_violation();
+        };
 
         if p == k {
             return;
@@ -216,7 +221,7 @@ fn median_of_medians<T, F: FnMut(&T, &T) -> bool>(mut v: &mut [T], is_less: &mut
 // Optimized for when `k` lies somewhere in the middle of the slice. Selects a pivot
 // as close as possible to the median of the slice. For more details on how the algorithm
 // operates, refer to the paper <https://drops.dagstuhl.de/opus/volltexte/2017/7612/pdf/LIPIcs-SEA-2017-24.pdf>.
-fn median_of_ninthers<T, F: FnMut(&T, &T) -> bool>(v: &mut [T], is_less: &mut F) -> usize {
+fn median_of_ninthers<T, F: FnMut(&T, &T) -> bool>(v: &mut [T], is_less: &mut F) -> Option<usize> {
     // use `saturating_mul` so the multiplication doesn't overflow on 16-bit platforms.
     let frac = if v.len() <= 1024 {
         v.len() / 12

library/core/src/slice/sort/shared/smallsort.rs

@@ -842,7 +842,7 @@ unsafe fn bidirectional_merge<T: FreezeMarker, F: FnMut(&T, &T) -> bool>(
 
 #[cfg_attr(not(panic = "immediate-abort"), inline(never), cold)]
 #[cfg_attr(panic = "immediate-abort", inline)]
-fn panic_on_ord_violation() -> ! {
+pub(crate) fn panic_on_ord_violation() -> ! {
     // This is indicative of a logic bug in the user-provided comparison function or Ord
     // implementation. They are expected to implement a total order as explained in the Ord
     // documentation.

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

@@ -6,10 +6,10 @@ use crate::mem::ManuallyDrop;
 #[cfg(not(feature = "optimize_for_size"))]
 use crate::slice::sort::shared::pivot::choose_pivot;
 #[cfg(not(feature = "optimize_for_size"))]
-use crate::slice::sort::shared::smallsort::UnstableSmallSortTypeImpl;
+use crate::slice::sort::shared::smallsort::{UnstableSmallSortTypeImpl, panic_on_ord_violation};
 #[cfg(not(feature = "optimize_for_size"))]
 use crate::slice::sort::unstable::heapsort;
-use crate::{cfg_select, intrinsics, ptr};
+use crate::{cfg_select, ptr};
 
 /// Sorts `v` recursively.
 ///
@@ -53,16 +53,18 @@ pub(crate) fn quicksort<'a, T, F>(
 
                 // Continue sorting elements greater than the pivot. We know that `num_lt` contains
                 // the pivot. So we can continue after `num_lt`.
-                v = &mut v[(num_lt + 1)..];
+                v = num_lt
+                    .and_then(|num_lt| v.get_mut(num_lt + 1..))
+                    .unwrap_or_else(|| panic_on_ord_violation());
                 ancestor_pivot = None;
                 continue;
             }
         }
 
         // Partition the slice.
-        let num_lt = partition(v, pivot_pos, is_less);
-        // SAFETY: partition ensures that `num_lt` will be in-bounds.
-        unsafe { intrinsics::assume(num_lt < v.len()) };
+        let Some(num_lt) = partition(v, pivot_pos, is_less) else {
+            panic_on_ord_violation();
+        };
 
         // Split the slice into `left`, `pivot`, and `right`.
         let (left, right) = v.split_at_mut(num_lt);
@@ -84,56 +86,46 @@ pub(crate) fn quicksort<'a, T, F>(
 /// on the left side of `v` followed by the other elements, notionally considered greater or
 /// equal to `pivot`.
 ///
-/// Returns the number of elements that are compared true for `is_less(elem, pivot)`.
+/// Returns the number of elements that are compared true for `is_less(elem, pivot)`
+/// if `is_less` implements a total order.
 ///
 /// If `is_less` does not implement a total order the resulting order and return value are
-/// unspecified. All original elements will remain in `v` and any possible modifications via
+/// unspecified, except that if `Some` is returned, the value will be in bounds of `v`.
+/// All original elements will remain in `v` and any possible modifications via
 /// interior mutability will be observable. Same is true if `is_less` panics or `v.len()`
 /// exceeds `scratch.len()`.
-pub(crate) fn partition<T, F>(v: &mut [T], pivot: usize, is_less: &mut F) -> usize
+pub(crate) fn partition<T, F>(v: &mut [T], pivot: usize, is_less: &mut F) -> Option<usize>
 where
     F: FnMut(&T, &T) -> bool,
 {
     let len = v.len();
 
     // Allows for panic-free code-gen by proving this property to the compiler.
-    if len == 0 {
-        return 0;
+    if len == 0 || pivot >= len {
+        return None;
     }
 
-    if pivot >= len {
-        intrinsics::abort();
-    }
-
-    // SAFETY: We checked that `pivot` is in-bounds.
-    unsafe {
-        // Place the pivot at the beginning of slice.
-        v.swap_unchecked(0, pivot);
-    }
-    let (pivot, v_without_pivot) = v.split_at_mut(1);
+    v.swap(0, pivot);
+    let (pivot, v_without_pivot) = v.split_first_mut()?;
 
     // Assuming that Rust generates noalias LLVM IR we can be sure that a partition function
     // signature of the form `(v: &mut [T], pivot: &T)` guarantees that pivot and v can't alias.
     // Having this guarantee is crucial for optimizations. It's possible to copy the pivot value
     // into a stack value, but this creates issues for types with interior mutability mandating
     // a drop guard.
-    let pivot = &mut pivot[0];
 
     // This construct is used to limit the LLVM IR generated, which saves large amounts of
     // compile-time by only instantiating the code that is needed. Idea by Frank Steffahn.
     let num_lt = (const { inst_partition::<T, F>() })(v_without_pivot, pivot, is_less);
 
     if num_lt >= len {
-        intrinsics::abort();
+        return None;
     }
 
-    // SAFETY: We checked that `num_lt` is in-bounds.
-    unsafe {
-        // Place the pivot between the two partitions.
-        v.swap_unchecked(0, num_lt);
-    }
+    // Place the pivot between the two partitions.
+    v.swap(0, num_lt);
 
-    num_lt
+    Some(num_lt)
 }
 
 const fn inst_partition<T, F: FnMut(&T, &T) -> bool>() -> fn(&mut [T], &T, &mut F) -> usize {