1use super::size_hint;
2use std::cmp::Ordering::{Equal, Greater, Less};
3use std::iter::{Fuse, FusedIterator};
4
5use 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///
14/// See [`.zip_longest()`](crate::Itertools::zip_longest) for more information.
15#[derive(Clone, Debug)]
16#[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
17pub struct ZipLongest<T, U> {
18 a: Fuse<T>,
19 b: Fuse<U>,
20}
21
22/// Create a new `ZipLongest` iterator.
23pub fn zip_longest<T, U>(a: T, b: U) -> ZipLongest<T, U>
24where
25 T: Iterator,
26 U: Iterator,
27{
28 ZipLongest {
29 a: a.fuse(),
30 b: b.fuse(),
31 }
32}
33
34impl<T, U> Iterator for ZipLongest<T, U>
35where
36 T: Iterator,
37 U: Iterator,
38{
39 type Item = EitherOrBoth<T::Item, U::Item>;
40
41 #[inline]
42 fn next(&mut self) -> Option<Self::Item> {
43 match (self.a.next(), self.b.next()) {
44 (None, None) => None,
45 (Some(a), None) => Some(EitherOrBoth::Left(a)),
46 (None, Some(b)) => Some(EitherOrBoth::Right(b)),
47 (Some(a), Some(b)) => Some(EitherOrBoth::Both(a, b)),
48 }
49 }
50
51 #[inline]
52 fn size_hint(&self) -> (usize, Option<usize>) {
53 size_hint::max(self.a.size_hint(), self.b.size_hint())
54 }
55
56 #[inline]
57 fn fold<B, F>(self, init: B, mut f: F) -> B
58 where
59 Self: Sized,
60 F: FnMut(B, Self::Item) -> B,
61 {
62 let Self { mut a, mut b } = self;
63 let res = a.try_fold(init, |init, a| match b.next() {
64 Some(b) => Ok(f(init, EitherOrBoth::Both(a, b))),
65 None => Err(f(init, EitherOrBoth::Left(a))),
66 });
67 match res {
68 Ok(acc) => b.map(EitherOrBoth::Right).fold(acc, f),
69 Err(acc) => a.map(EitherOrBoth::Left).fold(acc, f),
70 }
71 }
72}
73
74impl<T, U> DoubleEndedIterator for ZipLongest<T, U>
75where
76 T: DoubleEndedIterator + ExactSizeIterator,
77 U: DoubleEndedIterator + ExactSizeIterator,
78{
79 #[inline]
80 fn next_back(&mut self) -> Option<Self::Item> {
81 match self.a.len().cmp(&self.b.len()) {
82 Equal => match (self.a.next_back(), self.b.next_back()) {
83 (None, None) => None,
84 (Some(a), Some(b)) => Some(EitherOrBoth::Both(a, b)),
85 // These can only happen if .len() is inconsistent with .next_back()
86 (Some(a), None) => Some(EitherOrBoth::Left(a)),
87 (None, Some(b)) => Some(EitherOrBoth::Right(b)),
88 },
89 Greater => self.a.next_back().map(EitherOrBoth::Left),
90 Less => self.b.next_back().map(EitherOrBoth::Right),
91 }
92 }
93
94 fn rfold<B, F>(self, mut init: B, mut f: F) -> B
95 where
96 F: FnMut(B, Self::Item) -> B,
97 {
98 let Self { mut a, mut b } = self;
99 let a_len = a.len();
100 let b_len = b.len();
101 match a_len.cmp(&b_len) {
102 Equal => {}
103 Greater => {
104 init = a
105 .by_ref()
106 .rev()
107 .take(a_len - b_len)
108 .map(EitherOrBoth::Left)
109 .fold(init, &mut f)
110 }
111 Less => {
112 init = b
113 .by_ref()
114 .rev()
115 .take(b_len - a_len)
116 .map(EitherOrBoth::Right)
117 .fold(init, &mut f)
118 }
119 }
120 a.rfold(init, |acc, item_a| {
121 f(acc, EitherOrBoth::Both(item_a, b.next_back().unwrap()))
122 })
123 }
124}
125
126impl<T, U> ExactSizeIterator for ZipLongest<T, U>
127where
128 T: ExactSizeIterator,
129 U: ExactSizeIterator,
130{
131}
132
133impl<T, U> FusedIterator for ZipLongest<T, U>
134where
135 T: Iterator,
136 U: Iterator,
137{
138}
139