David LeGare | b72e905 | 2022-03-02 16:21:08 +0000 | [diff] [blame] | 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); |