resiter/
filter_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 and map Oks and Errors.
8pub trait FilterMap<O, E>: Sized {
9    fn filter_map_ok<F, O2>(self, _: F) -> FilterMapOk<Self, F>
10    where
11        F: FnMut(O) -> Option<O2>;
12    fn filter_map_err<F, E2>(self, _: F) -> FilterMapErr<Self, F>
13    where
14        F: FnMut(E) -> Option<E2>;
15}
16
17impl<I, O, E> FilterMap<O, E> for I
18where
19    I: Iterator<Item = Result<O, E>> + Sized,
20{
21    /// `filter_map` every `Ok` value
22    ///
23    /// ```
24    /// use std::str::FromStr;
25    /// use resiter::filter_map::FilterMap;
26    ///
27    /// let filter_mapped: Vec<_> = vec![
28    ///     Ok("1"),
29    ///     Err("2"),
30    ///     Ok("a"),
31    ///     Err("4"),
32    ///     Ok("5"),
33    ///     Err("b"),
34    ///     Err("8"),
35    /// ].into_iter()
36    ///     .filter_map_ok(|txt| usize::from_str(txt).ok())
37    ///     .collect();
38    /// assert_eq!(filter_mapped[0], Ok(1));
39    /// assert_eq!(filter_mapped[1], Err("2"));
40    /// assert_eq!(filter_mapped[2], Err("4"));
41    /// assert_eq!(filter_mapped[3], Ok(5));
42    /// assert_eq!(filter_mapped[4], Err("b"));
43    /// assert_eq!(filter_mapped[5], Err("8"));
44    /// ```
45    #[inline]
46    fn filter_map_ok<F, O2>(self, f: F) -> FilterMapOk<Self, F>
47    where
48        F: FnMut(O) -> Option<O2>,
49    {
50        FilterMapOk { iter: self, f }
51    }
52
53    /// `filter_map` every `Err(v)`
54    ///
55    /// ```
56    /// use std::str::FromStr;
57    /// use resiter::filter_map::FilterMap;
58    ///
59    /// let filter_mapped: Vec<_> = vec![
60    ///     Ok("1"),
61    ///     Err("2"),
62    ///     Ok("a"),
63    ///     Err("4"),
64    ///     Ok("5"),
65    ///     Err("b"),
66    ///     Err("8"),
67    /// ].into_iter()
68    ///     .filter_map_err(|txt| usize::from_str(txt).ok())
69    ///     .collect();
70    ///
71    /// assert_eq!(filter_mapped[0], Ok("1"));
72    /// assert_eq!(filter_mapped[1], Err(2));
73    /// assert_eq!(filter_mapped[2], Ok("a"));
74    /// assert_eq!(filter_mapped[3], Err(4));
75    /// assert_eq!(filter_mapped[4], Ok("5"));
76    /// assert_eq!(filter_mapped[5], Err(8));
77    /// ```
78    #[inline]
79    fn filter_map_err<F, E2>(self, f: F) -> FilterMapErr<Self, F>
80    where
81        F: FnMut(E) -> Option<E2>,
82    {
83        FilterMapErr { iter: self, f }
84    }
85}
86
87#[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
88pub struct FilterMapOk<I, F> {
89    iter: I,
90    f: F,
91}
92
93impl<I, O, E, F, O2> Iterator for FilterMapOk<I, F>
94where
95    I: Iterator<Item = Result<O, E>>,
96    F: FnMut(O) -> Option<O2>,
97{
98    type Item = Result<O2, E>;
99
100    fn next(&mut self) -> Option<Self::Item> {
101        loop {
102            match self.iter.next() {
103                Some(Ok(x)) => {
104                    if let Some(x) = (self.f)(x) {
105                        return Some(Ok(x));
106                    }
107                }
108                Some(Err(e)) => return Some(Err(e)),
109                None => return None,
110            }
111        }
112    }
113
114    #[inline]
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 FilterMapErr<I, F> {
122    iter: I,
123    f: F,
124}
125
126impl<I, O, E, F, E2> Iterator for FilterMapErr<I, F>
127where
128    I: Iterator<Item = Result<O, E>>,
129    F: FnMut(E) -> Option<E2>,
130{
131    type Item = Result<O, E2>;
132
133    fn next(&mut self) -> Option<Self::Item> {
134        loop {
135            match self.iter.next() {
136                Some(Ok(x)) => return Some(Ok(x)),
137                Some(Err(e)) => {
138                    if let Some(e) = (self.f)(e) {
139                        return Some(Err(e));
140                    }
141                }
142                None => return None,
143            }
144        }
145    }
146    #[inline]
147    fn size_hint(&self) -> (usize, Option<usize>) {
148        self.iter.size_hint()
149    }
150}
151
152#[cfg(test)]
153mod tests {
154    use super::*;
155
156    #[test]
157    fn test_filter_map_ok_hint() {
158        use std::str::FromStr;
159
160        let hint = ["1", "2", "a", "4", "5"]
161            .iter()
162            .map(|txt| usize::from_str(txt))
163            .filter_map_ok(|i| Some(2 * i))
164            .size_hint();
165
166        assert_eq!(hint, (5, Some(5)));
167    }
168
169    #[test]
170    fn test_filter_map_err_hint() {
171        use std::str::FromStr;
172
173        let hint = ["1", "2", "a", "4", "5"]
174            .iter()
175            .map(|txt| usize::from_str(txt))
176            .filter_map_err(|e| Some(format!("{:?}", e)))
177            .size_hint();
178
179        assert_eq!(hint, (5, Some(5)));
180    }
181}