iota_common/
try_iterator_ext.rs

1// Copyright (c) 2025 IOTA Stiftung
2// SPDX-License-Identifier: Apache-2.0
3
4pub trait TryIteratorExt<T, E>: Iterator<Item = Result<T, E>> + Sized {
5    /// Tries to map and collect items from a fallible iterator, stopping early
6    /// on the first error.
7    fn try_map_while_and_collect<U, P, B>(self, mut predicate: P) -> B
8    where
9        P: FnMut(T) -> Option<U>,
10        B: FromIterator<Result<U, E>>,
11    {
12        FromIterator::from_iter(self.map_while(|result| result.map(&mut predicate).transpose()))
13    }
14
15    /// Try taking at most limit items while predicate holds collecting mapped
16    /// values and breaking early on errors.
17    fn try_take_map_while_and_collect<U, F, P, B>(
18        self,
19        limit: Option<usize>,
20        mut predicate: P,
21        map_fn: F,
22    ) -> Result<B, E>
23    where
24        F: Fn(T) -> U,
25        P: FnMut(&T) -> bool,
26        B: FromIterator<U>,
27    {
28        let predicate = |v| predicate(&v).then(|| map_fn(v));
29
30        if let Some(limit) = limit {
31            self.take(limit).try_map_while_and_collect(predicate)
32        } else {
33            self.try_map_while_and_collect(predicate)
34        }
35    }
36}
37
38impl<I, T, E> TryIteratorExt<T, E> for I where I: Iterator<Item = Result<T, E>> {}
39
40#[cfg(test)]
41mod tests {
42    use super::*;
43
44    #[test]
45    fn test_try_skip_filter_map_and_collect() {
46        let result: Result<Vec<_>, &str> = [1, 2, 3, 8]
47            .into_iter()
48            .map(Ok)
49            .chain(std::iter::from_fn(|| panic!()))
50            .try_take_map_while_and_collect(Some(5), |&x: &i32| x < 8, |x| x * 2);
51        assert_eq!(result, Ok(vec![2, 4, 6])); // stops at 8
52    }
53
54    #[test]
55    fn test_try_skip_filter_map_and_collect_with_error() {
56        let result: Result<Vec<_>, _> = [Ok(1), Ok(2), Err("error")]
57            .into_iter()
58            .chain(std::iter::from_fn(|| panic!()))
59            .try_take_map_while_and_collect(None, |&x: &i32| x < 8, |x| x * 2);
60        assert_eq!(result, Err("error")); // stops on the first error
61    }
62
63    #[test]
64    fn test_try_skip_filter_map_and_collect_with_limit() {
65        let result: Result<Vec<_>, &str> = [Ok(1), Ok(2)]
66            .into_iter()
67            .chain(std::iter::from_fn(|| panic!()))
68            .try_take_map_while_and_collect(Some(2), |&x: &i32| x < 8, |x| x * 2);
69        assert_eq!(result, Ok(vec![2, 4])); // respects limit stopping before panic
70    }
71}