flux_common/
iter.rs

1use std::ops::Try;
2
3use rustc_errors::ErrorGuaranteed;
4
5use crate::result::ErrorCollector;
6
7pub trait IterExt: Iterator {
8    fn try_collect_vec<T, E>(self) -> Result<Vec<T>, E>
9    where
10        Self: Sized + Iterator<Item = Result<T, E>>,
11    {
12        self.collect()
13    }
14
15    fn collect_errors<T, E, C>(self, collector: &mut C) -> CollectErrors<Self, C>
16    where
17        Self: Iterator<Item = Result<T, E>> + Sized,
18        C: ErrorCollector<E>,
19    {
20        CollectErrors { iter: self, collector }
21    }
22
23    fn try_collect_exhaust<T, V>(self) -> Result<V, ErrorGuaranteed>
24    where
25        V: FromIterator<T>,
26        Self: Iterator<Item = Result<T, ErrorGuaranteed>> + Sized,
27    {
28        let mut acc: Option<ErrorGuaranteed> = None;
29        let v = self.collect_errors(&mut acc).collect();
30        match acc {
31            Some(e) => Err(e),
32            None => Ok(v),
33        }
34    }
35
36    fn try_for_each_exhaust<T, F>(self, mut f: F) -> Result<(), ErrorGuaranteed>
37    where
38        Self: Iterator<Item = T> + Sized,
39        F: FnMut(T) -> Result<(), ErrorGuaranteed>,
40    {
41        let mut acc: Option<ErrorGuaranteed> = None;
42        for v in self {
43            if let Err(e) = f(v) {
44                acc = Some(e).or(acc);
45            }
46        }
47        match acc {
48            Some(e) => Err(e),
49            None => Ok(()),
50        }
51    }
52
53    fn map_take_while<F, R>(&mut self, f: F) -> MapTakeWhile<Self, F>
54    where
55        Self: Clone,
56        F: FnMut(&Self::Item) -> Option<R>,
57    {
58        MapTakeWhile { iter: self, f }
59    }
60}
61
62impl<I: ?Sized> IterExt for I where I: Iterator {}
63
64pub struct CollectErrors<'a, I, C> {
65    iter: I,
66    collector: &'a mut C,
67}
68
69impl<I, T, E, F> Iterator for CollectErrors<'_, I, F>
70where
71    I: Iterator<Item = Result<T, E>>,
72    F: ErrorCollector<E>,
73{
74    type Item = T;
75
76    fn next(&mut self) -> Option<Self::Item> {
77        self.find(|_| true)
78    }
79
80    fn try_fold<B, F2, R2>(&mut self, init: B, mut f: F2) -> R2
81    where
82        F2: FnMut(B, Self::Item) -> R2,
83        R2: Try<Output = B>,
84    {
85        self.iter.try_fold(init, |acc, x| {
86            match x {
87                Ok(x) => f(acc, x),
88                Err(e) => {
89                    self.collector.collect(e);
90                    try { acc }
91                }
92            }
93        })
94    }
95
96    fn fold<B, F2>(mut self, init: B, fold: F2) -> B
97    where
98        Self: Sized,
99        F2: FnMut(B, Self::Item) -> B,
100    {
101        #[inline]
102        fn ok<B, T>(mut f: impl FnMut(B, T) -> B) -> impl FnMut(B, T) -> Result<B, !> {
103            move |acc, x| Ok(f(acc, x))
104        }
105
106        self.try_fold(init, ok(fold)).unwrap()
107    }
108}
109
110pub struct MapTakeWhile<'a, I: 'a, F> {
111    iter: &'a mut I,
112    f: F,
113}
114
115impl<'a, I, F, R> Iterator for MapTakeWhile<'a, I, F>
116where
117    I: 'a + Iterator + Clone,
118    F: FnMut(&I::Item) -> Option<R>,
119{
120    type Item = R;
121
122    fn next(&mut self) -> Option<Self::Item> {
123        let old = self.iter.clone();
124        match self.iter.next() {
125            None => None,
126            Some(elt) => {
127                if let Some(elt) = (self.f)(&elt) {
128                    Some(elt)
129                } else {
130                    *self.iter = old;
131                    None
132                }
133            }
134        }
135    }
136}
flux_common/iter.rs

flux_common/
iter.rs
