imgref/
iter.rs

1use core::iter::FusedIterator;
2use core::marker::PhantomData;
3use core::num::NonZeroUsize;
4use core::slice;
5
6#[cfg(test)]
7use alloc::vec;
8
9/// Rows of the image. Call `Img.rows()` to create it.
10///
11/// Each element is a slice `width` pixels wide. Ignores padding, if there's any.
12#[derive(Debug)]
13#[must_use]
14pub struct RowsIter<'a, T> {
15    pub(crate) inner: slice::Chunks<'a, T>,
16    pub(crate) width: usize,
17}
18
19impl<'a, T: 'a> Iterator for RowsIter<'a, T> {
20    type Item = &'a [T];
21
22    #[inline]
23    fn next(&mut self) -> Option<Self::Item> {
24        match self.inner.next() {
25            Some(s) => {
26                // guaranteed during creation of chunks iterator
27                debug_assert!(s.len() >= self.width);
28                unsafe {
29                    Some(s.get_unchecked(0..self.width))
30                }
31            },
32            None => None,
33        }
34    }
35
36    #[inline]
37    fn size_hint(&self) -> (usize, Option<usize>) {
38        self.inner.size_hint()
39    }
40
41    #[inline]
42    fn nth(&mut self, n: usize) -> Option<Self::Item> {
43        match self.inner.nth(n) {
44            Some(s) => {
45                // guaranteed during creation of chunks iterator
46                debug_assert!(s.len() >= self.width);
47                unsafe {
48                    Some(s.get_unchecked(0..self.width))
49                }
50            },
51            None => None,
52        }
53    }
54
55    #[inline]
56    fn count(self) -> usize {
57        self.inner.count()
58    }
59}
60
61impl<T> ExactSizeIterator for RowsIter<'_, T> {
62    #[inline]
63    fn len(&self) -> usize {
64        self.inner.len()
65    }
66}
67
68impl<T> FusedIterator for RowsIter<'_, T> {}
69
70impl<'a, T: 'a> DoubleEndedIterator for RowsIter<'a, T> {
71    #[inline]
72    fn next_back(&mut self) -> Option<Self::Item> {
73        match self.inner.next_back() {
74            Some(s) => {
75                // guaranteed during creation of chunks iterator
76                debug_assert!(s.len() >= self.width);
77                unsafe {
78                    Some(s.get_unchecked(0..self.width))
79                }
80            },
81            None => None,
82        }
83    }
84}
85
86/// Rows of the image. Call `Img.rows_mut()` to create it.
87///
88/// Each element is a slice `width` pixels wide. Ignores padding, if there's any.
89#[derive(Debug)]
90#[must_use]
91pub struct RowsIterMut<'a, T> {
92    pub(crate) width: usize,
93    pub(crate) inner: slice::ChunksMut<'a, T>,
94}
95
96impl<'a, T: 'a> Iterator for RowsIterMut<'a, T> {
97    type Item = &'a mut [T];
98
99    #[inline]
100    fn next(&mut self) -> Option<Self::Item> {
101        match self.inner.next() {
102            Some(s) => Some(&mut s[0..self.width]),
103            None => None,
104        }
105    }
106
107    #[inline]
108    fn size_hint(&self) -> (usize, Option<usize>) {
109        self.inner.size_hint()
110    }
111
112    #[inline]
113    fn nth(&mut self, n: usize) -> Option<Self::Item> {
114        match self.inner.nth(n) {
115            Some(s) => Some(&mut s[0..self.width]),
116            None => None,
117        }
118    }
119
120    #[inline]
121    fn count(self) -> usize {
122        self.inner.count()
123    }
124}
125
126impl<T> ExactSizeIterator for RowsIterMut<'_, T> {}
127impl<T> FusedIterator for RowsIterMut<'_, T> {}
128
129impl<'a, T: 'a> DoubleEndedIterator for RowsIterMut<'a, T> {
130    #[inline]
131    fn next_back(&mut self) -> Option<Self::Item> {
132        match self.inner.next_back() {
133            Some(s) => Some(&mut s[0..self.width]),
134            None => None,
135        }
136    }
137}
138
139/// Iterates over pixels in the (sub)image. Call `Img.pixels()` to create it.
140///
141/// Ignores padding, if there's any.
142#[must_use]
143pub struct PixelsIter<'a, T: Copy> {
144    inner: PixelsRefIter<'a, T>,
145}
146
147impl<'a, T: Copy + 'a> PixelsIter<'a, T> {
148    #[inline(always)]
149    #[track_caller]
150    pub(crate) fn new(img: super::ImgRef<'a, T>) -> Self {
151        Self {
152            inner: PixelsRefIter::new(img)
153        }
154    }
155}
156
157impl<'a, T: Copy + 'a> Iterator for PixelsIter<'a, T> {
158    type Item = T;
159
160    #[inline(always)]
161    fn next(&mut self) -> Option<Self::Item> {
162        self.inner.next().copied()
163    }
164}
165
166impl<T: Copy> ExactSizeIterator for PixelsIter<'_, T> {
167    #[inline]
168    fn len(&self) -> usize {
169        self.inner.len()
170    }
171}
172
173/// Iterates over pixels in the (sub)image. Call `Img.pixels_ref()` to create it.
174///
175/// Ignores padding, if there's any.
176#[derive(Debug)]
177#[must_use]
178pub struct PixelsRefIter<'a, T> {
179    current: *const T,
180    current_line_end: *const T,
181    rows_left: usize,
182    width: NonZeroUsize,
183    pad: usize,
184    _dat: PhantomData<&'a [T]>,
185}
186
187unsafe impl<T> Send for PixelsRefIter<'_, T> where T: Send {}
188unsafe impl<T> Sync for PixelsRefIter<'_, T> where T: Sync {}
189
190impl<'a, T: 'a> PixelsRefIter<'a, T> {
191    #[inline]
192    #[track_caller]
193    pub(crate) fn new(img: super::ImgRef<'a, T>) -> Self {
194        let width = NonZeroUsize::new(img.width()).expect("width > 0");
195        let height = img.height();
196        let stride = img.stride();
197        assert!(stride >= width.get());
198        let pad = stride - width.get();
199        debug_assert!(img.buf().len() + stride >= stride * height + width.get(),
200            "buffer len {} is less than {} (({}+{})x{})", img.buf().len(),
201            stride * height - pad, width, pad, height);
202        Self {
203            current: img.buf().as_ptr(),
204            current_line_end: img.buf()[width.get()..].as_ptr(),
205            width,
206            rows_left: height,
207            pad,
208            _dat: PhantomData,
209        }
210    }
211}
212
213impl<'a, T: 'a> Iterator for PixelsRefIter<'a, T> {
214    type Item = &'a T;
215
216    #[inline(always)]
217    fn next(&mut self) -> Option<Self::Item> {
218        unsafe {
219            if self.current >= self.current_line_end {
220                if self.rows_left <= 1 {
221                    return None;
222                }
223                self.rows_left -= 1;
224                self.current = self.current_line_end.add(self.pad);
225                self.current_line_end = self.current.add(self.width.get());
226            }
227            let px = &*self.current;
228            self.current = self.current.add(1);
229            Some(px)
230        }
231    }
232
233    #[inline]
234    #[cfg_attr(debug_assertions, track_caller)]
235    fn size_hint(&self) -> (usize, Option<usize>) {
236        let this_line = unsafe {
237            self.current_line_end.offset_from(self.current)
238        };
239        debug_assert!(this_line >= 0);
240        let len = this_line as usize + (self.rows_left - 1) * self.width.get();
241        (len, Some(len))
242    }
243}
244
245impl<T: Copy> ExactSizeIterator for PixelsRefIter<'_, T> {
246}
247
248/// Iterates over pixels in the (sub)image. Call `Img.pixels_mut()` to create it.
249///
250/// Ignores padding, if there's any.
251#[derive(Debug)]
252#[must_use]
253pub struct PixelsIterMut<'a, T> {
254    current: *mut T,
255    current_line_end: *mut T,
256    y: usize,
257    width: NonZeroUsize,
258    pad: usize,
259    _dat: PhantomData<&'a mut [T]>,
260}
261
262unsafe impl<T> Send for PixelsIterMut<'_, T> where T: Send {}
263unsafe impl<T> Sync for PixelsIterMut<'_, T> where T: Sync {}
264
265impl<'a, T: 'a> PixelsIterMut<'a, T> {
266    #[inline]
267    #[track_caller]
268    pub(crate) fn new(img: &mut super::ImgRefMut<'a, T>) -> Self {
269        let width = NonZeroUsize::new(img.width()).expect("width > 0");
270        let stride = img.stride();
271        debug_assert!(!img.buf().is_empty() && img.buf().len() + stride >= stride * img.height() + width.get());
272        Self {
273            current: img.buf_mut().as_mut_ptr(),
274            current_line_end: img.buf_mut()[width.get()..].as_mut_ptr(),
275            width,
276            y: img.height(),
277            pad: stride - width.get(),
278            _dat: PhantomData,
279        }
280    }
281}
282
283impl<'a, T: 'a> Iterator for PixelsIterMut<'a, T> {
284    type Item = &'a mut T;
285
286    #[inline(always)]
287    fn next(&mut self) -> Option<Self::Item> {
288        unsafe {
289            if self.current >= self.current_line_end {
290                self.y -= 1;
291                if self.y == 0 {
292                    return None;
293                }
294                self.current = self.current_line_end.add(self.pad);
295                self.current_line_end = self.current.add(self.width.get());
296            }
297            let px = &mut *self.current;
298            self.current = self.current.add(1);
299            Some(px)
300        }
301    }
302}
303
304#[test]
305fn iter() {
306    let img = super::Img::new(vec![1u8, 2], 1, 2);
307    let mut it = img.pixels();
308    assert_eq!(Some(1), it.next());
309    assert_eq!(Some(2), it.next());
310    assert_eq!(None, it.next());
311
312    let buf = [1u8; (16 + 3) * (8 + 1)];
313    for width in 1..16 {
314        for height in 1..8 {
315            for pad in 0..3 {
316                let stride = width + pad;
317                let img = super::Img::new_stride(&buf[..stride * height + stride - width], width, height, stride);
318                assert_eq!(width * height, img.pixels().map(|a| a as usize).sum(), "{width}x{height}");
319                assert_eq!(width * height, img.pixels().count(), "{width}x{height}");
320                assert_eq!(height, img.rows().count());
321
322                let mut iter1 = img.pixels();
323                let mut left = width * height;
324                while let Some(_px) = iter1.next() {
325                    left -= 1;
326                    assert_eq!(left, iter1.len());
327                }
328                assert_eq!(0, iter1.len());
329                iter1.next();
330                assert_eq!(0, iter1.len());
331
332                let mut iter2 = img.rows();
333                match iter2.next() {
334                    Some(_) => {
335                        assert_eq!(height - 1, iter2.size_hint().0);
336                        assert_eq!(height - 1, iter2.filter(|_| true).count());
337                    },
338                    None => {
339                        assert_eq!(height, 0);
340                    },
341                };
342            }
343        }
344    }
345}