1 use std::cmp;
2
3 pub trait ZipLongestIteratorExt: Iterator + Sized {
4 /// Creates an iterator which iterates over both this and the specified
5 /// iterators simultaneously, yielding pairs of two optional elements.
6 /// When both iterators return None, all further invocations of next() will
7 /// return None.
8 ///
9 /// # Example
10 ///
11 /// ```rust
12 /// let a = [0i];
13 /// let b = [1i, 2i];
14 /// let mut it = a.iter().zip(b.iter());
15 /// let (x0, x1, x2) = (0i, 1i, 2i);
16 /// assert_eq!(it.next().unwrap(), (Some(&x0), Some(&x1)));
17 /// assert_eq!(it.next().unwrap(), (None, Some(&x2)));
18 /// assert!(it.next().is_none());
19 /// ```
20 #[inline]
zip_longest<U: Iterator>(self, other: U) -> ZipLongest<Self, U>21 fn zip_longest<U: Iterator>(self, other: U) -> ZipLongest<Self, U> {
22 ZipLongest{a: self, b: other}
23 }
24 }
25
26
27 /// An iterator which iterates two other iterators simultaneously
28 #[derive(Clone)]
29 #[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
30 pub struct ZipLongest<T, U> {
31 a: T,
32 b: U
33 }
34
35 impl<A, B, T: Iterator<Item = A>, U: Iterator<Item = B>> Iterator for ZipLongest<T, U> {
36 type Item = EitherOrBoth<A, B>;
37
38 #[inline]
next(&mut self) -> Option<EitherOrBoth<A, B>>39 fn next(&mut self) -> Option<EitherOrBoth<A, B>> {
40 match (self.a.next(), self.b.next()) {
41 (None, None) => None,
42 (Some(a), None) => Some(EitherOrBoth::Left(a)),
43 (None, Some(b)) => Some(EitherOrBoth::Right(b)),
44 (Some(a), Some(b)) => Some(EitherOrBoth::Both(a, b)),
45 }
46 }
47
48 #[inline]
size_hint(&self) -> (usize, Option<usize>)49 fn size_hint(&self) -> (usize, Option<usize>) {
50 let (a_lower, a_upper) = self.a.size_hint();
51 let (b_lower, b_upper) = self.b.size_hint();
52
53 let lower = cmp::max(a_lower, b_lower);
54
55 let upper = match (a_upper, b_upper) {
56 (Some(x), Some(y)) => Some(cmp::max(x,y)),
57 _ => None
58 };
59
60 (lower, upper)
61 }
62 }
63
64 impl<T, U> DoubleEndedIterator for ZipLongest<T, U>
65 where T: DoubleEndedIterator + ExactSizeIterator, U: DoubleEndedIterator + ExactSizeIterator {
66 #[inline]
next_back(&mut self) -> Option<<Self as Iterator>::Item>67 fn next_back(&mut self) -> Option<<Self as Iterator>::Item> {
68 use std::cmp::Ordering::{Equal, Greater, Less};
69 match self.a.len().cmp(&self.b.len()) {
70 Equal => match (self.a.next_back(), self.b.next_back()) {
71 (None, None) => None,
72 (Some(a), Some(b)) => Some(EitherOrBoth::Both(a, b)),
73 // XXX these can only happen if .len() is inconsistent with .next_back()
74 (Some(a), None) => Some(EitherOrBoth::Left(a)),
75 (None, Some(b)) => Some(EitherOrBoth::Right(b)),
76 },
77 Greater => self.a.next_back().map(EitherOrBoth::Left),
78 Less => self.b.next_back().map(EitherOrBoth::Right),
79 }
80 }
81 }
82
83 impl<T: ExactSizeIterator, U: ExactSizeIterator> ExactSizeIterator for ZipLongest<T, U> {}
84
85
86 impl<I> ZipLongestIteratorExt for I where I: Iterator {}
87
88
89 /// A value yielded by `ZipLongest`.
90 /// Contains one or two values,
91 /// depending on which of the input iterators are exhausted.
92 #[derive(Clone, PartialEq, Eq, Debug)]
93 pub enum EitherOrBoth<A, B> {
94 /// Neither input iterator is exhausted yet, yielding two values.
95 Both(A, B),
96 /// The parameter iterator of `.zip_longest()` is exhausted,
97 /// only yielding a value from the `self` iterator.
98 Left(A),
99 /// The `self` iterator of `.zip_longest()` is exhausted,
100 /// only yielding a value from the parameter iterator.
101 Right(B),
102 }
103
104
105 #[test]
test_iterator_size_hint()106 fn test_iterator_size_hint() {
107 use std::usize;
108
109 let c = std::iter::repeat(0);
110 let v: &[_] = &[0i32, 1, 2, 3, 4, 5, 6, 7, 8, 9];
111 let v2 = &[10i32, 11, 12];
112 let vi = v.iter();
113 assert_eq!(c.zip_longest(vi.clone()).size_hint(), (usize::MAX, None));
114 assert_eq!(vi.zip_longest(v2.iter()).size_hint(), (10, Some(10)));
115 }
116
117 #[test]
test_double_ended()118 fn test_double_ended() {
119 let xs = [1i32, 2, 3, 4, 5, 6];
120 let ys = [1i32, 2, 3, 7];
121 let a = xs.iter().map(|&x| x);
122 let b = ys.iter().map(|&x| x);
123 let mut it = a.zip_longest(b);
124 assert_eq!(it.next(), Some(EitherOrBoth::Both(1, 1)));
125 assert_eq!(it.next(), Some(EitherOrBoth::Both(2, 2)));
126 assert_eq!(it.next_back(), Some(EitherOrBoth::Left(6)));
127 assert_eq!(it.next_back(), Some(EitherOrBoth::Left(5)));
128 assert_eq!(it.next_back(), Some(EitherOrBoth::Both(4, 7)));
129 assert_eq!(it.next(), Some(EitherOrBoth::Both(3, 3)));
130 assert_eq!(it.next(), None);
131 }
132