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