1//! This module contains the implementation for `slice::select_nth_unstable`.
2//! It uses an introselect algorithm based on ipnsort by Lukas Bergdoll and Orson Peters,
3//! published at: <https://github.com/Voultapher/sort-research-rs/tree/main/ipnsort>
4//!
5//! The fallback algorithm used for introselect is Median of Medians using Tukey's Ninther
6//! for pivot selection. Using this as a fallback ensures O(n) worst case running time with
7//! better performance than one would get using heapsort as fallback.
8
9use crate::cfg_select;
10use crate::mem::{self, SizedTypeProperties};
11#[cfg(not(feature = "optimize_for_size"))]
12use crate::slice::sort::shared::pivot::choose_pivot;
13use crate::slice::sort::shared::smallsort::insertion_sort_shift_left;
14use crate::slice::sort::unstable::quicksort::partition;
15
16/// Reorders the slice such that the element at `index` is at its final sorted position.
17pub(crate) fn partition_at_index<T, F>(
18 v: &mut [T],
19 index: usize,
20 mut is_less: F,
21) -> (&mut [T], &mut T, &mut [T])
22where
23 F: FnMut(&T, &T) -> bool,
24{
25 let len = v.len();
26
27 // Puts a lower limit of 1 on `len`.
28 if index >= len {
29 panic!("partition_at_index index {} greater than length of slice {}", index, len);
30 }
31
32 if T::IS_ZST {
33 // Sorting has no meaningful behavior on zero-sized types. Do nothing.
34 } else if index == len - 1 {
35 // Find max element and place it in the last position of the array. We're free to use
36 // `unwrap()` here because we checked that `v` is not empty.
37 let max_idx = max_index(v, &mut is_less).unwrap();
38 v.swap(max_idx, index);
39 } else if index == 0 {
40 // Find min element and place it in the first position of the array. We're free to use
41 // `unwrap()` here because we checked that `v` is not empty.
42 let min_idx = min_index(v, &mut is_less).unwrap();
43 v.swap(min_idx, index);
44 } else {
45 cfg_select! {
46 feature = "optimize_for_size" => {
47 median_of_medians(v, &mut is_less, index);
48 }
49 _ => {
50 partition_at_index_loop(v, index, None, &mut is_less);
51 }
52 }
53 }
54
55 let (left, right) = v.split_at_mut(index);
56 let (pivot, right) = right.split_at_mut(1);
57 let pivot = &mut pivot[0];
58 (left, pivot, right)
59}
60
61// For small sub-slices it's faster to use a dedicated small-sort, but because it is only called at
62// most once, it doesn't make sense to use something more sophisticated than insertion-sort.
63const INSERTION_SORT_THRESHOLD: usize = 16;
64
65#[cfg(not(feature = "optimize_for_size"))]
66fn partition_at_index_loop<'a, T, F>(
67 mut v: &'a mut [T],
68 mut index: usize,
69 mut ancestor_pivot: Option<&'a T>,
70 is_less: &mut F,
71) where
72 F: FnMut(&T, &T) -> bool,
73{
74 // Limit the amount of iterations and fall back to fast deterministic selection to ensure O(n)
75 // worst case running time. This limit needs to be constant, because using `ilog2(len)` like in
76 // `sort` would result in O(n log n) time complexity. The exact value of the limit is chosen
77 // somewhat arbitrarily, but for most inputs bad pivot selections should be relatively rare, so
78 // the limit is reached for sub-slices len / (2^limit or less). Which makes the remaining work
79 // with the fallback minimal in relative terms.
80 let mut limit = 16;
81
82 loop {
83 if v.len() <= INSERTION_SORT_THRESHOLD {
84 if v.len() >= 2 {
85 insertion_sort_shift_left(v, 1, is_less);
86 }
87 return;
88 }
89
90 if limit == 0 {
91 median_of_medians(v, is_less, index);
92 return;
93 }
94
95 limit -= 1;
96
97 // Choose a pivot
98 let pivot_pos = choose_pivot(v, is_less);
99
100 // If the chosen pivot is equal to the predecessor, then it's the smallest element in the
101 // slice. Partition the slice into elements equal to and elements greater than the pivot.
102 // This case is usually hit when the slice contains many duplicate elements.
103 if let Some(p) = ancestor_pivot {
104 // SAFETY: choose_pivot promises to return a valid pivot position.
105 let pivot = unsafe { v.get_unchecked(pivot_pos) };
106
107 if !is_less(p, pivot) {
108 let num_lt = partition(v, pivot_pos, &mut |a, b| !is_less(b, a));
109
110 // Continue sorting elements greater than the pivot. We know that `mid` contains
111 // the pivot. So we can continue after `mid`.
112 let mid = num_lt + 1;
113
114 // If we've passed our index, then we're good.
115 if mid > index {
116 return;
117 }
118
119 v = &mut v[mid..];
120 index = index - mid;
121 ancestor_pivot = None;
122 continue;
123 }
124 }
125
126 let mid = partition(v, pivot_pos, is_less);
127
128 // Split the slice into `left`, `pivot`, and `right`.
129 let (left, right) = v.split_at_mut(mid);
130 let (pivot, right) = right.split_at_mut(1);
131 let pivot = &pivot[0];
132
133 if mid < index {
134 v = right;
135 index = index - mid - 1;
136 ancestor_pivot = Some(pivot);
137 } else if mid > index {
138 v = left;
139 } else {
140 // If mid == index, then we're done, since partition() guaranteed that all elements
141 // after mid are greater than or equal to mid.
142 return;
143 }
144 }
145}
146
147/// Helper function that returns the index of the minimum element in the slice using the given
148/// comparator function
149fn min_index<T, F: FnMut(&T, &T) -> bool>(slice: &[T], is_less: &mut F) -> Option<usize> {
150 sliceOption<(usize, &T)>
151 .iter()
152 .enumerate()
153 .reduce(|acc: (usize, &T), t: (usize, &T)| if is_less(t.1, acc.1) { t } else { acc })
154 .map(|(i: usize, _)| i)
155}
156
157/// Helper function that returns the index of the maximum element in the slice using the given
158/// comparator function
159fn max_index<T, F: FnMut(&T, &T) -> bool>(slice: &[T], is_less: &mut F) -> Option<usize> {
160 sliceOption<(usize, &T)>
161 .iter()
162 .enumerate()
163 .reduce(|acc: (usize, &T), t: (usize, &T)| if is_less(acc.1, t.1) { t } else { acc })
164 .map(|(i: usize, _)| i)
165}
166
167/// Selection algorithm to select the k-th element from the slice in guaranteed O(n) time.
168/// This is essentially a quickselect that uses Tukey's Ninther for pivot selection
169fn median_of_medians<T, F: FnMut(&T, &T) -> bool>(mut v: &mut [T], is_less: &mut F, mut k: usize) {
170 // Since this function isn't public, it should never be called with an out-of-bounds index.
171 debug_assert!(k < v.len());
172
173 // If T is as ZST, `partition_at_index` will already return early.
174 debug_assert!(!T::IS_ZST);
175
176 // We now know that `k < v.len() <= isize::MAX`
177 loop {
178 if v.len() <= INSERTION_SORT_THRESHOLD {
179 if v.len() >= 2 {
180 insertion_sort_shift_left(v, 1, is_less);
181 }
182
183 return;
184 }
185
186 // `median_of_{minima,maxima}` can't handle the extreme cases of the first/last element,
187 // so we catch them here and just do a linear search.
188 if k == v.len() - 1 {
189 // Find max element and place it in the last position of the array. We're free to use
190 // `unwrap()` here because we know v must not be empty.
191 let max_idx = max_index(v, is_less).unwrap();
192 v.swap(max_idx, k);
193 return;
194 } else if k == 0 {
195 // Find min element and place it in the first position of the array. We're free to use
196 // `unwrap()` here because we know v must not be empty.
197 let min_idx = min_index(v, is_less).unwrap();
198 v.swap(min_idx, k);
199 return;
200 }
201
202 let p = median_of_ninthers(v, is_less);
203
204 if p == k {
205 return;
206 } else if p > k {
207 v = &mut v[..p];
208 } else {
209 // Since `p < k < v.len()`, `p + 1` doesn't overflow and is
210 // a valid index into the slice.
211 v = &mut v[p + 1..];
212 k -= p + 1;
213 }
214 }
215}
216
217// Optimized for when `k` lies somewhere in the middle of the slice. Selects a pivot
218// as close as possible to the median of the slice. For more details on how the algorithm
219// operates, refer to the paper <https://drops.dagstuhl.de/opus/volltexte/2017/7612/pdf/LIPIcs-SEA-2017-24.pdf>.
220fn median_of_ninthers<T, F: FnMut(&T, &T) -> bool>(v: &mut [T], is_less: &mut F) -> usize {
221 // use `saturating_mul` so the multiplication doesn't overflow on 16-bit platforms.
222 let frac = if v.len() <= 1024 {
223 v.len() / 12
224 } else if v.len() <= 128_usize.saturating_mul(1024) {
225 v.len() / 64
226 } else {
227 v.len() / 1024
228 };
229
230 let pivot = frac / 2;
231 let lo = v.len() / 2 - pivot;
232 let hi = frac + lo;
233 let gap = (v.len() - 9 * frac) / 4;
234 let mut a = lo - 4 * frac - gap;
235 let mut b = hi + gap;
236 for i in lo..hi {
237 ninther(v, is_less, a, i - frac, b, a + 1, i, b + 1, a + 2, i + frac, b + 2);
238 a += 3;
239 b += 3;
240 }
241
242 median_of_medians(&mut v[lo..lo + frac], is_less, pivot);
243
244 partition(v, lo + pivot, is_less)
245}
246
247/// Moves around the 9 elements at the indices a..i, such that
248/// `v[d]` contains the median of the 9 elements and the other
249/// elements are partitioned around it.
250fn ninther<T, F: FnMut(&T, &T) -> bool>(
251 v: &mut [T],
252 is_less: &mut F,
253 a: usize,
254 mut b: usize,
255 c: usize,
256 mut d: usize,
257 e: usize,
258 mut f: usize,
259 g: usize,
260 mut h: usize,
261 i: usize,
262) {
263 b = median_idx(v, is_less, a, b, c);
264 h = median_idx(v, is_less, g, h, i);
265 if is_less(&v[h], &v[b]) {
266 mem::swap(&mut b, &mut h);
267 }
268 if is_less(&v[f], &v[d]) {
269 mem::swap(&mut d, &mut f);
270 }
271 if is_less(&v[e], &v[d]) {
272 // do nothing
273 } else if is_less(&v[f], &v[e]) {
274 d = f;
275 } else {
276 if is_less(&v[e], &v[b]) {
277 v.swap(e, b);
278 } else if is_less(&v[h], &v[e]) {
279 v.swap(e, h);
280 }
281 return;
282 }
283 if is_less(&v[d], &v[b]) {
284 d = b;
285 } else if is_less(&v[h], &v[d]) {
286 d = h;
287 }
288
289 v.swap(d, e);
290}
291
292/// returns the index pointing to the median of the 3
293/// elements `v[a]`, `v[b]` and `v[c]`
294fn median_idx<T, F: FnMut(&T, &T) -> bool>(
295 v: &[T],
296 is_less: &mut F,
297 mut a: usize,
298 b: usize,
299 mut c: usize,
300) -> usize {
301 if is_less(&v[c], &v[a]) {
302 mem::swap(&mut a, &mut c);
303 }
304 if is_less(&v[c], &v[b]) {
305 return c;
306 }
307 if is_less(&v[b], &v[a]) {
308 return a;
309 }
310 b
311}
312