Skip to content

Commit

Permalink
Rollup merge of rust-lang#85835 - Seppel3210:master, r=yaahc
Browse files Browse the repository at this point in the history
Implement Extend<(A, B)> for (Extend<A>, Extend<B>)

I oriented myself at the implementation of `Iterator::unzip` and also rewrote the impl in terms of `(A, B)::extend` after that.

Since (A, B) now also implements Extend we could also mention in the documentation of unzip that it can do "nested unzipping" (you could unzip `Iterator<Item=(A, (B, C))>` into `(Vec<A>, (Vec<B>, Vec<C>))` for example) but I'm not sure of that so I'm asking here 🙂

(P.S. I saw a couple of people asking if there is an unzip3 but there isn't. So this could be a way to get equivalent functionality)
  • Loading branch information
JohnTitor authored Aug 11, 2021
2 parents 362e0f5 + 3d0c5d0 commit 9979921
Show file tree
Hide file tree
Showing 2 changed files with 69 additions and 22 deletions.
58 changes: 58 additions & 0 deletions library/core/src/iter/traits/collect.rs
Original file line number Diff line number Diff line change
Expand Up @@ -360,3 +360,61 @@ impl Extend<()> for () {
}
fn extend_one(&mut self, _item: ()) {}
}

#[stable(feature = "extend_for_tuple", since = "1.56.0")]
impl<A, B, ExtendA, ExtendB> Extend<(A, B)> for (ExtendA, ExtendB)
where
ExtendA: Extend<A>,
ExtendB: Extend<B>,
{
/// Allows to `extend` a tuple of collections that also implement `Extend`.
///
/// See also: [`Iterator::unzip`]
///
/// # Examples
/// ```
/// let mut tuple = (vec![0], vec![1]);
/// tuple.extend(vec![(2, 3), (4, 5), (6, 7)]);
/// assert_eq!(tuple.0, vec![0, 2, 4, 6]);
/// assert_eq!(tuple.1, vec![1, 3, 5, 7]);
///
/// // also allows for arbitrarily nested tuples
/// let mut nested_tuple = (vec![(1, -1)], vec![(2, -2)]);
/// nested_tuple.extend(vec![((3, -3), (4, -4)), ((5, -5), (6, -6))]);
///
/// assert_eq!(nested_tuple.0, vec![(1, -1), (3, -3), (5, -5)]);
/// assert_eq!(nested_tuple.1, vec![(2, -2), (4, -4), (6, -6)]);
/// ```
fn extend<T: IntoIterator<Item = (A, B)>>(&mut self, into_iter: T) {
let (a, b) = self;
let iter = into_iter.into_iter();

fn extend<'a, A, B>(
a: &'a mut impl Extend<A>,
b: &'a mut impl Extend<B>,
) -> impl FnMut((), (A, B)) + 'a {
move |(), (t, u)| {
a.extend_one(t);
b.extend_one(u);
}
}

let (lower_bound, _) = iter.size_hint();
if lower_bound > 0 {
a.extend_reserve(lower_bound);
b.extend_reserve(lower_bound);
}

iter.fold((), extend(a, b));
}

fn extend_one(&mut self, item: (A, B)) {
self.0.extend_one(item.0);
self.1.extend_one(item.1);
}

fn extend_reserve(&mut self, additional: usize) {
self.0.extend_reserve(additional);
self.1.extend_reserve(additional);
}
}
33 changes: 11 additions & 22 deletions library/core/src/iter/traits/iterator.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2841,6 +2841,14 @@ pub trait Iterator {
///
/// assert_eq!(left, [1, 3]);
/// assert_eq!(right, [2, 4]);
///
/// // you can also unzip multiple nested tuples at once
/// let a = [(1, (2, 3)), (4, (5, 6))];
///
/// let (x, (y, z)): (Vec<_>, (Vec<_>, Vec<_>)) = a.iter().cloned().unzip();
/// assert_eq!(x, [1, 4]);
/// assert_eq!(y, [2, 5]);
/// assert_eq!(z, [3, 6]);
/// ```
#[stable(feature = "rust1", since = "1.0.0")]
fn unzip<A, B, FromA, FromB>(self) -> (FromA, FromB)
Expand All @@ -2849,28 +2857,9 @@ pub trait Iterator {
FromB: Default + Extend<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_one(t);
us.extend_one(u);
}
}

let mut ts: FromA = Default::default();
let mut us: FromB = Default::default();

let (lower_bound, _) = self.size_hint();
if lower_bound > 0 {
ts.extend_reserve(lower_bound);
us.extend_reserve(lower_bound);
}

self.fold((), extend(&mut ts, &mut us));

(ts, us)
let mut unzipped: (FromA, FromB) = Default::default();
unzipped.extend(self);
unzipped
}

/// Creates an iterator which copies all of its elements.
Expand Down

0 comments on commit 9979921

Please sign in to comment.