Skip to content

Commit

Permalink
Fix off-by-one error causing driftsort to crash
Browse files Browse the repository at this point in the history
Fixes rust-lang#136103.
Based on the analysis by @jonathan-gruber-jg and @orlp.
  • Loading branch information
uellenberg authored and gitbot committed Feb 20, 2025
1 parent 7bfe9c9 commit c5f1867
Show file tree
Hide file tree
Showing 3 changed files with 22 additions and 8 deletions.
4 changes: 2 additions & 2 deletions core/src/slice/sort/stable/drift.rs
Original file line number Diff line number Diff line change
Expand Up @@ -10,8 +10,8 @@ use crate::{cmp, intrinsics};

/// Sorts `v` based on comparison function `is_less`. If `eager_sort` is true,
/// it will only do small-sorts and physical merges, ensuring O(N * log(N))
/// worst-case complexity. `scratch.len()` must be at least `max(v.len() / 2,
/// MIN_SMALL_SORT_SCRATCH_LEN)` otherwise the implementation may abort.
/// worst-case complexity. `scratch.len()` must be at least
/// `max(v.len() - v.len() / 2, SMALL_SORT_GENERAL_SCRATCH_LEN)` otherwise the implementation may abort.
/// Fully ascending and descending inputs will be sorted with exactly N - 1
/// comparisons.
///
Expand Down
24 changes: 18 additions & 6 deletions core/src/slice/sort/stable/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,8 @@ pub fn sort<T, F: FnMut(&T, &T) -> bool, BufT: BufGuard<T>>(v: &mut [T], is_less

cfg_if! {
if #[cfg(any(feature = "optimize_for_size", target_pointer_width = "16"))] {
// Unlike driftsort, mergesort only requires len / 2,
// not len - len / 2.
let alloc_len = len / 2;

cfg_if! {
Expand Down Expand Up @@ -91,16 +93,26 @@ fn driftsort_main<T, F: FnMut(&T, &T) -> bool, BufT: BufGuard<T>>(v: &mut [T], i
// By allocating n elements of memory we can ensure the entire input can
// be sorted using stable quicksort, which allows better performance on
// random and low-cardinality distributions. However, we still want to
// reduce our memory usage to n / 2 for large inputs. We do this by scaling
// our allocation as max(n / 2, min(n, 8MB)), ensuring we scale like n for
// small inputs and n / 2 for large inputs, without a sudden drop off. We
// also need to ensure our alloc >= MIN_SMALL_SORT_SCRATCH_LEN, as the
// reduce our memory usage to n - n / 2 for large inputs. We do this by scaling
// our allocation as max(n - n / 2, min(n, 8MB)), ensuring we scale like n for
// small inputs and n - n / 2 for large inputs, without a sudden drop off. We
// also need to ensure our alloc >= SMALL_SORT_GENERAL_SCRATCH_LEN, as the
// small-sort always needs this much memory.
//
// driftsort will produce unsorted runs of up to min_good_run_len, which
// is at most len - len / 2.
// Unsorted runs need to be processed by quicksort, which requires as much
// scratch space as the run length, therefore the scratch space must be at
// least len - len / 2.
// If min_good_run_len is ever modified, this code must be updated to allocate
// the correct scratch size for it.
const MAX_FULL_ALLOC_BYTES: usize = 8_000_000; // 8MB
let max_full_alloc = MAX_FULL_ALLOC_BYTES / mem::size_of::<T>();
let len = v.len();
let alloc_len =
cmp::max(cmp::max(len / 2, cmp::min(len, max_full_alloc)), SMALL_SORT_GENERAL_SCRATCH_LEN);
let alloc_len = cmp::max(
cmp::max(len - len / 2, cmp::min(len, max_full_alloc)),
SMALL_SORT_GENERAL_SCRATCH_LEN,
);

// For small inputs 4KiB of stack storage suffices, which allows us to avoid
// calling the (de-)allocator. Benchmarks showed this was quite beneficial.
Expand Down
2 changes: 2 additions & 0 deletions core/src/slice/sort/stable/quicksort.rs
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,8 @@ use crate::slice::sort::shared::smallsort::StableSmallSortTypeImpl;
use crate::{intrinsics, ptr};

/// Sorts `v` recursively using quicksort.
/// `scratch.len()` must be at least `max(v.len() - v.len() / 2, SMALL_SORT_GENERAL_SCRATCH_LEN)`
/// otherwise the implementation may abort.
///
/// `limit` when initialized with `c*log(v.len())` for some c ensures we do not
/// overflow the stack or go quadratic.
Expand Down

0 comments on commit c5f1867

Please sign in to comment.