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

resiter_dpc_tmp/
flat_map.rs
