core/slice/sort/unstable/quicksort.rs
1//! This module contains an unstable quicksort and two partition implementations.
2
3#[cfg(not(feature = "optimize_for_size"))]
4use crate::mem;
5use crate::mem::ManuallyDrop;
6#[cfg(not(feature = "optimize_for_size"))]
7use crate::slice::sort::shared::pivot::choose_pivot;
8#[cfg(not(feature = "optimize_for_size"))]
9use crate::slice::sort::shared::smallsort::UnstableSmallSortTypeImpl;
10#[cfg(not(feature = "optimize_for_size"))]
11use crate::slice::sort::unstable::heapsort;
12use crate::{intrinsics, ptr};
13
14/// Sorts `v` recursively.
15///
16/// If the slice had a predecessor in the original array, it is specified as `ancestor_pivot`.
17///
18/// `limit` is the number of allowed imbalanced partitions before switching to `heapsort`. If zero,
19/// this function will immediately switch to heapsort.
20#[cfg(not(feature = "optimize_for_size"))]
21pub(crate) fn quicksort<'a, T, F>(
22 mut v: &'a mut [T],
23 mut ancestor_pivot: Option<&'a T>,
24 mut limit: u32,
25 is_less: &mut F,
26) where
27 F: FnMut(&T, &T) -> bool,
28{
29 loop {
30 if v.len() <= T::small_sort_threshold() {
31 T::small_sort(v, is_less);
32 return;
33 }
34
35 // If too many bad pivot choices were made, simply fall back to heapsort in order to
36 // guarantee `O(N x log(N))` worst-case.
37 if limit == 0 {
38 heapsort::heapsort(v, is_less);
39 return;
40 }
41
42 limit -= 1;
43
44 // Choose a pivot and try guessing whether the slice is already sorted.
45 let pivot_pos = choose_pivot(v, is_less);
46
47 // If the chosen pivot is equal to the predecessor, then it's the smallest element in the
48 // slice. Partition the slice into elements equal to and elements greater than the pivot.
49 // This case is usually hit when the slice contains many duplicate elements.
50 if let Some(p) = ancestor_pivot {
51 // SAFETY: We assume choose_pivot yields an in-bounds position.
52 if !is_less(p, unsafe { v.get_unchecked(pivot_pos) }) {
53 let num_lt = partition(v, pivot_pos, &mut |a, b| !is_less(b, a));
54
55 // Continue sorting elements greater than the pivot. We know that `num_lt` contains
56 // the pivot. So we can continue after `num_lt`.
57 v = &mut v[(num_lt + 1)..];
58 ancestor_pivot = None;
59 continue;
60 }
61 }
62
63 // Partition the slice.
64 let num_lt = partition(v, pivot_pos, is_less);
65 // SAFETY: partition ensures that `num_lt` will be in-bounds.
66 unsafe { intrinsics::assume(num_lt < v.len()) };
67
68 // Split the slice into `left`, `pivot`, and `right`.
69 let (left, right) = v.split_at_mut(num_lt);
70 let (pivot, right) = right.split_at_mut(1);
71 let pivot = &pivot[0];
72
73 // Recurse into the left side. We have a fixed recursion limit, testing shows no real
74 // benefit for recursing into the shorter side.
75 quicksort(left, ancestor_pivot, limit, is_less);
76
77 // Continue with the right side.
78 v = right;
79 ancestor_pivot = Some(pivot);
80 }
81}
82
83/// Takes the input slice `v` and re-arranges elements such that when the call returns normally
84/// all elements that compare true for `is_less(elem, pivot)` where `pivot == v[pivot_pos]` are
85/// on the left side of `v` followed by the other elements, notionally considered greater or
86/// equal to `pivot`.
87///
88/// Returns the number of elements that are compared true for `is_less(elem, pivot)`.
89///
90/// If `is_less` does not implement a total order the resulting order and return value are
91/// unspecified. All original elements will remain in `v` and any possible modifications via
92/// interior mutability will be observable. Same is true if `is_less` panics or `v.len()`
93/// exceeds `scratch.len()`.
94pub(crate) fn partition<T, F>(v: &mut [T], pivot: usize, is_less: &mut F) -> usize
95where
96 F: FnMut(&T, &T) -> bool,
97{
98 let len = v.len();
99
100 // Allows for panic-free code-gen by proving this property to the compiler.
101 if len == 0 {
102 return 0;
103 }
104
105 if pivot >= len {
106 intrinsics::abort();
107 }
108
109 // SAFETY: We checked that `pivot` is in-bounds.
110 unsafe {
111 // Place the pivot at the beginning of slice.
112 v.swap_unchecked(0, pivot);
113 }
114 let (pivot, v_without_pivot) = v.split_at_mut(1);
115
116 // Assuming that Rust generates noalias LLVM IR we can be sure that a partition function
117 // signature of the form `(v: &mut [T], pivot: &T)` guarantees that pivot and v can't alias.
118 // Having this guarantee is crucial for optimizations. It's possible to copy the pivot value
119 // into a stack value, but this creates issues for types with interior mutability mandating
120 // a drop guard.
121 let pivot = &mut pivot[0];
122
123 // This construct is used to limit the LLVM IR generated, which saves large amounts of
124 // compile-time by only instantiating the code that is needed. Idea by Frank Steffahn.
125 let num_lt = (const { inst_partition::<T, F>() })(v_without_pivot, pivot, is_less);
126
127 if num_lt >= len {
128 intrinsics::abort();
129 }
130
131 // SAFETY: We checked that `num_lt` is in-bounds.
132 unsafe {
133 // Place the pivot between the two partitions.
134 v.swap_unchecked(0, num_lt);
135 }
136
137 num_lt
138}
139
140const fn inst_partition<T, F: FnMut(&T, &T) -> bool>() -> fn(&mut [T], &T, &mut F) -> usize {
141 const MAX_BRANCHLESS_PARTITION_SIZE: usize = 96;
142 if size_of::<T>() <= MAX_BRANCHLESS_PARTITION_SIZE {
143 // Specialize for types that are relatively cheap to copy, where branchless optimizations
144 // have large leverage e.g. `u64` and `String`.
145 cfg_if! {
146 if #[cfg(feature = "optimize_for_size")] {
147 partition_lomuto_branchless_simple::<T, F>
148 } else {
149 partition_lomuto_branchless_cyclic::<T, F>
150 }
151 }
152 } else {
153 partition_hoare_branchy_cyclic::<T, F>
154 }
155}
156
157/// See [`partition`].
158fn partition_hoare_branchy_cyclic<T, F>(v: &mut [T], pivot: &T, is_less: &mut F) -> usize
159where
160 F: FnMut(&T, &T) -> bool,
161{
162 let len = v.len();
163
164 if len == 0 {
165 return 0;
166 }
167
168 // Optimized for large types that are expensive to move. Not optimized for integers. Optimized
169 // for small code-gen, assuming that is_less is an expensive operation that generates
170 // substantial amounts of code or a call. And that copying elements will likely be a call to
171 // memcpy. Using 2 `ptr::copy_nonoverlapping` has the chance to be faster than
172 // `ptr::swap_nonoverlapping` because `memcpy` can use wide SIMD based on runtime feature
173 // detection. Benchmarks support this analysis.
174
175 let mut gap_opt: Option<GapGuard<T>> = None;
176
177 // SAFETY: The left-to-right scanning loop performs a bounds check, where we know that `left >=
178 // v_base && left < right && right <= v_base.add(len)`. The right-to-left scanning loop performs
179 // a bounds check ensuring that `right` is in-bounds. We checked that `len` is more than zero,
180 // which means that unconditional `right = right.sub(1)` is safe to do. The exit check makes
181 // sure that `left` and `right` never alias, making `ptr::copy_nonoverlapping` safe. The
182 // drop-guard `gap` ensures that should `is_less` panic we always overwrite the duplicate in the
183 // input. `gap.pos` stores the previous value of `right` and starts at `right` and so it too is
184 // in-bounds. We never pass the saved `gap.value` to `is_less` while it is inside the `GapGuard`
185 // thus any changes via interior mutability will be observed.
186 unsafe {
187 let v_base = v.as_mut_ptr();
188
189 let mut left = v_base;
190 let mut right = v_base.add(len);
191
192 loop {
193 // Find the first element greater than the pivot.
194 while left < right && is_less(&*left, pivot) {
195 left = left.add(1);
196 }
197
198 // Find the last element equal to the pivot.
199 loop {
200 right = right.sub(1);
201 if left >= right || is_less(&*right, pivot) {
202 break;
203 }
204 }
205
206 if left >= right {
207 break;
208 }
209
210 // Swap the found pair of out-of-order elements via cyclic permutation.
211 let is_first_swap_pair = gap_opt.is_none();
212
213 if is_first_swap_pair {
214 gap_opt = Some(GapGuard { pos: right, value: ManuallyDrop::new(ptr::read(left)) });
215 }
216
217 let gap = gap_opt.as_mut().unwrap_unchecked();
218
219 // Single place where we instantiate ptr::copy_nonoverlapping in the partition.
220 if !is_first_swap_pair {
221 ptr::copy_nonoverlapping(left, gap.pos, 1);
222 }
223 gap.pos = right;
224 ptr::copy_nonoverlapping(right, left, 1);
225
226 left = left.add(1);
227 }
228
229 left.offset_from_unsigned(v_base)
230
231 // `gap_opt` goes out of scope and overwrites the last wrong-side element on the right side
232 // with the first wrong-side element of the left side that was initially overwritten by the
233 // first wrong-side element on the right side element.
234 }
235}
236
237#[cfg(not(feature = "optimize_for_size"))]
238struct PartitionState<T> {
239 // The current element that is being looked at, scans left to right through slice.
240 right: *mut T,
241 // Counts the number of elements that compared less-than, also works around:
242 // https://github.com/rust-lang/rust/issues/117128
243 num_lt: usize,
244 // Gap guard that tracks the temporary duplicate in the input.
245 gap: GapGuardRaw<T>,
246}
247
248#[cfg(not(feature = "optimize_for_size"))]
249fn partition_lomuto_branchless_cyclic<T, F>(v: &mut [T], pivot: &T, is_less: &mut F) -> usize
250where
251 F: FnMut(&T, &T) -> bool,
252{
253 // Novel partition implementation by Lukas Bergdoll and Orson Peters. Branchless Lomuto
254 // partition paired with a cyclic permutation.
255 // https://github.com/Voultapher/sort-research-rs/blob/main/writeup/lomcyc_partition/text.md
256
257 let len = v.len();
258 let v_base = v.as_mut_ptr();
259
260 if len == 0 {
261 return 0;
262 }
263
264 // SAFETY: We checked that `len` is more than zero, which means that reading `v_base` is safe to
265 // do. From there we have a bounded loop where `v_base.add(i)` is guaranteed in-bounds. `v` and
266 // `pivot` can't alias because of type system rules. The drop-guard `gap` ensures that should
267 // `is_less` panic we always overwrite the duplicate in the input. `gap.pos` stores the previous
268 // value of `right` and starts at `v_base` and so it too is in-bounds. Given `UNROLL_LEN == 2`
269 // after the main loop we either have A) the last element in `v` that has not yet been processed
270 // because `len % 2 != 0`, or B) all elements have been processed except the gap value that was
271 // saved at the beginning with `ptr::read(v_base)`. In the case A) the loop will iterate twice,
272 // first performing loop_body to take care of the last element that didn't fit into the unroll.
273 // After that the behavior is the same as for B) where we use the saved value as `right` to
274 // overwrite the duplicate. If this very last call to `is_less` panics the saved value will be
275 // copied back including all possible changes via interior mutability. If `is_less` does not
276 // panic and the code continues we overwrite the duplicate and do `right = right.add(1)`, this
277 // is safe to do with `&mut *gap.value` because `T` is the same as `[T; 1]` and generating a
278 // pointer one past the allocation is safe.
279 unsafe {
280 let mut loop_body = |state: &mut PartitionState<T>| {
281 let right_is_lt = is_less(&*state.right, pivot);
282 let left = v_base.add(state.num_lt);
283
284 ptr::copy(left, state.gap.pos, 1);
285 ptr::copy_nonoverlapping(state.right, left, 1);
286
287 state.gap.pos = state.right;
288 state.num_lt += right_is_lt as usize;
289
290 state.right = state.right.add(1);
291 };
292
293 // Ideally we could just use GapGuard in PartitionState, but the reference that is
294 // materialized with `&mut state` when calling `loop_body` would create a mutable reference
295 // to the parent struct that contains the gap value, invalidating the reference pointer
296 // created from a reference to the gap value in the cleanup loop. This is only an issue
297 // under Stacked Borrows, Tree Borrows accepts the intuitive code using GapGuard as valid.
298 let mut gap_value = ManuallyDrop::new(ptr::read(v_base));
299
300 let mut state = PartitionState {
301 num_lt: 0,
302 right: v_base.add(1),
303
304 gap: GapGuardRaw { pos: v_base, value: &mut *gap_value },
305 };
306
307 // Manual unrolling that works well on x86, Arm and with opt-level=s without murdering
308 // compile-times. Leaving this to the compiler yields ok to bad results.
309 let unroll_len = const { if size_of::<T>() <= 16 { 2 } else { 1 } };
310
311 let unroll_end = v_base.add(len - (unroll_len - 1));
312 while state.right < unroll_end {
313 if unroll_len == 2 {
314 loop_body(&mut state);
315 loop_body(&mut state);
316 } else {
317 loop_body(&mut state);
318 }
319 }
320
321 // Single instantiate `loop_body` for both the unroll cleanup and cyclic permutation
322 // cleanup. Optimizes binary-size and compile-time.
323 let end = v_base.add(len);
324 loop {
325 let is_done = state.right == end;
326 state.right = if is_done { state.gap.value } else { state.right };
327
328 loop_body(&mut state);
329
330 if is_done {
331 mem::forget(state.gap);
332 break;
333 }
334 }
335
336 state.num_lt
337 }
338}
339
340#[cfg(feature = "optimize_for_size")]
341fn partition_lomuto_branchless_simple<T, F: FnMut(&T, &T) -> bool>(
342 v: &mut [T],
343 pivot: &T,
344 is_less: &mut F,
345) -> usize {
346 let mut left = 0;
347
348 for right in 0..v.len() {
349 // SAFETY: `left` can at max be incremented by 1 each loop iteration, which implies that
350 // left <= right and that both are in-bounds.
351 unsafe {
352 let right_is_lt = is_less(v.get_unchecked(right), pivot);
353 v.swap_unchecked(left, right);
354 left += right_is_lt as usize;
355 }
356 }
357
358 left
359}
360
361struct GapGuard<T> {
362 pos: *mut T,
363 value: ManuallyDrop<T>,
364}
365
366impl<T> Drop for GapGuard<T> {
367 fn drop(&mut self) {
368 // SAFETY: `self` MUST be constructed in a way that makes copying the gap value into
369 // `self.pos` sound.
370 unsafe {
371 ptr::copy_nonoverlapping(&*self.value, self.pos, 1);
372 }
373 }
374}
375
376/// Ideally this wouldn't be needed and we could just use the regular GapGuard.
377/// See comment in [`partition_lomuto_branchless_cyclic`].
378#[cfg(not(feature = "optimize_for_size"))]
379struct GapGuardRaw<T> {
380 pos: *mut T,
381 value: *mut T,
382}
383
384#[cfg(not(feature = "optimize_for_size"))]
385impl<T> Drop for GapGuardRaw<T> {
386 fn drop(&mut self) {
387 // SAFETY: `self` MUST be constructed in a way that makes copying the gap value into
388 // `self.pos` sound.
389 unsafe {
390 ptr::copy_nonoverlapping(self.value, self.pos, 1);
391 }
392 }
393}