Jakub Kotur | a425e55 | 2020-12-21 17:28:15 +0100 | [diff] [blame] | 1 | use std::cmp::Ordering::{Equal, Greater, Less}; |
| 2 | use super::size_hint; |
Joel Galenson | b593e25 | 2021-06-21 13:15:57 -0700 | [diff] [blame] | 3 | use std::iter::{Fuse, FusedIterator}; |
Jakub Kotur | a425e55 | 2020-12-21 17:28:15 +0100 | [diff] [blame] | 4 | |
| 5 | use crate::either_or_both::EitherOrBoth; |
| 6 | |
| 7 | // ZipLongest originally written by SimonSapin, |
| 8 | // and dedicated to itertools https://github.com/rust-lang/rust/pull/19283 |
| 9 | |
| 10 | /// An iterator which iterates two other iterators simultaneously |
| 11 | /// |
| 12 | /// This iterator is *fused*. |
| 13 | /// |
Joel Galenson | b593e25 | 2021-06-21 13:15:57 -0700 | [diff] [blame] | 14 | /// See [`.zip_longest()`](crate::Itertools::zip_longest) for more information. |
Jakub Kotur | a425e55 | 2020-12-21 17:28:15 +0100 | [diff] [blame] | 15 | #[derive(Clone, Debug)] |
| 16 | #[must_use = "iterator adaptors are lazy and do nothing unless consumed"] |
| 17 | pub struct ZipLongest<T, U> { |
| 18 | a: Fuse<T>, |
| 19 | b: Fuse<U>, |
| 20 | } |
| 21 | |
| 22 | /// Create a new `ZipLongest` iterator. |
Joel Galenson | b593e25 | 2021-06-21 13:15:57 -0700 | [diff] [blame] | 23 | pub fn zip_longest<T, U>(a: T, b: U) -> ZipLongest<T, U> |
Jakub Kotur | a425e55 | 2020-12-21 17:28:15 +0100 | [diff] [blame] | 24 | where T: Iterator, |
| 25 | U: Iterator |
| 26 | { |
| 27 | ZipLongest { |
| 28 | a: a.fuse(), |
| 29 | b: b.fuse(), |
| 30 | } |
| 31 | } |
| 32 | |
| 33 | impl<T, U> Iterator for ZipLongest<T, U> |
| 34 | where T: Iterator, |
| 35 | U: Iterator |
| 36 | { |
| 37 | type Item = EitherOrBoth<T::Item, U::Item>; |
| 38 | |
| 39 | #[inline] |
| 40 | fn next(&mut self) -> Option<Self::Item> { |
| 41 | match (self.a.next(), self.b.next()) { |
| 42 | (None, None) => None, |
| 43 | (Some(a), None) => Some(EitherOrBoth::Left(a)), |
| 44 | (None, Some(b)) => Some(EitherOrBoth::Right(b)), |
| 45 | (Some(a), Some(b)) => Some(EitherOrBoth::Both(a, b)), |
| 46 | } |
| 47 | } |
| 48 | |
| 49 | #[inline] |
| 50 | fn size_hint(&self) -> (usize, Option<usize>) { |
| 51 | size_hint::max(self.a.size_hint(), self.b.size_hint()) |
| 52 | } |
| 53 | } |
| 54 | |
| 55 | impl<T, U> DoubleEndedIterator for ZipLongest<T, U> |
| 56 | where T: DoubleEndedIterator + ExactSizeIterator, |
| 57 | U: DoubleEndedIterator + ExactSizeIterator |
| 58 | { |
| 59 | #[inline] |
| 60 | fn next_back(&mut self) -> Option<Self::Item> { |
| 61 | match self.a.len().cmp(&self.b.len()) { |
| 62 | Equal => match (self.a.next_back(), self.b.next_back()) { |
| 63 | (None, None) => None, |
| 64 | (Some(a), Some(b)) => Some(EitherOrBoth::Both(a, b)), |
| 65 | // These can only happen if .len() is inconsistent with .next_back() |
| 66 | (Some(a), None) => Some(EitherOrBoth::Left(a)), |
| 67 | (None, Some(b)) => Some(EitherOrBoth::Right(b)), |
| 68 | }, |
| 69 | Greater => self.a.next_back().map(EitherOrBoth::Left), |
| 70 | Less => self.b.next_back().map(EitherOrBoth::Right), |
| 71 | } |
| 72 | } |
| 73 | } |
| 74 | |
| 75 | impl<T, U> ExactSizeIterator for ZipLongest<T, U> |
| 76 | where T: ExactSizeIterator, |
| 77 | U: ExactSizeIterator |
| 78 | {} |
Joel Galenson | b593e25 | 2021-06-21 13:15:57 -0700 | [diff] [blame] | 79 | |
| 80 | impl<T, U> FusedIterator for ZipLongest<T, U> |
| 81 | where T: Iterator, |
| 82 | U: Iterator |
| 83 | {} |