Jakub Kotur | a425e55 | 2020-12-21 17:28:15 +0100 | [diff] [blame] | 1 | use super::size_hint; |
| 2 | |
Joel Galenson | b593e25 | 2021-06-21 13:15:57 -0700 | [diff] [blame] | 3 | /// See [`multizip`] for more information. |
Jakub Kotur | a425e55 | 2020-12-21 17:28:15 +0100 | [diff] [blame] | 4 | #[derive(Clone, Debug)] |
| 5 | #[must_use = "iterator adaptors are lazy and do nothing unless consumed"] |
| 6 | pub struct Zip<T> { |
| 7 | t: T, |
| 8 | } |
| 9 | |
| 10 | /// An iterator that generalizes *.zip()* and allows running multiple iterators in lockstep. |
| 11 | /// |
| 12 | /// The iterator `Zip<(I, J, ..., M)>` is formed from a tuple of iterators (or values that |
Joel Galenson | b593e25 | 2021-06-21 13:15:57 -0700 | [diff] [blame] | 13 | /// implement [`IntoIterator`]) and yields elements |
Jakub Kotur | a425e55 | 2020-12-21 17:28:15 +0100 | [diff] [blame] | 14 | /// until any of the subiterators yields `None`. |
| 15 | /// |
| 16 | /// The iterator element type is a tuple like like `(A, B, ..., E)` where `A` to `E` are the |
| 17 | /// element types of the subiterator. |
| 18 | /// |
| 19 | /// **Note:** The result of this macro is a value of a named type (`Zip<(I, J, |
| 20 | /// ..)>` of each component iterator `I, J, ...`) if each component iterator is |
| 21 | /// nameable. |
| 22 | /// |
| 23 | /// Prefer [`izip!()`] over `multizip` for the performance benefits of using the |
| 24 | /// standard library `.zip()`. Prefer `multizip` if a nameable type is needed. |
| 25 | /// |
Jakub Kotur | a425e55 | 2020-12-21 17:28:15 +0100 | [diff] [blame] | 26 | /// ``` |
| 27 | /// use itertools::multizip; |
| 28 | /// |
| 29 | /// // iterate over three sequences side-by-side |
| 30 | /// let mut results = [0, 0, 0, 0]; |
| 31 | /// let inputs = [3, 7, 9, 6]; |
| 32 | /// |
| 33 | /// for (r, index, input) in multizip((&mut results, 0..10, &inputs)) { |
| 34 | /// *r = index * 10 + input; |
| 35 | /// } |
| 36 | /// |
| 37 | /// assert_eq!(results, [0 + 3, 10 + 7, 29, 36]); |
| 38 | /// ``` |
| 39 | pub fn multizip<T, U>(t: U) -> Zip<T> |
| 40 | where Zip<T>: From<U>, |
| 41 | Zip<T>: Iterator, |
| 42 | { |
| 43 | Zip::from(t) |
| 44 | } |
| 45 | |
| 46 | macro_rules! impl_zip_iter { |
| 47 | ($($B:ident),*) => ( |
| 48 | #[allow(non_snake_case)] |
| 49 | impl<$($B: IntoIterator),*> From<($($B,)*)> for Zip<($($B::IntoIter,)*)> { |
| 50 | fn from(t: ($($B,)*)) -> Self { |
| 51 | let ($($B,)*) = t; |
| 52 | Zip { t: ($($B.into_iter(),)*) } |
| 53 | } |
| 54 | } |
| 55 | |
| 56 | #[allow(non_snake_case)] |
| 57 | #[allow(unused_assignments)] |
| 58 | impl<$($B),*> Iterator for Zip<($($B,)*)> |
| 59 | where |
| 60 | $( |
| 61 | $B: Iterator, |
| 62 | )* |
| 63 | { |
| 64 | type Item = ($($B::Item,)*); |
| 65 | |
| 66 | fn next(&mut self) -> Option<Self::Item> |
| 67 | { |
| 68 | let ($(ref mut $B,)*) = self.t; |
| 69 | |
| 70 | // NOTE: Just like iter::Zip, we check the iterators |
| 71 | // for None in order. We may finish unevenly (some |
| 72 | // iterators gave n + 1 elements, some only n). |
| 73 | $( |
| 74 | let $B = match $B.next() { |
| 75 | None => return None, |
| 76 | Some(elt) => elt |
| 77 | }; |
| 78 | )* |
| 79 | Some(($($B,)*)) |
| 80 | } |
| 81 | |
| 82 | fn size_hint(&self) -> (usize, Option<usize>) |
| 83 | { |
| 84 | let sh = (::std::usize::MAX, None); |
| 85 | let ($(ref $B,)*) = self.t; |
| 86 | $( |
| 87 | let sh = size_hint::min($B.size_hint(), sh); |
| 88 | )* |
| 89 | sh |
| 90 | } |
| 91 | } |
| 92 | |
| 93 | #[allow(non_snake_case)] |
| 94 | impl<$($B),*> ExactSizeIterator for Zip<($($B,)*)> where |
| 95 | $( |
| 96 | $B: ExactSizeIterator, |
| 97 | )* |
| 98 | { } |
Joel Galenson | 6f79871 | 2021-04-01 17:03:06 -0700 | [diff] [blame] | 99 | |
| 100 | #[allow(non_snake_case)] |
| 101 | impl<$($B),*> DoubleEndedIterator for Zip<($($B,)*)> where |
| 102 | $( |
| 103 | $B: DoubleEndedIterator + ExactSizeIterator, |
| 104 | )* |
| 105 | { |
| 106 | #[inline] |
| 107 | fn next_back(&mut self) -> Option<Self::Item> { |
| 108 | let ($(ref mut $B,)*) = self.t; |
| 109 | let size = *[$( $B.len(), )*].iter().min().unwrap(); |
| 110 | |
| 111 | $( |
| 112 | if $B.len() != size { |
| 113 | for _ in 0..$B.len() - size { $B.next_back(); } |
| 114 | } |
| 115 | )* |
| 116 | |
| 117 | match ($($B.next_back(),)*) { |
| 118 | ($(Some($B),)*) => Some(($($B,)*)), |
| 119 | _ => None, |
| 120 | } |
| 121 | } |
| 122 | } |
Jakub Kotur | a425e55 | 2020-12-21 17:28:15 +0100 | [diff] [blame] | 123 | ); |
| 124 | } |
| 125 | |
| 126 | impl_zip_iter!(A); |
| 127 | impl_zip_iter!(A, B); |
| 128 | impl_zip_iter!(A, B, C); |
| 129 | impl_zip_iter!(A, B, C, D); |
| 130 | impl_zip_iter!(A, B, C, D, E); |
| 131 | impl_zip_iter!(A, B, C, D, E, F); |
| 132 | impl_zip_iter!(A, B, C, D, E, F, G); |
| 133 | impl_zip_iter!(A, B, C, D, E, F, G, H); |
Joel Galenson | 6f79871 | 2021-04-01 17:03:06 -0700 | [diff] [blame] | 134 | impl_zip_iter!(A, B, C, D, E, F, G, H, I); |
| 135 | impl_zip_iter!(A, B, C, D, E, F, G, H, I, J); |
| 136 | impl_zip_iter!(A, B, C, D, E, F, G, H, I, J, K); |
| 137 | impl_zip_iter!(A, B, C, D, E, F, G, H, I, J, K, L); |