Rollup merge of #144314 - kornelski:pivot-safely, r=jhpratt

Hint that choose_pivot returns index in bounds

Instead of using `unsafe` in multiple places, one `hint::assert_unchecked` allows use of safe code instead.

Part of #rust-lang/rust#144326

4 files changed, 10 insertions, 10 deletions

diff --git a/library/core/src/slice/sort/select.rs b/library/core/src/slice/sort/select.rs
index 82194db7fd8..fc31013caf8 100644
--- a/library/core/src/slice/sort/select.rs
+++ b/library/core/src/slice/sort/select.rs
@@ -101,8 +101,7 @@ fn partition_at_index_loop<'a, T, F>(
         // slice. Partition the slice into elements equal to and elements greater than the pivot.
         // This case is usually hit when the slice contains many duplicate elements.
         if let Some(p) = ancestor_pivot {
-            // SAFETY: choose_pivot promises to return a valid pivot position.
-            let pivot = unsafe { v.get_unchecked(pivot_pos) };
+            let pivot = &v[pivot_pos];
 
             if !is_less(p, pivot) {
                 let num_lt = partition(v, pivot_pos, &mut |a, b| !is_less(b, a));
diff --git a/library/core/src/slice/sort/shared/pivot.rs b/library/core/src/slice/sort/shared/pivot.rs
index 3aace484b6a..9eb60f854ce 100644
--- a/library/core/src/slice/sort/shared/pivot.rs
+++ b/library/core/src/slice/sort/shared/pivot.rs
@@ -1,6 +1,6 @@
 //! This module contains the logic for pivot selection.
 
-use crate::intrinsics;
+use crate::{hint, intrinsics};
 
 // Recursively select a pseudomedian if above this threshold.
 const PSEUDO_MEDIAN_REC_THRESHOLD: usize = 64;
@@ -9,6 +9,7 @@ const PSEUDO_MEDIAN_REC_THRESHOLD: usize = 64;
 ///
 /// This chooses a pivot by sampling an adaptive amount of points, approximating
 /// the quality of a median of sqrt(n) elements.
+#[inline]
 pub fn choose_pivot<T, F: FnMut(&T, &T) -> bool>(v: &[T], is_less: &mut F) -> usize {
     // We use unsafe code and raw pointers here because we're dealing with
     // heavy recursion. Passing safe slices around would involve a lot of
@@ -22,7 +23,7 @@ pub fn choose_pivot<T, F: FnMut(&T, &T) -> bool>(v: &[T], is_less: &mut F) -> us
     // SAFETY: a, b, c point to initialized regions of len_div_8 elements,
     // satisfying median3 and median3_rec's preconditions as v_base points
     // to an initialized region of n = len elements.
-    unsafe {
+    let index = unsafe {
         let v_base = v.as_ptr();
         let len_div_8 = len / 8;
 
@@ -35,6 +36,11 @@ pub fn choose_pivot<T, F: FnMut(&T, &T) -> bool>(v: &[T], is_less: &mut F) -> us
         } else {
             median3_rec(a, b, c, len_div_8, is_less).offset_from_unsigned(v_base)
         }
+    };
+    // SAFETY: preconditions must have been met for offset_from_unsigned()
+    unsafe {
+        hint::assert_unchecked(index < v.len());
+        index
     }
 }
 
diff --git a/library/core/src/slice/sort/stable/quicksort.rs b/library/core/src/slice/sort/stable/quicksort.rs
index 3c9688790c4..0439ba870bd 100644
--- a/library/core/src/slice/sort/stable/quicksort.rs
+++ b/library/core/src/slice/sort/stable/quicksort.rs
@@ -37,10 +37,6 @@ pub fn quicksort<T, F: FnMut(&T, &T) -> bool>(
         limit -= 1;
 
         let pivot_pos = choose_pivot(v, is_less);
-        // SAFETY: choose_pivot promises to return a valid pivot index.
-        unsafe {
-            intrinsics::assume(pivot_pos < v.len());
-        }
 
         // SAFETY: We only access the temporary copy for Freeze types, otherwise
         // self-modifications via `is_less` would not be observed and this would
diff --git a/library/core/src/slice/sort/unstable/quicksort.rs b/library/core/src/slice/sort/unstable/quicksort.rs
index 98efee242eb..bdf56a80803 100644
--- a/library/core/src/slice/sort/unstable/quicksort.rs
+++ b/library/core/src/slice/sort/unstable/quicksort.rs
@@ -48,8 +48,7 @@ pub(crate) fn quicksort<'a, T, F>(
         // slice. Partition the slice into elements equal to and elements greater than the pivot.
         // This case is usually hit when the slice contains many duplicate elements.
         if let Some(p) = ancestor_pivot {
-            // SAFETY: We assume choose_pivot yields an in-bounds position.
-            if !is_less(p, unsafe { v.get_unchecked(pivot_pos) }) {
+            if !is_less(p, &v[pivot_pos]) {
                 let num_lt = partition(v, pivot_pos, &mut |a, b| !is_less(b, a));
 
                 // Continue sorting elements greater than the pivot. We know that `num_lt` contains