Rollup merge of #55448 - Mokosha:SortAtIndex, r=bluss

Add 'partition_at_index/_by/_by_key' for slices.

This is an analog to C++'s std::nth_element (a.k.a. quickselect).

Corresponds to tracking bug #55300.

2 files changed, 236 insertions, 0 deletions

diff --git a/src/libcore/slice/mod.rs b/src/libcore/slice/mod.rs
index 4eb5bddb5d2..122ef9c79c2 100644
--- a/src/libcore/slice/mod.rs
+++ b/src/libcore/slice/mod.rs
@@ -1585,6 +1585,153 @@ impl<T> [T] {
         sort::quicksort(self, |a, b| f(a).lt(&f(b)));
     }
 
+    /// Reorder the slice such that the element at `index` is at its final sorted position.
+    ///
+    /// This reordering has the additional property that any value at position `i < index` will be
+    /// less than or equal to any value at a position `j > index`. Additionally, this reordering is
+    /// unstable (i.e. any number of equal elements may end up at position `index`), in-place
+    /// (i.e. does not allocate), and `O(n)` worst-case. This function is also/ known as "kth
+    /// element" in other libraries. It returns a triplet of the following values: all elements less
+    /// than the one at the given index, the value at the given index, and all elements greater than
+    /// the one at the given index.
+    ///
+    /// # Current implementation
+    ///
+    /// The current algorithm is based on the quickselect portion of the same quicksort algorithm
+    /// used for [`sort_unstable`].
+    ///
+    /// [`sort_unstable`]: #method.sort_unstable
+    ///
+    /// # Panics
+    ///
+    /// Panics when `index >= len()`, meaning it always panics on empty slices.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// #![feature(slice_partition_at_index)]
+    ///
+    /// let mut v = [-5i32, 4, 1, -3, 2];
+    ///
+    /// // Find the median
+    /// v.partition_at_index(2);
+    ///
+    /// // We are only guaranteed the slice will be one of the following, based on the way we sort
+    /// // about the specified index.
+    /// assert!(v == [-3, -5, 1, 2, 4] ||
+    ///         v == [-5, -3, 1, 2, 4] ||
+    ///         v == [-3, -5, 1, 4, 2] ||
+    ///         v == [-5, -3, 1, 4, 2]);
+    /// ```
+    #[unstable(feature = "slice_partition_at_index", issue = "55300")]
+    #[inline]
+    pub fn partition_at_index(&mut self, index: usize) -> (&mut [T], &mut T, &mut [T])
+        where T: Ord
+    {
+        let mut f = |a: &T, b: &T| a.lt(b);
+        sort::partition_at_index(self, index, &mut f)
+    }
+
+    /// Reorder the slice with a comparator function such that the element at `index` is at its
+    /// final sorted position.
+    ///
+    /// This reordering has the additional property that any value at position `i < index` will be
+    /// less than or equal to any value at a position `j > index` using the comparator function.
+    /// Additionally, this reordering is unstable (i.e. any number of equal elements may end up at
+    /// position `index`), in-place (i.e. does not allocate), and `O(n)` worst-case. This function
+    /// is also known as "kth element" in other libraries. It returns a triplet of the following
+    /// values: all elements less than the one at the given index, the value at the given index,
+    /// and all elements greater than the one at the given index, using the provided comparator
+    /// function.
+    ///
+    /// # Current implementation
+    ///
+    /// The current algorithm is based on the quickselect portion of the same quicksort algorithm
+    /// used for [`sort_unstable`].
+    ///
+    /// [`sort_unstable`]: #method.sort_unstable
+    ///
+    /// # Panics
+    ///
+    /// Panics when `index >= len()`, meaning it always panics on empty slices.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// #![feature(slice_partition_at_index)]
+    ///
+    /// let mut v = [-5i32, 4, 1, -3, 2];
+    ///
+    /// // Find the median as if the slice were sorted in descending order.
+    /// v.partition_at_index_by(2, |a, b| b.cmp(a));
+    ///
+    /// // We are only guaranteed the slice will be one of the following, based on the way we sort
+    /// // about the specified index.
+    /// assert!(v == [2, 4, 1, -5, -3] ||
+    ///         v == [2, 4, 1, -3, -5] ||
+    ///         v == [4, 2, 1, -5, -3] ||
+    ///         v == [4, 2, 1, -3, -5]);
+    /// ```
+    #[unstable(feature = "slice_partition_at_index", issue = "55300")]
+    #[inline]
+    pub fn partition_at_index_by<F>(&mut self, index: usize, mut compare: F)
+                                    -> (&mut [T], &mut T, &mut [T])
+        where F: FnMut(&T, &T) -> Ordering
+    {
+        let mut f = |a: &T, b: &T| compare(a, b) == Less;
+        sort::partition_at_index(self, index, &mut f)
+    }
+
+    /// Reorder the slice with a key extraction function such that the element at `index` is at its
+    /// final sorted position.
+    ///
+    /// This reordering has the additional property that any value at position `i < index` will be
+    /// less than or equal to any value at a position `j > index` using the key extraction function.
+    /// Additionally, this reordering is unstable (i.e. any number of equal elements may end up at
+    /// position `index`), in-place (i.e. does not allocate), and `O(n)` worst-case. This function
+    /// is also known as "kth element" in other libraries. It returns a triplet of the following
+    /// values: all elements less than the one at the given index, the value at the given index, and
+    /// all elements greater than the one at the given index, using the provided key extraction
+    /// function.
+    ///
+    /// # Current implementation
+    ///
+    /// The current algorithm is based on the quickselect portion of the same quicksort algorithm
+    /// used for [`sort_unstable`].
+    ///
+    /// [`sort_unstable`]: #method.sort_unstable
+    ///
+    /// # Panics
+    ///
+    /// Panics when `index >= len()`, meaning it always panics on empty slices.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// #![feature(slice_partition_at_index)]
+    ///
+    /// let mut v = [-5i32, 4, 1, -3, 2];
+    ///
+    /// // Return the median as if the array were sorted according to absolute value.
+    /// v.partition_at_index_by_key(2, |a| a.abs());
+    ///
+    /// // We are only guaranteed the slice will be one of the following, based on the way we sort
+    /// // about the specified index.
+    /// assert!(v == [1, 2, -3, 4, -5] ||
+    ///         v == [1, 2, -3, -5, 4] ||
+    ///         v == [2, 1, -3, 4, -5] ||
+    ///         v == [2, 1, -3, -5, 4]);
+    /// ```
+    #[unstable(feature = "slice_partition_at_index", issue = "55300")]
+    #[inline]
+    pub fn partition_at_index_by_key<K, F>(&mut self, index: usize, mut f: F)
+                                           -> (&mut [T], &mut T, &mut [T])
+        where F: FnMut(&T) -> K, K: Ord
+    {
+        let mut g = |a: &T, b: &T| f(a).lt(&f(b));
+        sort::partition_at_index(self, index, &mut g)
+    }
+
     /// Moves all consecutive repeated elements to the end of the slice according to the
     /// [`PartialEq`] trait implementation.
     ///
