1use crate::size_hint;
2use std::{
3 fmt,
4 iter::{DoubleEndedIterator, FusedIterator},
5};
6
7pub fn flatten_ok<I, T, E>(iter: I) -> FlattenOk<I, T, E>
8where
9 I: Iterator<Item = Result<T, E>>,
10 T: IntoIterator,
11{
12 FlattenOk {
13 iter,
14 inner_front: None,
15 inner_back: None,
16 }
17}
18
19/// An iterator adaptor that flattens `Result::Ok` values and
20/// allows `Result::Err` values through unchanged.
21///
22/// See [`.flatten_ok()`](crate::Itertools::flatten_ok) for more information.
23#[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
24pub struct FlattenOk<I, T, E>
25where
26 I: Iterator<Item = Result<T, E>>,
27 T: IntoIterator,
28{
29 iter: I,
30 inner_front: Option<T::IntoIter>,
31 inner_back: Option<T::IntoIter>,
32}
33
34impl<I, T, E> Iterator for FlattenOk<I, T, E>
35where
36 I: Iterator<Item = Result<T, E>>,
37 T: IntoIterator,
38{
39 type Item = Result<T::Item, E>;
40
41 fn next(&mut self) -> Option<Self::Item> {
42 loop {
43 // Handle the front inner iterator.
44 if let Some(inner) = &mut self.inner_front {
45 if let Some(item) = inner.next() {
46 return Some(Ok(item));
47 }
48
49 // This is necessary for the iterator to implement `FusedIterator`
50 // with only the original iterator being fused.
51 self.inner_front = None;
52 }
53
54 match self.iter.next() {
55 Some(Ok(ok)) => self.inner_front = Some(ok.into_iter()),
56 Some(Err(e)) => return Some(Err(e)),
57 None => {
58 // Handle the back inner iterator.
59 if let Some(inner) = &mut self.inner_back {
60 if let Some(item) = inner.next() {
61 return Some(Ok(item));
62 }
63
64 // This is necessary for the iterator to implement `FusedIterator`
65 // with only the original iterator being fused.
66 self.inner_back = None;
67 } else {
68 return None;
69 }
70 }
71 }
72 }
73 }
74
75 fn fold<B, F>(self, init: B, mut f: F) -> B
76 where
77 Self: Sized,
78 F: FnMut(B, Self::Item) -> B,
79 {
80 // Front
81 let mut acc = match self.inner_front {
82 Some(x) => x.fold(init, |a, o| f(a, Ok(o))),
83 None => init,
84 };
85
86 acc = self.iter.fold(acc, |acc, x| match x {
87 Ok(it) => it.into_iter().fold(acc, |a, o| f(a, Ok(o))),
88 Err(e) => f(acc, Err(e)),
89 });
90
91 // Back
92 match self.inner_back {
93 Some(x) => x.fold(acc, |a, o| f(a, Ok(o))),
94 None => acc,
95 }
96 }
97
98 fn size_hint(&self) -> (usize, Option<usize>) {
99 let inner_hint = |inner: &Option<T::IntoIter>| {
100 inner
101 .as_ref()
102 .map(Iterator::size_hint)
103 .unwrap_or((0, Some(0)))
104 };
105 let inner_front = inner_hint(&self.inner_front);
106 let inner_back = inner_hint(&self.inner_back);
107 // The outer iterator `Ok` case could be (0, None) as we don't know its size_hint yet.
108 let outer = match self.iter.size_hint() {
109 (0, Some(0)) => (0, Some(0)),
110 _ => (0, None),
111 };
112
113 size_hint::add(size_hint::add(inner_front, inner_back), outer)
114 }
115}
116
117impl<I, T, E> DoubleEndedIterator for FlattenOk<I, T, E>
118where
119 I: DoubleEndedIterator<Item = Result<T, E>>,
120 T: IntoIterator,
121 T::IntoIter: DoubleEndedIterator,
122{
123 fn next_back(&mut self) -> Option<Self::Item> {
124 loop {
125 // Handle the back inner iterator.
126 if let Some(inner) = &mut self.inner_back {
127 if let Some(item) = inner.next_back() {
128 return Some(Ok(item));
129 }
130
131 // This is necessary for the iterator to implement `FusedIterator`
132 // with only the original iterator being fused.
133 self.inner_back = None;
134 }
135
136 match self.iter.next_back() {
137 Some(Ok(ok)) => self.inner_back = Some(ok.into_iter()),
138 Some(Err(e)) => return Some(Err(e)),
139 None => {
140 // Handle the front inner iterator.
141 if let Some(inner) = &mut self.inner_front {
142 if let Some(item) = inner.next_back() {
143 return Some(Ok(item));
144 }
145
146 // This is necessary for the iterator to implement `FusedIterator`
147 // with only the original iterator being fused.
148 self.inner_front = None;
149 } else {
150 return None;
151 }
152 }
153 }
154 }
155 }
156
157 fn rfold<B, F>(self, init: B, mut f: F) -> B
158 where
159 Self: Sized,
160 F: FnMut(B, Self::Item) -> B,
161 {
162 // Back
163 let mut acc = match self.inner_back {
164 Some(x) => x.rfold(init, |a, o| f(a, Ok(o))),
165 None => init,
166 };
167
168 acc = self.iter.rfold(acc, |acc, x| match x {
169 Ok(it) => it.into_iter().rfold(acc, |a, o| f(a, Ok(o))),
170 Err(e) => f(acc, Err(e)),
171 });
172
173 // Front
174 match self.inner_front {
175 Some(x) => x.rfold(acc, |a, o| f(a, Ok(o))),
176 None => acc,
177 }
178 }
179}
180
181impl<I, T, E> Clone for FlattenOk<I, T, E>
182where
183 I: Iterator<Item = Result<T, E>> + Clone,
184 T: IntoIterator,
185 T::IntoIter: Clone,
186{
187 clone_fields!(iter, inner_front, inner_back);
188}
189
190impl<I, T, E> fmt::Debug for FlattenOk<I, T, E>
191where
192 I: Iterator<Item = Result<T, E>> + fmt::Debug,
193 T: IntoIterator,
194 T::IntoIter: fmt::Debug,
195{
196 debug_fmt_fields!(FlattenOk, iter, inner_front, inner_back);
197}
198
199/// Only the iterator being flattened needs to implement [`FusedIterator`].
200impl<I, T, E> FusedIterator for FlattenOk<I, T, E>
201where
202 I: FusedIterator<Item = Result<T, E>>,
203 T: IntoIterator,
204{
205}
206