Skip to content

Hint that choose_pivot returns index in bounds #144314

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Jul 25, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 1 addition & 2 deletions library/core/src/slice/sort/select.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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));
Expand Down
10 changes: 8 additions & 2 deletions library/core/src/slice/sort/shared/pivot.rs
Original file line number Diff line number Diff line change
@@ -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;
Expand All @@ -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
Expand All @@ -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;

Expand All @@ -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
}
}

Expand Down
4 changes: 0 additions & 4 deletions library/core/src/slice/sort/stable/quicksort.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
3 changes: 1 addition & 2 deletions library/core/src/slice/sort/unstable/quicksort.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Loading