Skip to content

Commit fa3fffb

Browse files
bors[bot]Veykril
andauthored
Merge #565
565: Implement `Itertools::multiunzip` r=phimuemue a=Veykril Simple implementation of [Iterator::unzip](https://doc.rust-lang.org/std/iter/trait.Iterator.html#method.unzip) but for tuples of sizes 0-12. This requires adding a new public trait to be able to express the dependency between the iterator item tuple to the collection output tuple. But given that a `multiunzip` function is added to the Itertools trait the `MultiUnzip` trait does not have to be imported to make use of this. Another option would be a macro but that would require giving it some extra syntax to declare the arity/return type which seems suboptimal, while this approach just works as a method call akin to std's iterator unzip. Closes #362 Co-authored-by: Lukas Wirth <lukastw97@gmail.com>
2 parents 599ae8c + f39d142 commit fa3fffb

File tree

3 files changed

+118
-0
lines changed

3 files changed

+118
-0
lines changed

src/lib.rs

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -179,6 +179,7 @@ pub use crate::repeatn::repeat_n;
179179
#[allow(deprecated)]
180180
pub use crate::sources::{repeat_call, unfold, iterate};
181181
pub use crate::with_position::Position;
182+
pub use crate::unziptuple::{multiunzip, MultiUnzip};
182183
pub use crate::ziptuple::multizip;
183184
mod adaptors;
184185
mod either_or_both;
@@ -237,6 +238,7 @@ mod tuple_impl;
237238
mod duplicates_impl;
238239
#[cfg(feature = "use_std")]
239240
mod unique_impl;
241+
mod unziptuple;
240242
mod with_position;
241243
mod zip_eq_impl;
242244
mod zip_longest;
@@ -3403,6 +3405,33 @@ pub trait Itertools : Iterator {
34033405
{
34043406
self.map(f).counts()
34053407
}
3408+
3409+
/// Converts an iterator of tuples into a tuple of containers.
3410+
///
3411+
/// `unzip()` consumes an entire iterator of n-ary tuples, producing `n` collections, one for each
3412+
/// column.
3413+
///
3414+
/// This function is, in some sense, the opposite of [`multizip`].
3415+
///
3416+
/// ```
3417+
/// use itertools::Itertools;
3418+
///
3419+
/// let inputs = vec![(1, 2, 3), (4, 5, 6), (7, 8, 9)];
3420+
///
3421+
/// let (a, b, c): (Vec<_>, Vec<_>, Vec<_>) = inputs
3422+
/// .into_iter()
3423+
/// .multiunzip();
3424+
///
3425+
/// assert_eq!(a, vec![1, 4, 7]);
3426+
/// assert_eq!(b, vec![2, 5, 8]);
3427+
/// assert_eq!(c, vec![3, 6, 9]);
3428+
/// ```
3429+
fn multiunzip<FromI>(self) -> FromI
3430+
where
3431+
Self: Sized + MultiUnzip<FromI>,
3432+
{
3433+
MultiUnzip::multiunzip(self)
3434+
}
34063435
}
34073436

34083437
impl<T: ?Sized> Itertools for T where T: Iterator { }

src/unziptuple.rs

Lines changed: 80 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,80 @@
1+
/// Converts an iterator of tuples into a tuple of containers.
2+
///
3+
/// `unzip()` consumes an entire iterator of n-ary tuples, producing `n` collections, one for each
4+
/// column.
5+
///
6+
/// This function is, in some sense, the opposite of [`multizip`].
7+
///
8+
/// ```
9+
/// use itertools::multiunzip;
10+
///
11+
/// let inputs = vec![(1, 2, 3), (4, 5, 6), (7, 8, 9)];
12+
///
13+
/// let (a, b, c): (Vec<_>, Vec<_>, Vec<_>) = multiunzip(inputs);
14+
///
15+
/// assert_eq!(a, vec![1, 4, 7]);
16+
/// assert_eq!(b, vec![2, 5, 8]);
17+
/// assert_eq!(c, vec![3, 6, 9]);
18+
/// ```
19+
///
20+
/// [`multizip`]: crate::multizip
21+
pub fn multiunzip<FromI, I>(i: I) -> FromI
22+
where
23+
I: IntoIterator,
24+
I::IntoIter: MultiUnzip<FromI>,
25+
{
26+
i.into_iter().multiunzip()
27+
}
28+
29+
/// An iterator that can be unzipped into multiple collections.
30+
///
31+
/// See [`.multiunzip()`](crate::Itertools::multiunzip) for more information.
32+
pub trait MultiUnzip<FromI>: Iterator {
33+
/// Unzip this iterator into multiple collections.
34+
fn multiunzip(self) -> FromI;
35+
}
36+
37+
macro_rules! impl_unzip_iter {
38+
($($T:ident => $FromT:ident),*) => (
39+
#[allow(non_snake_case)]
40+
impl<IT: Iterator<Item = ($($T,)*)>, $($T, $FromT: Default + Extend<$T>),* > MultiUnzip<($($FromT,)*)> for IT {
41+
fn multiunzip(self) -> ($($FromT,)*) {
42+
// This implementation mirrors the logic of Iterator::unzip as close as possible.
43+
// Unfortunately a lot of the used api there is still unstable represented by
44+
// the commented out parts that follow.
45+
//
46+
// https://doc.rust-lang.org/src/core/iter/traits/iterator.rs.html#2816-2844
47+
48+
let mut res = ($($FromT::default(),)*);
49+
let ($($FromT,)*) = &mut res;
50+
51+
// Still unstable #72631
52+
// let (lower_bound, _) = self.size_hint();
53+
// if lower_bound > 0 {
54+
// $($FromT.extend_reserve(lower_bound);)*
55+
// }
56+
57+
self.fold((), |(), ($($T,)*)| {
58+
// Still unstable #72631
59+
// $( $FromT.extend_one($T); )*
60+
$( $FromT.extend(std::iter::once($T)); )*
61+
});
62+
res
63+
}
64+
}
65+
);
66+
}
67+
68+
impl_unzip_iter!();
69+
impl_unzip_iter!(A => FromA);
70+
impl_unzip_iter!(A => FromA, B => FromB);
71+
impl_unzip_iter!(A => FromA, B => FromB, C => FromC);
72+
impl_unzip_iter!(A => FromA, B => FromB, C => FromC, D => FromD);
73+
impl_unzip_iter!(A => FromA, B => FromB, C => FromC, D => FromD, E => FromE);
74+
impl_unzip_iter!(A => FromA, B => FromB, C => FromC, D => FromD, E => FromE, F => FromF);
75+
impl_unzip_iter!(A => FromA, B => FromB, C => FromC, D => FromD, E => FromE, F => FromF, G => FromG);
76+
impl_unzip_iter!(A => FromA, B => FromB, C => FromC, D => FromD, E => FromE, F => FromF, G => FromG, H => FromH);
77+
impl_unzip_iter!(A => FromA, B => FromB, C => FromC, D => FromD, E => FromE, F => FromF, G => FromG, H => FromH, I => FromI);
78+
impl_unzip_iter!(A => FromA, B => FromB, C => FromC, D => FromD, E => FromE, F => FromF, G => FromG, H => FromH, I => FromI, J => FromJ);
79+
impl_unzip_iter!(A => FromA, B => FromB, C => FromC, D => FromD, E => FromE, F => FromF, G => FromG, H => FromH, I => FromI, J => FromJ, K => FromK);
80+
impl_unzip_iter!(A => FromA, B => FromB, C => FromC, D => FromD, E => FromE, F => FromF, G => FromG, H => FromH, I => FromI, J => FromJ, K => FromK, L => FromL);

tests/test_std.rs

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1080,3 +1080,12 @@ fn exactly_one_question_mark_return() -> Result<(), ExactlyOneError<std::slice::
10801080
[].iter().exactly_one()?;
10811081
Ok(())
10821082
}
1083+
1084+
#[test]
1085+
fn multiunzip() {
1086+
let (a, b, c): (Vec<_>, Vec<_>, Vec<_>) = [(0, 1, 2), (3, 4, 5), (6, 7, 8)].iter().cloned().multiunzip();
1087+
assert_eq!((a, b, c), (vec![0, 3, 6], vec![1, 4, 7], vec![2, 5, 8]));
1088+
let (): () = [(), (), ()].iter().cloned().multiunzip();
1089+
let t: (Vec<_>, Vec<_>, Vec<_>, Vec<_>, Vec<_>, Vec<_>, Vec<_>, Vec<_>, Vec<_>, Vec<_>, Vec<_>, Vec<_>) = [(0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11)].iter().cloned().multiunzip();
1090+
assert_eq!(t, (vec![0], vec![1], vec![2], vec![3], vec![4], vec![5], vec![6], vec![7], vec![8], vec![9], vec![10], vec![11]));
1091+
}

0 commit comments

Comments
 (0)