Approximate Median-of-Medians in Select

library/core/benches/slice.rs

@@ -127,3 +127,44 @@ fn fill_byte_sized(b: &mut Bencher) {
         black_box(slice.fill(black_box(NewType(42))));
     });
 }
+
+fn quickselect_adversarial(b: &mut Bencher, cache: Cache) {
+    use core::cmp::Ordering;
+    let prepare_worst_case = |n: usize| -> Vec<u64> {
+        let mut values = vec![None; n];
+        let mut indices: Vec<usize> = (0..n).collect();
+        let mut ctr = 0;
+
+        indices.select_nth_unstable_by(n - 2, |lhs, rhs| match (values[*lhs], values[*rhs]) {
+            (None, None) => {
+                values[*lhs] = Some(ctr);
+                ctr += 1;
+                Ordering::Less
+            }
+            (None, Some(_)) => Ordering::Greater,
+            (Some(_), None) => Ordering::Less,
+            (Some(a), Some(b)) => a.cmp(&b),
+        });
+
+        values.into_iter().map(|v| v.unwrap_or(ctr)).collect()
+    };
+    let n = cache.size();
+    let mut worst_case = prepare_worst_case(n);
+
+    b.iter(move || {
+        worst_case.select_nth_unstable(n - 2);
+    });
+}
+
+#[bench]
+fn quickselect_l1(b: &mut Bencher) {
+    quickselect_adversarial(b, Cache::L1);
+}
+#[bench]
+fn quickselect_l2(b: &mut Bencher) {
+    quickselect_adversarial(b, Cache::L2);
+}
+#[bench]
+fn quickselect_l3(b: &mut Bencher) {
+    quickselect_adversarial(b, Cache::L3);
+}

library/core/src/slice/sort.rs

@@ -581,8 +581,7 @@ where
 
             // Swap the found pair of out-of-order elements.
             r -= 1;
-            let ptr = v.as_mut_ptr();
-            ptr::swap(ptr.add(l), ptr.add(r));
+            v.swap(l, r);
             l += 1;
         }
     }
@@ -639,6 +638,43 @@ fn break_patterns<T>(v: &mut [T]) {
     }
 }
 
+fn approx_median_of_medians<T, F>(v: &mut [T], is_less: &mut F) -> usize
+where
+    F: FnMut(&T, &T) -> bool,
+{
+    use cmp::Ordering::{Greater, Less};
+    const N: usize = 7;
+    const BUF_SIZE: usize = 32;
+
+    let mut buf = [0; BUF_SIZE];
+    let mut n = 0;
+    let len = v.len();
+    let mut curr_rng = len.wrapping_mul(0xdea1).wrapping_add(3) % BUF_SIZE;
+    for i in 0..len / N {
+        if n < BUF_SIZE {
+            buf[n] = i;
+            n += 1;
+        } else {
+            buf[curr_rng] = i;
+            curr_rng = curr_rng.wrapping_mul(0xdea1).wrapping_add(3) % BUF_SIZE;
+        }
+    }
+    let idxs = &mut buf[..n];
+    // SAFETY: Since it only uses indeces up to floor(len/N)-1, it will always be safe to
+    // index by i*N, i*(N+1).
+    unsafe {
+        for i in idxs.iter_mut() {
+            let start = *i * N;
+            insertion_sort(&mut v.get_unchecked_mut(start..start + N), is_less);
+            *i = start + N / 2;
+        }
+        idxs.sort_unstable_by(|&a, &b| {
+            if is_less(&v.get_unchecked(a), &v.get_unchecked(b)) { Less } else { Greater }
+        });
+    }
+    idxs[n / 2]
+}
+
 /// Chooses a pivot in `v` and returns the index and `true` if the slice is likely already sorted.
 ///
 /// Elements in `v` might be reordered in the process.
@@ -784,7 +820,7 @@ where
         was_partitioned = was_p;
 
         // Split the slice into `left`, `pivot`, and `right`.
-        let (left, right) = { v }.split_at_mut(mid);
+        let (left, right) = v.split_at_mut(mid);
         let (pivot, right) = right.split_at_mut(1);
         let pivot = &pivot[0];
 
@@ -835,7 +871,7 @@ fn partition_at_index_loop<'a, T, F>(
         }
 
         // Choose a pivot
-        let (pivot, _) = choose_pivot(v, is_less);
+        let pivot = approx_median_of_medians(v, is_less);
 
         // If the chosen pivot is equal to the predecessor, then it's the smallest element in the
         // slice. Partition the slice into elements equal to and elements greater than the pivot.
@@ -860,7 +896,7 @@ fn partition_at_index_loop<'a, T, F>(
         let (mid, _) = partition(v, pivot, is_less);
 
         // Split the slice into `left`, `pivot`, and `right`.
-        let (left, right) = { v }.split_at_mut(mid);
+        let (left, right) = v.split_at_mut(mid);
         let (pivot, right) = right.split_at_mut(1);
         let pivot = &pivot[0];
 
@@ -886,12 +922,14 @@ pub fn partition_at_index<T, F>(
 where
     F: FnMut(&T, &T) -> bool,
 {
-    use cmp::Ordering::Greater;
-    use cmp::Ordering::Less;
-
-    if index >= v.len() {
-        panic!("partition_at_index index {} greater than length of slice {}", index, v.len());
-    }
+    use cmp::Ordering::{Greater, Less};
+
+    assert!(
+        index < v.len(),
+        "partition_at_index index {} greater than length of slice {}",
+        index,
+        v.len()
+    );
 
     if mem::size_of::<T>() == 0 {
         // Sorting has no meaningful behavior on zero-sized types. Do nothing.