Skip to main content

moirai_iter/parallel/adapters/
filter.rs

1use super::super::{Consumer, FilterConsumer, ParallelIterator, VecParIter, VecRefParIter};
2use super::flat::Flatten;
3use super::map::Map;
4use super::pair::ZipEq;
5use super::ref_ops::Copied;
6
7/// Filter adapter for parallel iterators.
8pub struct Filter<I, F> {
9    pub(super) base: I,
10    pub(super) filter_fn: F,
11}
12
13impl<I, F> Filter<I, F> {
14    pub(crate) fn new(base: I, filter_fn: F) -> Self {
15        Self { base, filter_fn }
16    }
17}
18
19impl<I, MapFn, FilterFn, Mapped> Filter<Map<Flatten<I>, MapFn>, FilterFn>
20where
21    I: ParallelIterator,
22    I::Item: IntoIterator,
23    MapFn: Fn(<I::Item as IntoIterator>::Item) -> Mapped + Send + Sync + Clone,
24    FilterFn: Fn(&Mapped) -> bool + Send + Sync + Clone,
25    Mapped: Send,
26{
27    /// Sum a flattened, mapped, filtered nested stream without intermediate vectors.
28    pub fn sum<S>(self) -> S
29    where
30        S: std::iter::Sum<Mapped> + Send,
31    {
32        let filter_fn = self.filter_fn;
33        let map = self.base;
34        let map_fn = map.map_fn;
35        let flatten = map.base;
36
37        flatten
38            .base
39            .seq_items()
40            .into_iter()
41            .flat_map(IntoIterator::into_iter)
42            .map(map_fn)
43            .filter(filter_fn)
44            .sum()
45    }
46}
47
48impl<I, J, MapFn, FilterFn, Mapped> Filter<Map<ZipEq<I, J>, MapFn>, FilterFn>
49where
50    I: ParallelIterator,
51    J: ParallelIterator,
52    I::Item: Sync + 'static,
53    J::Item: Sync + 'static,
54    MapFn: Fn((I::Item, J::Item)) -> Mapped + Send + Sync + Clone,
55    FilterFn: Fn(&Mapped) -> bool + Send + Sync + Clone,
56    Mapped: Send,
57{
58    /// Collect a zipped, mapped, filtered stream without intermediate pair vectors.
59    pub fn collect<C>(self) -> C
60    where
61        C: FromIterator<Mapped> + Send,
62    {
63        let filter_fn = self.filter_fn;
64        let map = self.base;
65        let map_fn = map.map_fn;
66        let zip = map.base;
67        let left = zip.left.seq_items();
68        let right = zip.right.seq_items();
69        assert_eq!(
70            left.len(),
71            right.len(),
72            "zip_eq requires equal input lengths"
73        );
74
75        left.into_iter()
76            .zip(right)
77            .map(map_fn)
78            .filter(filter_fn)
79            .collect()
80    }
81}
82
83impl<'data, T, MapFn, FilterFn, Mapped>
84    Filter<Map<Copied<VecRefParIter<'data, T>>, MapFn>, FilterFn>
85where
86    T: Copy + Send + Sync + 'data,
87    MapFn: Fn(T) -> Mapped + Send + Sync + Clone,
88    FilterFn: Fn(&Mapped) -> bool + Send + Sync + Clone,
89    Mapped: Send,
90{
91    /// Sum a borrowed copied-map-filter stream without materializing references.
92    pub fn sum<S>(self) -> S
93    where
94        S: std::iter::Sum<Mapped> + Send,
95    {
96        let filter_fn = self.filter_fn;
97        let map = self.base;
98        let map_fn = map.map_fn;
99        let copied = map.base;
100
101        copied
102            .base
103            .into_slice()
104            .iter()
105            .copied()
106            .map(map_fn)
107            .filter(filter_fn)
108            .sum()
109    }
110}
111
112impl<I, F> ParallelIterator for Filter<I, F>
113where
114    I: ParallelIterator,
115    F: Fn(&I::Item) -> bool + Send + Sync + Clone,
116{
117    type Item = I::Item;
118
119    fn seq_items(self) -> Vec<Self::Item> {
120        self.base
121            .seq_items()
122            .into_iter()
123            .filter(|x| (self.filter_fn)(x))
124            .collect()
125    }
126
127    fn drive<C, R>(self, consumer: C) -> R
128    where
129        C: Consumer<Self::Item, Result = R> + Send + Sync,
130        R: Send,
131    {
132        self.base
133            .drive(FilterConsumer::new(consumer, self.filter_fn))
134    }
135}
136
137/// Filter-map adapter with value-semantic optional output.
138pub struct FilterMap<I, F> {
139    pub(super) base: I,
140    pub(super) filter_map_fn: F,
141}
142
143impl<I, F> FilterMap<I, F> {
144    pub(crate) fn new(base: I, filter_map_fn: F) -> Self {
145        Self {
146            base,
147            filter_map_fn,
148        }
149    }
150}
151
152impl<I, F, R> ParallelIterator for FilterMap<I, F>
153where
154    I: ParallelIterator,
155    F: Fn(I::Item) -> Option<R> + Send + Sync + Clone,
156    R: Send + Sync + 'static,
157{
158    type Item = R;
159
160    fn seq_items(self) -> Vec<Self::Item> {
161        self.base
162            .seq_items()
163            .into_iter()
164            .filter_map(self.filter_map_fn)
165            .collect()
166    }
167
168    fn drive<C, R2>(self, consumer: C) -> R2
169    where
170        C: Consumer<Self::Item, Result = R2> + Send + Sync,
171        R2: Send,
172    {
173        consumer.consume(VecParIter::new(self.seq_items()))
174    }
175}
176
177/// While-some adapter with prefix-unwrapping semantics for optional streams.
178pub struct WhileSome<I> {
179    pub(super) base: I,
180}
181
182impl<I> WhileSome<I> {
183    pub(crate) fn new(base: I) -> Self {
184        Self { base }
185    }
186}
187
188impl<I, T> ParallelIterator for WhileSome<I>
189where
190    I: ParallelIterator<Item = Option<T>>,
191    T: Send + Sync + 'static,
192{
193    type Item = T;
194
195    fn seq_items(self) -> Vec<Self::Item> {
196        self.base
197            .seq_items()
198            .into_iter()
199            .map_while(|item| item)
200            .collect()
201    }
202
203    fn drive<C, R>(self, consumer: C) -> R
204    where
205        C: Consumer<Self::Item, Result = R> + Send + Sync,
206        R: Send,
207    {
208        consumer.consume(VecParIter::new(self.seq_items()))
209    }
210}