1 | use crate::size_hint; |
2 | use std::{ |
3 | fmt, |
4 | iter::{DoubleEndedIterator, FusedIterator}, |
5 | }; |
6 | |
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 | |
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 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 | { |
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 | { |
150 | fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { |
151 | f&mut DebugStruct<'_, '_>.debug_struct("FlattenOk" ) |
152 | .field("iter" , &self.iter) |
153 | .field("inner_front" , &self.inner_front) |
154 | .field(name:"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 | |