bevy_tasks/iter/
adapters.rs

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

bevy_tasks/iter/
adapters.rs
