1 use crate::size_hint;
2 use std::{
3     fmt,
4     iter::{DoubleEndedIterator, FusedIterator},
5 };
6 
flatten_ok<I, T, E>(iter: I) -> FlattenOk<I, T, E> where I: Iterator<Item = Result<T, E>>, T: IntoIterator,7 pub fn flatten_ok<I, T, E>(iter: I) -> FlattenOk<I, T, E>
8 where
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"]
24 pub struct FlattenOk<I, T, E>
25 where
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 
34 impl<I, T, E> Iterator for FlattenOk<I, T, E>
35 where
36     I: Iterator<Item = Result<T, E>>,
37     T: IntoIterator,
38 {
39     type Item = Result<T::Item, E>;
40 
next(&mut self) -> Option<Self::Item>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 
size_hint(&self) -> (usize, Option<usize>)75     fn size_hint(&self) -> (usize, Option<usize>) {
76         let inner_hint = |inner: &Option<T::IntoIter>| {
77             inner
78                 .as_ref()
79                 .map(Iterator::size_hint)
80                 .unwrap_or((0, Some(0)))
81         };
82         let inner_front = inner_hint(&self.inner_front);
83         let inner_back = inner_hint(&self.inner_back);
84         // The outer iterator `Ok` case could be (0, None) as we don't know its size_hint yet.
85         let outer = match self.iter.size_hint() {
86             (0, Some(0)) => (0, Some(0)),
87             _ => (0, None),
88         };
89 
90         size_hint::add(size_hint::add(inner_front, inner_back), outer)
91     }
92 }
93 
94 impl<I, T, E> DoubleEndedIterator for FlattenOk<I, T, E>
95 where
96     I: DoubleEndedIterator<Item = Result<T, E>>,
97     T: IntoIterator,
98     T::IntoIter: DoubleEndedIterator,
99 {
next_back(&mut self) -> Option<Self::Item>100     fn next_back(&mut self) -> Option<Self::Item> {
101         loop {
102             // Handle the back inner iterator.
103             if let Some(inner) = &mut self.inner_back {
104                 if let Some(item) = inner.next_back() {
105                     return Some(Ok(item));
106                 }
107 
108                 // This is necessary for the iterator to implement `FusedIterator`
109                 // with only the original iterator being fused.
110                 self.inner_back = None;
111             }
112 
113             match self.iter.next_back() {
114                 Some(Ok(ok)) => self.inner_back = Some(ok.into_iter()),
115                 Some(Err(e)) => return Some(Err(e)),
116                 None => {
117                     // Handle the front inner iterator.
118                     if let Some(inner) = &mut self.inner_front {
119                         if let Some(item) = inner.next_back() {
120                             return Some(Ok(item));
121                         }
122 
123                         // This is necessary for the iterator to implement `FusedIterator`
124                         // with only the original iterator being fused.
125                         self.inner_front = None;
126                     } else {
127                         return None;
128                     }
129                 }
130             }
131         }
132     }
133 }
134 
135 impl<I, T, E> Clone for FlattenOk<I, T, E>
136 where
137     I: Iterator<Item = Result<T, E>> + Clone,
138     T: IntoIterator,
139     T::IntoIter: Clone,
140 {
141     clone_fields!(iter, inner_front, inner_back);
142 }
143 
144 impl<I, T, E> fmt::Debug for FlattenOk<I, T, E>
145 where
146     I: Iterator<Item = Result<T, E>> + fmt::Debug,
147     T: IntoIterator,
148     T::IntoIter: fmt::Debug,
149 {
fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result150     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
151         f.debug_struct("FlattenOk")
152             .field("iter", &self.iter)
153             .field("inner_front", &self.inner_front)
154             .field("inner_back", &self.inner_back)
155             .finish()
156     }
157 }
158 
159 /// Only the iterator being flattened needs to implement [`FusedIterator`].
160 impl<I, T, E> FusedIterator for FlattenOk<I, T, E>
161 where
162     I: FusedIterator<Item = Result<T, E>>,
163     T: IntoIterator,
164 {
165 }
166