resiter/
flat_map.rs

Help
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 FlatMap<O, E>: Sized {
9    /// [flat_map](Iterator::flat_map) every `Ok` value and leave all `Err` as is
10    ///
11    /// ```
12    /// use resiter::flat_map::FlatMap;
13    ///
14    /// let mapped: Vec<_> = vec![Ok(1), Ok(2), Err(2), Err(0), Ok(2)]
15    ///     .into_iter()
16    ///     .flat_map_ok(|i| (0..i))
17    ///     .collect();
18    /// assert_eq!(mapped, [Ok(0), Ok(0), Ok(1), Err(2), Err(0), Ok(0), Ok(1)]);
19    /// ```
20    fn flat_map_ok<U, F, O2>(self, _: F) -> FlatMapOk<Self, U, F>
21    where
22        F: FnMut(O) -> U,
23        U: IntoIterator<Item = O2>;
24
25    /// [flat_map](Iterator::flat_map) every `Err` value and leave all `Ok` as is
26    ///
27    /// ```
28    /// use resiter::flat_map::FlatMap;
29    ///
30    /// let mapped: Vec<_> = vec![Ok(1), Ok(2), Err(2), Err(0), Ok(2)]
31    ///     .into_iter()
32    ///     .flat_map_err(|i| 0..(i * 2))
33    ///     .collect();
34    ///
35    /// assert_eq!(
36    ///     mapped,
37    ///     [Ok(1), Ok(2), Err(0), Err(1), Err(2), Err(3), Ok(2)]
38    /// );
39    /// ```
40    fn flat_map_err<U, F, E2>(self, _: F) -> FlatMapErr<Self, U, F>
41    where
42        F: FnMut(E) -> U,
43        U: IntoIterator<Item = E2>;
44}
45
46impl<I, O, E> FlatMap<O, E> for I
47where
48    I: Iterator<Item = Result<O, E>> + Sized,
49{
50    #[inline]
51    fn flat_map_ok<U, F, O2>(self, f: F) -> FlatMapOk<Self, U, F>
52    where
53        F: FnMut(O) -> U,
54        U: IntoIterator<Item = O2>,
55    {
56        FlatMapOk {
57            frontiter: None,
58            iter: self,
59            f,
60        }
61    }
62
63    #[inline]
64    fn flat_map_err<U, F, E2>(self, f: F) -> FlatMapErr<Self, U, F>
65    where
66        F: FnMut(E) -> U,
67        U: IntoIterator<Item = E2>,
68    {
69        FlatMapErr {
70            frontiter: None,
71            iter: self,
72            f,
73        }
74    }
75}
76
77#[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
78pub struct FlatMapOk<I, U, F>
79where
80    U: IntoIterator,
81{
82    frontiter: Option<<U as IntoIterator>::IntoIter>,
83    iter: I,
84    f: F,
85}
86
87impl<I, O, E, F, O2, U> Iterator for FlatMapOk<I, U, F>
88where
89    I: Iterator<Item = Result<O, E>>,
90    F: FnMut(O) -> U,
91    U: IntoIterator<Item = O2>,
92{
93    type Item = Result<O2, E>;
94
95    fn next(&mut self) -> Option<Self::Item> {
96        loop {
97            if let Some(ref mut inner) = self.frontiter {
98                if let elt @ Some(_) = inner.next() {
99                    return elt.map(Ok);
100                }
101            }
102            match self.iter.next() {
103                None => return None,
104                Some(Ok(x)) => {
105                    self.frontiter = Some((self.f)(x).into_iter());
106                }
107                Some(Err(e)) => return Some(Err(e)),
108            }
109        }
110    }
111
112    #[inline]
113    // TODO: Oh dear, this hint could be much better
114    // https://doc.rust-lang.org/src/core/iter/mod.rs.html#2694
115    fn size_hint(&self) -> (usize, Option<usize>) {
116        self.iter.size_hint()
117    }
118}
119
120#[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
121pub struct FlatMapErr<I, U: IntoIterator, F> {
122    frontiter: Option<<U as IntoIterator>::IntoIter>,
123    iter: I,
124    f: F,
125}
126
127impl<I, O, E, F, E2, U> Iterator for FlatMapErr<I, U, F>
128where
129    I: Iterator<Item = Result<O, E>>,
130    F: FnMut(E) -> U,
131    U: IntoIterator<Item = E2>,
132{
133    type Item = Result<O, E2>;
134
135    fn next(&mut self) -> Option<Self::Item> {
136        loop {
137            if let Some(ref mut inner) = self.frontiter {
138                if let elt @ Some(_) = inner.next() {
139                    return elt.map(Err);
140                }
141            }
142            match self.iter.next() {
143                None => return None,
144                Some(Err(e)) => {
145                    self.frontiter = Some((self.f)(e).into_iter());
146                }
147                Some(Ok(o)) => return Some(Ok(o)),
148            }
149        }
150    }
151
152    #[inline]
153    fn size_hint(&self) -> (usize, Option<usize>) {
154        self.iter.size_hint()
155    }
156}
resiter/flat_map.rs

resiter/
flat_map.rs
