diff --git a/src/libcore/iter/traits/iterator.rs b/src/libcore/iter/traits/iterator.rs index b9a98236f18a7..6eddac672c158 100644 --- a/src/libcore/iter/traits/iterator.rs +++ b/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, + 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( + 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

(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. /// diff --git a/src/libcore/tests/iter.rs b/src/libcore/tests/iter.rs index 4d840ef24c8e6..b7b0849e2129b 100644 --- a/src/libcore/tests/iter.rs +++ b/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 +} diff --git a/src/libcore/tests/lib.rs b/src/libcore/tests/lib.rs index 4b48d1225902b..cbb6423d71010 100644 --- a/src/libcore/tests/lib.rs +++ b/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;