image/
buffer_par.rs

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/// Parallel iterator over pixel refs.
11#[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
75/// Parallel iterator over mutable pixel refs.
76pub 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/// Parallel iterator over pixel refs and their coordinates.
140#[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
227/// Parallel iterator over mutable pixel refs and their coordinates.
228pub 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    /// Returns a parallel iterator over the pixels of this image, usable with `rayon`.
321    /// See [`pixels`] for more information.
322    ///
323    /// [`pixels`]: #method.pixels
324    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    /// Returns a parallel iterator over the pixels of this image and their coordinates, usable with `rayon`.
333    /// See [`enumerate_pixels`] for more information.
334    ///
335    /// [`enumerate_pixels`]: #method.enumerate_pixels
336    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    /// Returns a parallel iterator over the mutable pixels of this image, usable with `rayon`.
351    /// See [`pixels_mut`] for more information.
352    ///
353    /// [`pixels_mut`]: #method.pixels_mut
354    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    /// Returns a parallel iterator over the mutable pixels of this image and their coordinates, usable with `rayon`.
363    /// See [`enumerate_pixels_mut`] for more information.
364    ///
365    /// [`enumerate_pixels_mut`]: #method.enumerate_pixels_mut
366    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    /// Constructs a new `ImageBuffer` by repeated application of the supplied function,
381    /// utilizing multi-threading via `rayon`.
382    ///
383    /// The arguments to the function are the pixel's x and y coordinates.
384    ///
385    /// # Panics
386    ///
387    /// Panics when the resulting image is larger than the maximum size of a vector.
388    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}