diff --git a/src/libcore/slice/sort.rs b/src/libcore/slice/sort.rs
index 3f84faa0499..68f1fb4b526 100644
--- a/src/libcore/slice/sort.rs
+++ b/src/libcore/slice/sort.rs
@@ -691,3 +691,92 @@ pub fn quicksort<T, F>(v: &mut [T], mut is_less: F)
 
     recurse(v, &mut is_less, None, limit);
 }
+
+fn partition_at_index_loop<'a, T, F>( mut v: &'a mut [T], mut index: usize, is_less: &mut F
+                                    , mut pred: Option<&'a T>) where F: FnMut(&T, &T) -> bool
+{
+    loop {
+        // For slices of up to this length it's probably faster to simply sort them.
+        const MAX_INSERTION: usize = 10;
+        if v.len() <= MAX_INSERTION {
+            insertion_sort(v, is_less);
+            return;
+        }
+
+        // Choose a pivot
+        let (pivot, _) = choose_pivot(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.
+        // This case is usually hit when the slice contains many duplicate elements.
+        if let Some(p) = pred {
+            if !is_less(p, &v[pivot]) {
+                let mid = partition_equal(v, pivot, is_less);
+
+                // If we've passed our index, then we're good.
+                if mid > index {
+                    return;
+                }
+
+                // Otherwise, continue sorting elements greater than the pivot.
+                v = &mut v[mid..];
+                index = index - mid;
+                pred = None;
+                continue;
+            }
+        }
+
+        let (mid, _) = partition(v, pivot, is_less);
+
+        // Split the slice into `left`, `pivot`, and `right`.
+        let (left, right) = {v}.split_at_mut(mid);
+        let (pivot, right) = right.split_at_mut(1);
+        let pivot = &pivot[0];
+
+        if mid < index {
+            v = right;
+            index = index - mid - 1;
+            pred = Some(pivot);
+        } else if mid > index {
+            v = left;
+        } else {
+            // If mid == index, then we're done, since partition() guaranteed that all elements
+            // after mid are greater than or equal to mid.
+            return;
+        }
+    }
+}
+
+pub fn partition_at_index<T, F>(v: &mut [T], index: usize, mut is_less: F)
+                                -> (&mut [T], &mut T, &mut [T]) where F: FnMut(&T, &T) -> bool
+{
+    use cmp::Ordering::Less;
+    use cmp::Ordering::Greater;
+
+    if index >= v.len() {
+        panic!("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.
+    } else if index == v.len() - 1 {
+        // Find max element and place it in the last position of the array. We're free to use
+        // `unwrap()` here because we know v must not be empty.
+        let (max_index, _) = v.iter().enumerate().max_by(
+            |&(_, x), &(_, y)| if is_less(x, y) { Less } else { Greater }).unwrap();
+        v.swap(max_index, index);
+    } else if index == 0 {
+        // Find min element and place it in the first position of the array. We're free to use
+        // `unwrap()` here because we know v must not be empty.
+        let (min_index, _) = v.iter().enumerate().min_by(
+            |&(_, x), &(_, y)| if is_less(x, y) { Less } else { Greater }).unwrap();
+        v.swap(min_index, index);
+    } else {
+        partition_at_index_loop(v, index, &mut is_less, None);
+    }
+
+    let (left, right) = v.split_at_mut(index);
+    let (pivot, right) = right.split_at_mut(1);
+    let pivot = &mut pivot[0];
+    (left, pivot, right)
+}