From 15652544786b6787af29cae9bc0bb16d49d48fb4 Mon Sep 17 00:00:00 2001
From: uellenberg <JonahUellenberg@hotmail.com>
Date: Mon, 27 Jan 2025 16:21:08 -0800
Subject: [PATCH] Fix off-by-one error causing driftsort to crash
Fixes #136103.
Based on the analysis by @jonathan-gruber-jg and @orlp.
---
library/core/src/slice/sort/stable/drift.rs | 4 ++--
library/core/src/slice/sort/stable/mod.rs | 24 ++++++++++++++-----
.../core/src/slice/sort/stable/quicksort.rs | 2 ++
.../driftsort-off-by-one-issue-136103.rs | 10 ++++++++
4 files changed, 32 insertions(+), 8 deletions(-)
create mode 100644 tests/ui/array-slice-vec/driftsort-off-by-one-issue-136103.rs
diff --git a/library/core/src/slice/sort/stable/drift.rs b/library/core/src/slice/sort/stable/drift.rs
index 644e75a4581e9..cf1df1e91a50d 100644
--- a/library/core/src/slice/sort/stable/drift.rs
+++ b/library/core/src/slice/sort/stable/drift.rs
@@ -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.
///
diff --git a/library/core/src/slice/sort/stable/mod.rs b/library/core/src/slice/sort/stable/mod.rs
index 7adcc83b818d1..3ff2e71fd05bc 100644
--- a/library/core/src/slice/sort/stable/mod.rs
+++ b/library/core/src/slice/sort/stable/mod.rs
@@ -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! {
@@ -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.
diff --git a/library/core/src/slice/sort/stable/quicksort.rs b/library/core/src/slice/sort/stable/quicksort.rs
index 0c8308bfce00e..630c6ff907703 100644
--- a/library/core/src/slice/sort/stable/quicksort.rs
+++ b/library/core/src/slice/sort/stable/quicksort.rs
@@ -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.
diff --git a/tests/ui/array-slice-vec/driftsort-off-by-one-issue-136103.rs b/tests/ui/array-slice-vec/driftsort-off-by-one-issue-136103.rs
new file mode 100644
index 0000000000000..42197ff102d72
--- /dev/null
+++ b/tests/ui/array-slice-vec/driftsort-off-by-one-issue-136103.rs
@@ -0,0 +1,10 @@
+//@ run-pass
+// Ensures that driftsort doesn't crash under specific slice
+// length and memory size.
+// Based on the example given in https://github.com/rust-lang/rust/issues/136103.
+fn main() {
+ let n = 127;
+ let mut objs: Vec<_> =
+ (0..n).map(|i| [(i % 2) as u8; 125001]).collect();
+ objs.sort();
+}