Fix off-by-one error causing driftsort to crash

Fixes #136103.
Based on the analysis by @jonathan-gruber-jg and @orlp.

Author: uellenberg

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, MIN_SMALL_SORT_SCRATCH_LEN)` otherwise the implementation may abort.
 /// Fully ascending and descending inputs will be sorted with exactly N - 1
 /// comparisons.
 ///

library/core/src/slice/sort/stable/mod.rs

@@ -91,16 +91,18 @@ 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
+    // 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 >= MIN_SMALL_SORT_SCRATCH_LEN, as the
     // small-sort always needs this much memory.
     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.

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, MIN_SMALL_SORT_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.

tests/ui/array-slice-vec/driftsort-off-by-one-issue-136103.rs

@@ -0,0 +1,53 @@
+//@ 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.
+use std::cmp::Ordering;
+use std::hint::black_box;
+
+#[derive(Eq)]
+struct S(bool, [u8; 999_999]);
+
+// Don't compare the large array to make
+// the test run faster.
+
+impl PartialOrd for S {
+    #[inline]
+    fn partial_cmp(&self, other: &S) -> Option<Ordering> {
+        self.0.partial_cmp(&other.0)
+    }
+}
+
+impl Ord for S {
+    #[inline]
+    fn cmp(&self, other: &S) -> Ordering {
+        self.0.cmp(&other.0)
+    }
+}
+
+impl PartialEq for S {
+    #[inline]
+    fn eq(&self, other: &S) -> bool {
+        self.0 == other.0
+    }
+}
+
+fn main() {
+    let n = 101;
+    let mut objs = Vec::with_capacity(n);
+
+    let mut on = false;
+
+    for _ in 0..n {
+        objs.push(S(on, [0; 999_999]));
+        on = !on;
+    }
+
+    objs.sort();
+
+    // Ensure that the large array isn't eliminated,
+    // as that's needed to trigger the bug.
+    black_box(&objs[0].1);
+
+    println!("Succeeded");
+}