Skip to content

Improved the use of Extend in a few Iterator methods #65219

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 7 commits into from
Closed
76 changes: 49 additions & 27 deletions src/libcore/iter/traits/iterator.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1504,34 +1504,44 @@ pub trait Iterator {
/// assert_eq!(odd, vec![1, 3]);
/// ```
#[stable(feature = "rust1", since = "1.0.0")]
fn partition<B, F>(self, f: F) -> (B, B) where
fn partition<B, F>(self, mut predicate: F) -> (B, B)
where
Self: Sized,
B: Default + Extend<Self::Item>,
F: FnMut(&Self::Item) -> bool
F: FnMut(&Self::Item) -> bool,
{
let mut left: B = Default::default();
let mut right: B = Default::default();
let mut fused = self.fuse();

#[inline]
fn extend<'a, T, B: Extend<T>>(
mut f: impl FnMut(&T) -> bool + 'a,
left: &'a mut B,
right: &'a mut B,
) -> impl FnMut(T) + 'a {
move |x| {
if f(&x) {
left.extend(Some(x));
fn extend_rhs<'a, A, T: Extend<A>, P: FnMut(&A) -> bool + 'a>(
rhs: &'a mut T,
mut predicate: P,
) -> impl FnMut(A) -> Option<A> + 'a {
move |item| {
if predicate(&item) {
Some(item)
} else {
right.extend(Some(x));
rhs.extend(Some(item));
None
}
}
}

let mut left: B = Default::default();
let mut right: B = Default::default();
left.extend((&mut fused).filter_map(extend_rhs(&mut right, &mut predicate)));

self.for_each(extend(f, &mut left, &mut right));
// left.extend may not have fully consumed self.
right.extend(fused.filter(move |item| !predicate(item)));

// right.extend may not have fully consumed self, but in that case we
// assume that we don't care about the remaining elements of `self`,
// since both `left` and `right` finsished being extended

(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.
Expand Down Expand Up @@ -2373,29 +2383,41 @@ pub trait Iterator {
/// assert_eq!(right, [2, 4]);
/// ```
#[stable(feature = "rust1", since = "1.0.0")]
fn unzip<A, B, FromA, FromB>(self) -> (FromA, FromB) where
fn unzip<A, B, FromA, FromB>(self) -> (FromA, FromB)
where
FromA: Default + Extend<A>,
FromB: Default + Extend<B>,
Self: Sized + Iterator<Item=(A, B)>,
Self: Sized + Iterator<Item = (A, B)>,
{
fn extend<'a, A, B>(
ts: &'a mut impl Extend<A>,
us: &'a mut impl Extend<B>,
) -> impl FnMut((A, B)) + 'a {
move |(t, u)| {
ts.extend(Some(t));
us.extend(Some(u));
let mut lhs = FromA::default();
let mut rhs = FromB::default();
let mut fused = self.fuse();

#[inline]
fn extend_rhs<'a, A, B, T: Extend<B>>(rhs: &'a mut T) -> impl FnMut((A, B)) -> A + 'a {
move |(a, b)| {
rhs.extend(Some(b));
a
}
}

let mut ts: FromA = Default::default();
let mut us: FromB = Default::default();
#[inline]
fn second<A, B>((_a, b): (A, B)) -> B {
b
}

lhs.extend((&mut fused).map(extend_rhs(&mut rhs)));

self.for_each(extend(&mut ts, &mut us));
// lhs.extend may not have fully consumed the iterator
rhs.extend((&mut fused).map(second));

(ts, us)
// rhs.extend may not have fully consumed the iterator
fused.for_each(drop);

(lhs, rhs)
}


/// Creates an iterator which copies all of its elements.
///
/// This is useful when you have an iterator over `&T`, but you need an
Expand Down