itertools/
flatten_ok.rs

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}
itertools/flatten_ok.rs

itertools/
flatten_ok.rs
