bevy_tasks/iter/
adapters.rs

1use crate::iter::ParallelIterator;
2
3#[derive(Debug)]
4pub struct Chain<T, U> {
5    pub(crate) left: T,
6    pub(crate) right: U,
7    pub(crate) left_in_progress: bool,
8}
9
10impl<B, T, U> ParallelIterator<B> for Chain<T, U>
11where
12    B: Iterator + Send,
13    T: ParallelIterator<B>,
14    U: ParallelIterator<B>,
15{
16    fn next_batch(&mut self) -> Option<B> {
17        if self.left_in_progress {
18            match self.left.next_batch() {
19                b @ Some(_) => return b,
20                None => self.left_in_progress = false,
21            }
22        }
23        self.right.next_batch()
24    }
25}
26
27#[derive(Debug)]
28pub struct Map<P, F> {
29    pub(crate) iter: P,
30    pub(crate) f: F,
31}
32
33impl<B, U, T, F> ParallelIterator<std::iter::Map<B, F>> for Map<U, F>
34where
35    B: Iterator + Send,
36    U: ParallelIterator<B>,
37    F: FnMut(B::Item) -> T + Send + Clone,
38{
39    fn next_batch(&mut self) -> Option<std::iter::Map<B, F>> {
40        self.iter.next_batch().map(|b| b.map(self.f.clone()))
41    }
42}
43
44#[derive(Debug)]
45pub struct Filter<P, F> {
46    pub(crate) iter: P,
47    pub(crate) predicate: F,
48}
49
50impl<B, P, F> ParallelIterator<std::iter::Filter<B, F>> for Filter<P, F>
51where
52    B: Iterator + Send,
53    P: ParallelIterator<B>,
54    F: FnMut(&B::Item) -> bool + Send + Clone,
55{
56    fn next_batch(&mut self) -> Option<std::iter::Filter<B, F>> {
57        self.iter
58            .next_batch()
59            .map(|b| b.filter(self.predicate.clone()))
60    }
61}
62
63#[derive(Debug)]
64pub struct FilterMap<P, F> {
65    pub(crate) iter: P,
66    pub(crate) f: F,
67}
68
69impl<B, P, R, F> ParallelIterator<std::iter::FilterMap<B, F>> for FilterMap<P, F>
70where
71    B: Iterator + Send,
72    P: ParallelIterator<B>,
73    F: FnMut(B::Item) -> Option<R> + Send + Clone,
74{
75    fn next_batch(&mut self) -> Option<std::iter::FilterMap<B, F>> {
76        self.iter.next_batch().map(|b| b.filter_map(self.f.clone()))
77    }
78}
79
80#[derive(Debug)]
81pub struct FlatMap<P, F> {
82    pub(crate) iter: P,
83    pub(crate) f: F,
84}
85
86impl<B, P, U, F> ParallelIterator<std::iter::FlatMap<B, U, F>> for FlatMap<P, F>
87where
88    B: Iterator + Send,
89    P: ParallelIterator<B>,
90    F: FnMut(B::Item) -> U + Send + Clone,
91    U: IntoIterator,
92    U::IntoIter: Send,
93{
94    // This extends each batch using the flat map. The other option is
95    // to turn each IntoIter into its own batch.
96    fn next_batch(&mut self) -> Option<std::iter::FlatMap<B, U, F>> {
97        self.iter.next_batch().map(|b| b.flat_map(self.f.clone()))
98    }
99}
100
101#[derive(Debug)]
102pub struct Flatten<P> {
103    pub(crate) iter: P,
104}
105
106impl<B, P> ParallelIterator<std::iter::Flatten<B>> for Flatten<P>
107where
108    B: Iterator + Send,
109    P: ParallelIterator<B>,
110    B::Item: IntoIterator,
111    <B::Item as IntoIterator>::IntoIter: Send,
112{
113    // This extends each batch using the flatten. The other option is to
114    // turn each IntoIter into its own batch.
115    fn next_batch(&mut self) -> Option<std::iter::Flatten<B>> {
116        self.iter.next_batch().map(|b| b.flatten())
117    }
118}
119
120#[derive(Debug)]
121pub struct Fuse<P> {
122    pub(crate) iter: Option<P>,
123}
124
125impl<B, P> ParallelIterator<B> for Fuse<P>
126where
127    B: Iterator + Send,
128    P: ParallelIterator<B>,
129{
130    fn next_batch(&mut self) -> Option<B> {
131        match &mut self.iter {
132            Some(iter) => iter.next_batch().or_else(|| {
133                self.iter = None;
134                None
135            }),
136            None => None,
137        }
138    }
139}
140
141#[derive(Debug)]
142pub struct Inspect<P, F> {
143    pub(crate) iter: P,
144    pub(crate) f: F,
145}
146
147impl<B, P, F> ParallelIterator<std::iter::Inspect<B, F>> for Inspect<P, F>
148where
149    B: Iterator + Send,
150    P: ParallelIterator<B>,
151    F: FnMut(&B::Item) + Send + Clone,
152{
153    fn next_batch(&mut self) -> Option<std::iter::Inspect<B, F>> {
154        self.iter.next_batch().map(|b| b.inspect(self.f.clone()))
155    }
156}
157
158#[derive(Debug)]
159pub struct Copied<P> {
160    pub(crate) iter: P,
161}
162
163impl<'a, B, P, T> ParallelIterator<std::iter::Copied<B>> for Copied<P>
164where
165    B: Iterator<Item = &'a T> + Send,
166    P: ParallelIterator<B>,
167    T: 'a + Copy,
168{
169    fn next_batch(&mut self) -> Option<std::iter::Copied<B>> {
170        self.iter.next_batch().map(|b| b.copied())
171    }
172}
173
174#[derive(Debug)]
175pub struct Cloned<P> {
176    pub(crate) iter: P,
177}
178
179impl<'a, B, P, T> ParallelIterator<std::iter::Cloned<B>> for Cloned<P>
180where
181    B: Iterator<Item = &'a T> + Send,
182    P: ParallelIterator<B>,
183    T: 'a + Copy,
184{
185    fn next_batch(&mut self) -> Option<std::iter::Cloned<B>> {
186        self.iter.next_batch().map(|b| b.cloned())
187    }
188}
189
190#[derive(Debug)]
191pub struct Cycle<P> {
192    pub(crate) iter: P,
193    pub(crate) curr: Option<P>,
194}
195
196impl<B, P> ParallelIterator<B> for Cycle<P>
197where
198    B: Iterator + Send,
199    P: ParallelIterator<B> + Clone,
200{
201    fn next_batch(&mut self) -> Option<B> {
202        self.curr.as_mut().and_then(|c| c.next_batch()).or_else(|| {
203            self.curr = Some(self.iter.clone());
204            self.next_batch()
205        })
206    }
207}