Rollup merge of rust-lang#62278 - cuviper:iter-partition, r=alexcrichton

Add Iterator::partition_in_place() and is_partitioned()

`partition_in_place()` swaps `&mut T` items in-place to satisfy the
predicate, so all `true` items precede all `false` items. This requires
a `DoubleEndedIterator` so we can search from front and back for items
that need swapping.

`is_partitioned()` checks whether the predicate is already satisfied.

Author: Centril

src/libcore/iter/traits/iterator.rs

@@ -1472,6 +1472,11 @@ pub trait Iterator {
     /// `partition()` returns a pair, all of the elements for which it returned
     /// `true`, and all of the elements for which it returned `false`.
     ///
+    /// See also [`is_partitioned()`] and [`partition_in_place()`].
+    ///
+    /// [`is_partitioned()`]: #method.is_partitioned
+    /// [`partition_in_place()`]: #method.partition_in_place
+    ///
     /// # Examples
     ///
     /// Basic usage:
@@ -1506,6 +1511,101 @@ pub trait Iterator {
         (left, right)
     }
 
+    /// Reorder the elements of this iterator *in-place* according to the given predicate,
+    /// such that all those that return `true` precede all those that return `false`.
+    /// Returns the number of `true` elements found.
+    ///
+    /// The relative order of partitioned items is not maintained.
+    ///
+    /// See also [`is_partitioned()`] and [`partition()`].
+    ///
+    /// [`is_partitioned()`]: #method.is_partitioned
+    /// [`partition()`]: #method.partition
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// #![feature(iter_partition_in_place)]
+    ///
+    /// let mut a = [1, 2, 3, 4, 5, 6, 7];
+    ///
+    /// // Partition in-place between evens and odds
+    /// let i = a.iter_mut().partition_in_place(|&n| n % 2 == 0);
+    ///
+    /// assert_eq!(i, 3);
+    /// assert!(a[..i].iter().all(|&n| n % 2 == 0)); // evens
+    /// assert!(a[i..].iter().all(|&n| n % 2 == 1)); // odds
+    /// ```
+    #[unstable(feature = "iter_partition_in_place", reason = "new API", issue = "62543")]
+    fn partition_in_place<'a, T: 'a, P>(mut self, ref mut predicate: P) -> usize
+    where
+        Self: Sized + DoubleEndedIterator<Item = &'a mut T>,
+        P: FnMut(&T) -> bool,
+    {
+        // FIXME: should we worry about the count overflowing? The only way to have more than
+        // `usize::MAX` mutable references is with ZSTs, which aren't useful to partition...
+
+        // These closure "factory" functions exist to avoid genericity in `Self`.
+
+        #[inline]
+        fn is_false<'a, T>(
+            predicate: &'a mut impl FnMut(&T) -> bool,
+            true_count: &'a mut usize,
+        ) -> impl FnMut(&&mut T) -> bool + 'a {
+            move |x| {
+                let p = predicate(&**x);
+                *true_count += p as usize;
+                !p
+            }
+        }
+
+        #[inline]
+        fn is_true<T>(
+            predicate: &mut impl FnMut(&T) -> bool
+        ) -> impl FnMut(&&mut T) -> bool + '_ {
+            move |x| predicate(&**x)
+        }
+
+        // Repeatedly find the first `false` and swap it with the last `true`.
+        let mut true_count = 0;
+        while let Some(head) = self.find(is_false(predicate, &mut true_count)) {
+            if let Some(tail) = self.rfind(is_true(predicate)) {
+                crate::mem::swap(head, tail);
+                true_count += 1;
+            } else {
+                break;
+            }
+        }
+        true_count
+    }
+
+    /// Checks if the elements of this iterator are partitioned according to the given predicate,
+    /// such that all those that return `true` precede all those that return `false`.
+    ///
+    /// See also [`partition()`] and [`partition_in_place()`].
+    ///
+    /// [`partition()`]: #method.partition
+    /// [`partition_in_place()`]: #method.partition_in_place
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// #![feature(iter_is_partitioned)]
+    ///
+    /// assert!("Iterator".chars().is_partitioned(char::is_uppercase));
+    /// assert!(!"IntoIterator".chars().is_partitioned(char::is_uppercase));
+    /// ```
+    #[unstable(feature = "iter_is_partitioned", reason = "new API", issue = "62544")]
+    fn is_partitioned<P>(mut self, mut predicate: P) -> bool
+    where
+        Self: Sized,
+        P: FnMut(Self::Item) -> bool,
+    {
+        // Either all items test `true`, or the first clause stops at `false`
+        // and we check that there are no more `true` items after that.
+        self.all(&mut predicate) || !self.any(predicate)
+    }
+
     /// An iterator method that applies a function as long as it returns
     /// successfully, producing a single, final value.
     ///

src/libcore/tests/iter.rs

@@ -2460,3 +2460,39 @@ fn test_is_sorted() {
     assert!(!["c", "bb", "aaa"].iter().is_sorted());
     assert!(["c", "bb", "aaa"].iter().is_sorted_by_key(|s| s.len()));
 }
+
+#[test]
+fn test_partition() {
+    fn check(xs: &mut [i32], ref p: impl Fn(&i32) -> bool, expected: usize) {
+        let i = xs.iter_mut().partition_in_place(p);
+        assert_eq!(expected, i);
+        assert!(xs[..i].iter().all(p));
+        assert!(!xs[i..].iter().any(p));
+        assert!(xs.iter().is_partitioned(p));
+        if i == 0 || i == xs.len() {
+            assert!(xs.iter().rev().is_partitioned(p));
+        } else {
+            assert!(!xs.iter().rev().is_partitioned(p));
+        }
+    }
+
+    check(&mut [], |_| true, 0);
+    check(&mut [], |_| false, 0);
+
+    check(&mut [0], |_| true, 1);
+    check(&mut [0], |_| false, 0);
+
+    check(&mut [-1, 1], |&x| x > 0, 1);
+    check(&mut [-1, 1], |&x| x < 0, 1);
+
+    let ref mut xs = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9];
+    check(xs, |_| true, 10);
+    check(xs, |_| false, 0);
+    check(xs, |&x| x % 2 == 0, 5); // evens
+    check(xs, |&x| x % 2 == 1, 5); // odds
+    check(xs, |&x| x % 3 == 0, 4); // multiple of 3
+    check(xs, |&x| x % 4 == 0, 3); // multiple of 4
+    check(xs, |&x| x % 5 == 0, 2); // multiple of 5
+    check(xs, |&x| x < 3, 3); // small
+    check(xs, |&x| x > 6, 3); // large
+}

src/libcore/tests/lib.rs

@@ -31,6 +31,8 @@
 #![feature(slice_partition_dedup)]
 #![feature(int_error_matching)]
 #![feature(const_fn)]
+#![feature(iter_partition_in_place)]
+#![feature(iter_is_partitioned)]
 #![warn(rust_2018_idioms)]
 
 extern crate test;