1use rayon::iter::plumbing::*;
2use rayon::iter::{IndexedParallelIterator, ParallelIterator};
3use rayon::slice::{ChunksExact, ChunksExactMut, ParallelSlice, ParallelSliceMut};
4use std::fmt;
5use std::ops::{Deref, DerefMut};
6
7use crate::traits::Pixel;
8use crate::ImageBuffer;
9
10#[derive(Clone)]
12pub struct PixelsPar<'a, P>
13where
14 P: Pixel + Sync + 'a,
15 P::Subpixel: Sync + 'a,
16{
17 chunks: ChunksExact<'a, P::Subpixel>,
18}
19
20impl<'a, P> ParallelIterator for PixelsPar<'a, P>
21where
22 P: Pixel + Sync + 'a,
23 P::Subpixel: Sync + 'a,
24{
25 type Item = &'a P;
26
27 fn drive_unindexed<C>(self, consumer: C) -> C::Result
28 where
29 C: UnindexedConsumer<Self::Item>,
30 {
31 self.chunks
32 .map(|v| <P as Pixel>::from_slice(v))
33 .drive_unindexed(consumer)
34 }
35
36 fn opt_len(&self) -> Option<usize> {
37 Some(self.len())
38 }
39}
40
41impl<'a, P> IndexedParallelIterator for PixelsPar<'a, P>
42where
43 P: Pixel + Sync + 'a,
44 P::Subpixel: Sync + 'a,
45{
46 fn drive<C: Consumer<Self::Item>>(self, consumer: C) -> C::Result {
47 self.chunks
48 .map(|v| <P as Pixel>::from_slice(v))
49 .drive(consumer)
50 }
51
52 fn len(&self) -> usize {
53 self.chunks.len()
54 }
55
56 fn with_producer<CB: ProducerCallback<Self::Item>>(self, callback: CB) -> CB::Output {
57 self.chunks
58 .map(|v| <P as Pixel>::from_slice(v))
59 .with_producer(callback)
60 }
61}
62
63impl<P> fmt::Debug for PixelsPar<'_, P>
64where
65 P: Pixel + Sync,
66 P::Subpixel: Sync + fmt::Debug,
67{
68 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
69 f.debug_struct("PixelsPar")
70 .field("chunks", &self.chunks)
71 .finish()
72 }
73}
74
75pub struct PixelsMutPar<'a, P>
77where
78 P: Pixel + Send + Sync + 'a,
79 P::Subpixel: Send + Sync + 'a,
80{
81 chunks: ChunksExactMut<'a, P::Subpixel>,
82}
83
84impl<'a, P> ParallelIterator for PixelsMutPar<'a, P>
85where
86 P: Pixel + Send + Sync + 'a,
87 P::Subpixel: Send + Sync + 'a,
88{
89 type Item = &'a mut P;
90
91 fn drive_unindexed<C>(self, consumer: C) -> C::Result
92 where
93 C: UnindexedConsumer<Self::Item>,
94 {
95 self.chunks
96 .map(|v| <P as Pixel>::from_slice_mut(v))
97 .drive_unindexed(consumer)
98 }
99
100 fn opt_len(&self) -> Option<usize> {
101 Some(self.len())
102 }
103}
104
105impl<'a, P> IndexedParallelIterator for PixelsMutPar<'a, P>
106where
107 P: Pixel + Send + Sync + 'a,
108 P::Subpixel: Send + Sync + 'a,
109{
110 fn drive<C: Consumer<Self::Item>>(self, consumer: C) -> C::Result {
111 self.chunks
112 .map(|v| <P as Pixel>::from_slice_mut(v))
113 .drive(consumer)
114 }
115
116 fn len(&self) -> usize {
117 self.chunks.len()
118 }
119
120 fn with_producer<CB: ProducerCallback<Self::Item>>(self, callback: CB) -> CB::Output {
121 self.chunks
122 .map(|v| <P as Pixel>::from_slice_mut(v))
123 .with_producer(callback)
124 }
125}
126
127impl<P> fmt::Debug for PixelsMutPar<'_, P>
128where
129 P: Pixel + Send + Sync,
130 P::Subpixel: Send + Sync + fmt::Debug,
131{
132 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
133 f.debug_struct("PixelsMutPar")
134 .field("chunks", &self.chunks)
135 .finish()
136 }
137}
138
139#[derive(Clone)]
141pub struct EnumeratePixelsPar<'a, P>
142where
143 P: Pixel + Sync + 'a,
144 P::Subpixel: Sync + 'a,
145{
146 pixels: PixelsPar<'a, P>,
147 width: u32,
148}
149
150impl<'a, P> ParallelIterator for EnumeratePixelsPar<'a, P>
151where
152 P: Pixel + Sync + 'a,
153 P::Subpixel: Sync + 'a,
154{
155 type Item = (u32, u32, &'a P);
156
157 fn drive_unindexed<C>(self, consumer: C) -> C::Result
158 where
159 C: UnindexedConsumer<Self::Item>,
160 {
161 self.pixels
162 .enumerate()
163 .map(|(i, p)| {
164 (
165 (i % self.width as usize) as u32,
166 (i / self.width as usize) as u32,
167 p,
168 )
169 })
170 .drive_unindexed(consumer)
171 }
172
173 fn opt_len(&self) -> Option<usize> {
174 Some(self.len())
175 }
176}
177
178impl<'a, P> IndexedParallelIterator for EnumeratePixelsPar<'a, P>
179where
180 P: Pixel + Sync + 'a,
181 P::Subpixel: Sync + 'a,
182{
183 fn drive<C: Consumer<Self::Item>>(self, consumer: C) -> C::Result {
184 self.pixels
185 .enumerate()
186 .map(|(i, p)| {
187 (
188 (i % self.width as usize) as u32,
189 (i / self.width as usize) as u32,
190 p,
191 )
192 })
193 .drive(consumer)
194 }
195
196 fn len(&self) -> usize {
197 self.pixels.len()
198 }
199
200 fn with_producer<CB: ProducerCallback<Self::Item>>(self, callback: CB) -> CB::Output {
201 self.pixels
202 .enumerate()
203 .map(|(i, p)| {
204 (
205 (i % self.width as usize) as u32,
206 (i / self.width as usize) as u32,
207 p,
208 )
209 })
210 .with_producer(callback)
211 }
212}
213
214impl<P> fmt::Debug for EnumeratePixelsPar<'_, P>
215where
216 P: Pixel + Sync,
217 P::Subpixel: Sync + fmt::Debug,
218{
219 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
220 f.debug_struct("EnumeratePixelsPar")
221 .field("pixels", &self.pixels)
222 .field("width", &self.width)
223 .finish()
224 }
225}
226
227pub struct EnumeratePixelsMutPar<'a, P>
229where
230 P: Pixel + Send + Sync + 'a,
231 P::Subpixel: Send + Sync + 'a,
232{
233 pixels: PixelsMutPar<'a, P>,
234 width: u32,
235}
236
237impl<'a, P> ParallelIterator for EnumeratePixelsMutPar<'a, P>
238where
239 P: Pixel + Send + Sync + 'a,
240 P::Subpixel: Send + Sync + 'a,
241{
242 type Item = (u32, u32, &'a mut P);
243
244 fn drive_unindexed<C>(self, consumer: C) -> C::Result
245 where
246 C: UnindexedConsumer<Self::Item>,
247 {
248 self.pixels
249 .enumerate()
250 .map(|(i, p)| {
251 (
252 (i % self.width as usize) as u32,
253 (i / self.width as usize) as u32,
254 p,
255 )
256 })
257 .drive_unindexed(consumer)
258 }
259
260 fn opt_len(&self) -> Option<usize> {
261 Some(self.len())
262 }
263}
264
265impl<'a, P> IndexedParallelIterator for EnumeratePixelsMutPar<'a, P>
266where
267 P: Pixel + Send + Sync + 'a,
268 P::Subpixel: Send + Sync + 'a,
269{
270 fn drive<C: Consumer<Self::Item>>(self, consumer: C) -> C::Result {
271 self.pixels
272 .enumerate()
273 .map(|(i, p)| {
274 (
275 (i % self.width as usize) as u32,
276 (i / self.width as usize) as u32,
277 p,
278 )
279 })
280 .drive(consumer)
281 }
282
283 fn len(&self) -> usize {
284 self.pixels.len()
285 }
286
287 fn with_producer<CB: ProducerCallback<Self::Item>>(self, callback: CB) -> CB::Output {
288 self.pixels
289 .enumerate()
290 .map(|(i, p)| {
291 (
292 (i % self.width as usize) as u32,
293 (i / self.width as usize) as u32,
294 p,
295 )
296 })
297 .with_producer(callback)
298 }
299}
300
301impl<P> fmt::Debug for EnumeratePixelsMutPar<'_, P>
302where
303 P: Pixel + Send + Sync,
304 P::Subpixel: Send + Sync + fmt::Debug,
305{
306 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
307 f.debug_struct("EnumeratePixelsMutPar")
308 .field("pixels", &self.pixels)
309 .field("width", &self.width)
310 .finish()
311 }
312}
313
314impl<P, Container> ImageBuffer<P, Container>
315where
316 P: Pixel + Sync,
317 P::Subpixel: Sync,
318 Container: Deref<Target = [P::Subpixel]>,
319{
320 pub fn par_pixels(&self) -> PixelsPar<P> {
325 PixelsPar {
326 chunks: self
327 .inner_pixels()
328 .par_chunks_exact(<P as Pixel>::CHANNEL_COUNT as usize),
329 }
330 }
331
332 pub fn par_enumerate_pixels(&self) -> EnumeratePixelsPar<P> {
337 EnumeratePixelsPar {
338 pixels: self.par_pixels(),
339 width: self.width(),
340 }
341 }
342}
343
344impl<P, Container> ImageBuffer<P, Container>
345where
346 P: Pixel + Send + Sync,
347 P::Subpixel: Send + Sync,
348 Container: Deref<Target = [P::Subpixel]> + DerefMut,
349{
350 pub fn par_pixels_mut(&mut self) -> PixelsMutPar<P> {
355 PixelsMutPar {
356 chunks: self
357 .inner_pixels_mut()
358 .par_chunks_exact_mut(<P as Pixel>::CHANNEL_COUNT as usize),
359 }
360 }
361
362 pub fn par_enumerate_pixels_mut(&mut self) -> EnumeratePixelsMutPar<P> {
367 let width = self.width();
368 EnumeratePixelsMutPar {
369 pixels: self.par_pixels_mut(),
370 width,
371 }
372 }
373}
374
375impl<P> ImageBuffer<P, Vec<P::Subpixel>>
376where
377 P: Pixel + Send + Sync,
378 P::Subpixel: Send + Sync,
379{
380 pub fn from_par_fn<F>(width: u32, height: u32, f: F) -> ImageBuffer<P, Vec<P::Subpixel>>
389 where
390 F: Fn(u32, u32) -> P + Send + Sync,
391 {
392 let mut buf = ImageBuffer::new(width, height);
393 buf.par_enumerate_pixels_mut().for_each(|(x, y, p)| {
394 *p = f(x, y);
395 });
396
397 buf
398 }
399}
400
401#[cfg(test)]
402mod test {
403 use crate::{Rgb, RgbImage};
404 use rayon::iter::{IndexedParallelIterator, ParallelIterator};
405
406 fn test_width_height(width: u32, height: u32, len: usize) {
407 let mut image = RgbImage::new(width, height);
408
409 assert_eq!(image.par_enumerate_pixels_mut().len(), len);
410 assert_eq!(image.par_enumerate_pixels().len(), len);
411 assert_eq!(image.par_pixels_mut().len(), len);
412 assert_eq!(image.par_pixels().len(), len);
413 }
414
415 #[test]
416 fn zero_width_zero_height() {
417 test_width_height(0, 0, 0);
418 }
419
420 #[test]
421 fn zero_width_nonzero_height() {
422 test_width_height(0, 2, 0);
423 }
424
425 #[test]
426 fn nonzero_width_zero_height() {
427 test_width_height(2, 0, 0);
428 }
429
430 #[test]
431 fn iter_parity() {
432 let mut image1 = RgbImage::from_fn(17, 29, |x, y| {
433 Rgb(std::array::from_fn(|i| {
434 ((x + y * 98 + i as u32 * 27) % 255) as u8
435 }))
436 });
437 let mut image2 = image1.clone();
438
439 assert_eq!(
440 image1.enumerate_pixels_mut().collect::<Vec<_>>(),
441 image2.par_enumerate_pixels_mut().collect::<Vec<_>>()
442 );
443 assert_eq!(
444 image1.enumerate_pixels().collect::<Vec<_>>(),
445 image2.par_enumerate_pixels().collect::<Vec<_>>()
446 );
447 assert_eq!(
448 image1.pixels_mut().collect::<Vec<_>>(),
449 image2.par_pixels_mut().collect::<Vec<_>>()
450 );
451 assert_eq!(
452 image1.pixels().collect::<Vec<_>>(),
453 image2.par_pixels().collect::<Vec<_>>()
454 );
455 }
456}
457
458#[cfg(test)]
459#[cfg(feature = "benchmarks")]
460mod benchmarks {
461 use crate::{Rgb, RgbImage};
462
463 const S: u32 = 1024;
464
465 #[bench]
466 fn creation(b: &mut test::Bencher) {
467 let mut bytes = 0;
468 b.iter(|| {
469 let img = RgbImage::from_fn(S, S, |_, _| test::black_box(pixel_func()));
470
471 bytes += img.as_raw().len() as u64;
472 });
473
474 b.bytes = bytes;
475 }
476
477 #[bench]
478 fn creation_par(b: &mut test::Bencher) {
479 let mut bytes = 0;
480 b.iter(|| {
481 let img = RgbImage::from_par_fn(S, S, |_, _| test::black_box(pixel_func()));
482
483 bytes += img.as_raw().len() as u64;
484 });
485
486 b.bytes = bytes;
487 }
488
489 fn pixel_func() -> Rgb<u8> {
490 use std::collections::hash_map::RandomState;
491 use std::hash::{BuildHasher, Hasher};
492 Rgb(std::array::from_fn(|_| {
493 RandomState::new().build_hasher().finish() as u8
494 }))
495 }
496}