iter_flow/
flatten.rs

1use crate::{
2    and_then::{self, AndThen, AndThenBorrow},
3    map_ok::{self, MapOk, MapOkBorrow},
4};
5
6/// An iterator adaptor that flattens `Result::Ok` values and allows `Result::Err` values through unchanged.
7#[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
8pub struct FlattenOk<I, T, E>
9where
10    I: Iterator<Item = Result<T, E>>,
11    T: IntoIterator,
12{
13    iter: I,
14    inner_front: Option<T::IntoIter>,
15}
16
17impl<I, T, E> Iterator for FlattenOk<I, T, E>
18where
19    I: Iterator<Item = Result<T, E>>,
20    T: IntoIterator,
21{
22    type Item = Result<T::Item, E>;
23
24    fn next(&mut self) -> Option<Self::Item> {
25        loop {
26            if let Some(inner) = &mut self.inner_front {
27                if let Some(item) = inner.next() {
28                    return Some(Ok(item));
29                }
30                self.inner_front = None;
31            }
32
33            match self.iter.next() {
34                Some(Ok(ok)) => self.inner_front = Some(ok.into_iter()),
35                Some(Err(e)) => return Some(Err(e)),
36                None => return None,
37            }
38        }
39    }
40
41    fn size_hint(&self) -> (usize, Option<usize>) {
42        let inner_front = self
43            .inner_front
44            .as_ref()
45            .map(Iterator::size_hint)
46            .unwrap_or((0, Some(0)));
47        // The outer iterator `Ok` case could be (0, None) as we don't know its size_hint yet.
48        let outer = match self.iter.size_hint() {
49            (0, Some(0)) => (0, Some(0)),
50            _ => (0, None),
51        };
52
53        size_hint_add(inner_front, outer)
54    }
55}
56
57#[inline]
58pub fn size_hint_add(a: (usize, Option<usize>), b: (usize, Option<usize>)) -> (usize, Option<usize>) {
59    let min = a.0.saturating_add(b.0);
60    let max = match (a.1, b.1) {
61        (Some(x), Some(y)) => x.checked_add(y),
62        _ => None,
63    };
64
65    (min, max)
66}
67
68/// Create a new `MapOk` wrapped in a `FlattenOk`.
69pub fn flat_map_ok<I, F, T, U, E>(iter: I, f: F) -> FlattenOk<MapOk<I, F>, U, E>
70where
71    I: Iterator<Item = Result<T, E>>,
72    F: FnMut(T) -> U,
73    U: IntoIterator,
74{
75    FlattenOk {
76        iter: map_ok::map_ok(iter, f),
77        inner_front: None,
78    }
79}
80
81/// Create a new `MapOkBorrow` wrapped in a `FlattenOk`.
82pub fn flat_map_ok_borrow<I, F, T, U, E>(iter: I, f: F) -> FlattenOk<MapOkBorrow<I, F>, U, E>
83where
84    I: Iterator<Item = Result<T, E>>,
85    F: FnMut(&T) -> U,
86    U: IntoIterator,
87{
88    FlattenOk {
89        iter: map_ok::map_ok_borrow(iter, f),
90        inner_front: None,
91    }
92}
93
94/// Create a new `AndThen` wrapped in a `FlattenOk`.
95pub fn and_then_flat<I, F, T, U, E, R>(iter: I, f: F) -> FlattenOk<AndThen<I, F>, U, E>
96where
97    I: Iterator<Item = Result<T, E>>,
98    F: FnMut(T) -> Result<U, R>,
99    E: From<R>,
100    U: IntoIterator,
101{
102    FlattenOk {
103        iter: and_then::and_then(iter, f),
104        inner_front: None,
105    }
106}
107
108/// Create a new `AndThen` wrapped in a `FlattenOk`.
109pub fn and_then_flat_borrow<I, F, T, U, E, R>(iter: I, f: F) -> FlattenOk<AndThenBorrow<I, F>, U, E>
110where
111    I: Iterator<Item = Result<T, E>>,
112    F: FnMut(&T) -> Result<U, R>,
113    E: From<R>,
114    U: IntoIterator,
115{
116    FlattenOk {
117        iter: and_then::and_then_borrow(iter, f),
118        inner_front: None,
119    }
120}