resiter/
flatten.rs

1//
2// This Source Code Form is subject to the terms of the Mozilla Public
3// License, v. 2.0. If a copy of the MPL was not distributed with this
4// file, You can obtain one at http://mozilla.org/MPL/2.0/.
5//
6
7/// Extension trait for `Iterator<Item = Result<O, E>>` to selectively transform Oks and Errors.
8pub trait Flatten<O, E>: Sized {
9    /// [flatten](Iterator::flatten) `Ok` values while leaving `Err`-values as is.
10    ///
11    /// ```
12    /// use resiter::flatten::Flatten;
13    /// use resiter::map::Map;
14    ///
15    /// let mapped: Vec<_> = vec![Ok(1), Ok(2), Err(2), Err(0), Ok(2)]
16    ///     .into_iter()
17    ///     .map_ok(|i| (0..i))
18    ///     .map_err(|i| 0..(i * 2))
19    ///     .flatten_ok()
20    ///     .collect();
21    ///
22    /// assert_eq!(
23    ///     mapped,
24    ///     [Ok(0), Ok(0), Ok(1), Err(0..4), Err(0..0), Ok(0), Ok(1)]
25    /// );
26    /// ```
27    fn flatten_ok<U, O2>(self) -> FlattenOk<Self, U>
28    where
29        U: IntoIterator<Item = O2>;
30    /// [flatten](Iterator::flatten) `Err` values while leaving `Ok`-values as is.
31    ///
32    /// ```
33    /// use resiter::flatten::Flatten;
34    /// use resiter::map::Map;
35    ///
36    /// let mapped: Vec<_> = vec![Ok(1), Ok(2), Err(2), Err(0), Ok(2)]
37    ///     .into_iter()
38    ///     .map_ok(|i| (0..i))
39    ///     .map_err(|i| 0..(i * 2))
40    ///     .flatten_err()
41    ///     .collect();
42    ///
43    /// assert_eq!(
44    ///     mapped,
45    ///     [
46    ///         Ok(0..1),
47    ///         Ok(0..2),
48    ///         Err(0),
49    ///         Err(1),
50    ///         Err(2),
51    ///         Err(3),
52    ///         Ok(0..2),
53    ///     ]
54    /// );
55    /// ```
56    fn flatten_err<U, E2>(self) -> FlattenErr<Self, U>
57    where
58        U: IntoIterator<Item = E2>;
59}
60
61impl<I, O, E> Flatten<O, E> for I
62where
63    I: Iterator<Item = Result<O, E>> + Sized,
64{
65    #[inline]
66    fn flatten_ok<U, O2>(self) -> FlattenOk<Self, U>
67    where
68        U: IntoIterator<Item = O2>,
69    {
70        FlattenOk {
71            frontiter: None,
72            iter: self,
73        }
74    }
75
76    #[inline]
77    fn flatten_err<U, E2>(self) -> FlattenErr<Self, U>
78    where
79        U: IntoIterator<Item = E2>,
80    {
81        FlattenErr {
82            frontiter: None,
83            iter: self,
84        }
85    }
86}
87
88#[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
89pub struct FlattenOk<I, U>
90where
91    U: IntoIterator,
92{
93    frontiter: Option<<U as IntoIterator>::IntoIter>,
94    iter: I,
95}
96
97impl<I, E, O2, U> Iterator for FlattenOk<I, U>
98where
99    I: Iterator<Item = Result<U, E>>,
100    U: IntoIterator<Item = O2>,
101{
102    type Item = Result<O2, E>;
103
104    fn next(&mut self) -> Option<Self::Item> {
105        loop {
106            if let Some(ref mut inner) = self.frontiter {
107                if let elt @ Some(_) = inner.next() {
108                    return elt.map(Ok);
109                }
110            }
111            match self.iter.next() {
112                None => return None,
113                Some(Ok(x)) => {
114                    self.frontiter = Some(x.into_iter());
115                }
116                Some(Err(e)) => return Some(Err(e)),
117            }
118        }
119    }
120
121    #[inline]
122    // TODO: Oh dear, this hint could be much better
123    // https://doc.rust-lang.org/src/core/iter/mod.rs.html#2694
124    fn size_hint(&self) -> (usize, Option<usize>) {
125        self.iter.size_hint()
126    }
127}
128
129#[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
130pub struct FlattenErr<I, U: IntoIterator> {
131    frontiter: Option<<U as IntoIterator>::IntoIter>,
132    iter: I,
133}
134
135impl<I, O, E2, U> Iterator for FlattenErr<I, U>
136where
137    I: Iterator<Item = Result<O, U>>,
138    U: IntoIterator<Item = E2>,
139{
140    type Item = Result<O, E2>;
141
142    fn next(&mut self) -> Option<Self::Item> {
143        loop {
144            if let Some(ref mut inner) = self.frontiter {
145                if let elt @ Some(_) = inner.next() {
146                    return elt.map(Err);
147                }
148            }
149            match self.iter.next() {
150                None => return None,
151                Some(Err(e)) => {
152                    self.frontiter = Some(e.into_iter());
153                }
154                Some(Ok(o)) => return Some(Ok(o)),
155            }
156        }
157    }
158
159    #[inline]
160    fn size_hint(&self) -> (usize, Option<usize>) {
161        self.iter.size_hint()
162    }
163}