Skip to main content

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 buf = img.valid_buf();
195        let height = img.height();
196        match NonZeroUsize::new(img.width()) {
197            Some(width) if height > 0 => {
198                let stride = img.stride();
199                let pad = stride - width.get();
200                Self {
201                    current: buf.as_ptr(),
202                    current_line_end: buf[width.get()..].as_ptr(),
203                    width,
204                    rows_left: height - 1,
205                    pad,
206                    _dat: PhantomData,
207                }
208            },
209            _ => {
210                Self {
211                    current: buf.as_ptr(),
212                    current_line_end: buf.as_ptr(),
213                    width: NonZeroUsize::new(1).unwrap(),
214                    rows_left: 0,
215                    pad: 0,
216                    _dat: PhantomData,
217                }
218            }
219        }
220    }
221}
222
223impl<'a, T: 'a> Iterator for PixelsRefIter<'a, T> {
224    type Item = &'a T;
225
226    #[inline(always)]
227    fn next(&mut self) -> Option<Self::Item> {
228        unsafe {
229            if self.current >= self.current_line_end {
230                if self.rows_left == 0 {
231                    return None;
232                }
233                self.rows_left -= 1;
234                self.current = self.current_line_end.add(self.pad);
235                self.current_line_end = self.current.add(self.width.get());
236            }
237            let px = &*self.current;
238            self.current = self.current.add(1);
239            Some(px)
240        }
241    }
242
243    #[inline]
244    #[cfg_attr(debug_assertions, track_caller)]
245    fn size_hint(&self) -> (usize, Option<usize>) {
246        let this_line = unsafe {
247            self.current_line_end.offset_from(self.current)
248        };
249        debug_assert!(this_line >= 0);
250        let len = this_line as usize + self.rows_left * self.width.get();
251        (len, Some(len))
252    }
253}
254
255impl<T: Copy> ExactSizeIterator for PixelsRefIter<'_, T> {
256}
257
258/// Iterates over pixels in the (sub)image. Call `Img.pixels_mut()` to create it.
259///
260/// Ignores padding, if there's any.
261#[derive(Debug)]
262#[must_use]
263pub struct PixelsIterMut<'a, T> {
264    current: *mut T,
265    current_line_end: *mut T,
266    rows_left: usize,
267    width: NonZeroUsize,
268    pad: usize,
269    _dat: PhantomData<&'a mut [T]>,
270}
271
272unsafe impl<T> Send for PixelsIterMut<'_, T> where T: Send {}
273unsafe impl<T> Sync for PixelsIterMut<'_, T> where T: Sync {}
274
275impl<'a, T: 'a> PixelsIterMut<'a, T> {
276    #[inline]
277    #[track_caller]
278    pub(crate) fn new(mut img: super::ImgRefMut<'a, T>) -> Self {
279        let width = img.width();
280        let height = img.height();
281        let stride = img.stride();
282        let buf = img.valid_buf_mut();
283        let ptr = buf.as_mut_ptr();
284        match NonZeroUsize::new(width) {
285            Some(width) if height > 0 => {
286                Self {
287                    current: ptr,
288                    current_line_end: unsafe { ptr.add(width.get()) },
289                    width,
290                    rows_left: height - 1,
291                    pad: stride - width.get(),
292                    _dat: PhantomData,
293                }
294            },
295            _ => {
296                Self {
297                    current: ptr,
298                    current_line_end: ptr,
299                    width: NonZeroUsize::new(1).unwrap(),
300                    rows_left: 0,
301                    pad: 0,
302                    _dat: PhantomData,
303                }
304            }
305        }
306    }
307}
308
309impl<'a, T: 'a> Iterator for PixelsIterMut<'a, T> {
310    type Item = &'a mut T;
311
312    #[inline(always)]
313    fn next(&mut self) -> Option<Self::Item> {
314        unsafe {
315            if self.current >= self.current_line_end {
316                if self.rows_left == 0 {
317                    return None;
318                }
319                self.rows_left -= 1;
320                self.current = self.current_line_end.add(self.pad);
321                self.current_line_end = self.current.add(self.width.get());
322            }
323            let px = &mut *self.current;
324            self.current = self.current.add(1);
325            Some(px)
326        }
327    }
328
329    #[inline]
330    #[cfg_attr(debug_assertions, track_caller)]
331    fn size_hint(&self) -> (usize, Option<usize>) {
332        let this_line = unsafe {
333            self.current_line_end.offset_from(self.current)
334        };
335        debug_assert!(this_line >= 0);
336        let len = this_line as usize + self.rows_left * self.width.get();
337        (len, Some(len))
338    }
339}
340
341impl<T: Copy> ExactSizeIterator for PixelsIterMut<'_, T> {
342}
343
344#[test]
345fn iter() {
346    let img = super::Img::new(vec![1u8, 2], 1, 2);
347    let mut it = img.pixels();
348    assert_eq!(Some(1), it.next());
349    assert_eq!(Some(2), it.next());
350    assert_eq!(None, it.next());
351
352    let buf = [1u8; (16 + 3) * (8 + 1)];
353    for width in 1..16 {
354        for height in 1..8 {
355            for pad in 0..3 {
356                let stride = width + pad;
357                let img = super::Img::new_stride(&buf[..stride * height + stride - width], width, height, stride);
358                assert_eq!(width * height, img.pixels().map(|a| a as usize).sum(), "{width}x{height}");
359                assert_eq!(width * height, img.pixels().count(), "{width}x{height}");
360                assert_eq!(height, img.rows().count());
361
362                let mut iter1 = img.pixels();
363                let mut left = width * height;
364                while let Some(_px) = iter1.next() {
365                    left -= 1;
366                    assert_eq!(left, iter1.len());
367                }
368                assert_eq!(0, iter1.len());
369                assert_eq!(0, left);
370                iter1.next();
371                assert_eq!(0, iter1.len());
372
373                let mut iter2 = img.rows();
374                match iter2.next() {
375                    Some(_) => {
376                        assert_eq!(height - 1, iter2.size_hint().0);
377                        assert_eq!(height - 1, iter2.filter(|_| true).count());
378                    },
379                    None => {
380                        assert_eq!(height, 0);
381                    },
382                }
383            }
384        }
385    }
386}