glance_core/img/
iterators.rs

1use rayon::{
2    iter::{IndexedParallelIterator, ParallelIterator},
3    slice::{ParallelSlice, ParallelSliceMut},
4};
5
6use super::Image;
7
8pub struct PixelIter<'a> {
9    image: &'a Image,
10    idx: usize,
11}
12
13impl<'a> Iterator for PixelIter<'a> {
14    type Item = (usize, usize, [u8; 4]);
15    fn next(&mut self) -> Option<Self::Item> {
16        if self.idx >= (self.image.width * self.image.height) as usize {
17            return None;
18        }
19        let x = self.idx % self.image.width as usize;
20        let y = self.idx / self.image.width as usize;
21        let base = self.idx * 4;
22        let px = [
23            self.image.data[base],
24            self.image.data[base + 1],
25            self.image.data[base + 2],
26            self.image.data[base + 3],
27        ];
28        self.idx += 1;
29        Some((x, y, px))
30    }
31}
32
33pub struct PixelIterMut<'a> {
34    image: &'a mut Image,
35    idx: usize,
36}
37
38impl<'a> Iterator for PixelIterMut<'a> {
39    type Item = (usize, usize, &'a mut [u8; 4]);
40    fn next(&mut self) -> Option<Self::Item> {
41        let width = self.image.width as usize;
42        let height = self.image.height as usize;
43        if self.idx >= width * height {
44            return None;
45        }
46        let x = self.idx % width;
47        let y = self.idx / width;
48        let base = self.idx * 4;
49
50        // Each pixel here is only yielded once. So this is safe.
51        let pixel = unsafe {
52            let ptr = self.image.data.as_mut_ptr().add(base) as *mut [u8; 4];
53            &mut *ptr
54        };
55
56        self.idx += 1;
57        Some((x, y, pixel))
58    }
59}
60
61impl Image {
62    pub fn pixels(&self) -> PixelIter<'_> {
63        PixelIter {
64            image: self,
65            idx: 0,
66        }
67    }
68
69    pub fn pixels_mut(&mut self) -> PixelIterMut<'_> {
70        PixelIterMut {
71            image: self,
72            idx: 0,
73        }
74    }
75
76    pub fn par_pixels(&self) -> impl ParallelIterator<Item = (usize, usize, [u8; 4])> {
77        let width = self.width as usize;
78        self.data
79            .par_chunks_exact(4)
80            .enumerate()
81            .map(move |(idx, chunk)| {
82                let x = idx % width;
83                let y = idx / width;
84                let pixel = [chunk[0], chunk[1], chunk[2], chunk[3]];
85                (x, y, pixel)
86            })
87    }
88
89    pub fn par_pixels_mut(&mut self) -> impl ParallelIterator<Item = (usize, usize, &mut [u8; 4])> {
90        let width = self.width as usize;
91
92        self.data
93            .par_chunks_exact_mut(4)
94            .enumerate()
95            .map(move |(idx, chunk)| {
96                let x = idx % width;
97                let y = idx / width;
98
99                // Chunk is exactly 4 bytes (from par_chunks_exact_mut),
100                // and Rayon ensures chunks are disjoint so it's safe to cast to &mut [u8; 4]
101                let pixel = unsafe { &mut *(chunk.as_mut_ptr() as *mut [u8; 4]) };
102
103                (x, y, pixel)
104            })
105    }
106}
glance_core/img/iterators.rs

glance_core/img/
iterators.rs
