Skip to content

Commit a2678f7

Browse files
shssoichiroshssoichiro
committed
Implement ExactSizeIterator and DoubleEndedIterator
1 parent f794f5f commit a2678f7

File tree

2 files changed

+331
-21
lines changed

2 files changed

+331
-21
lines changed

‎src/plane.rs‎

Lines changed: 80 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -135,7 +135,8 @@ where
135135
/// Returns an iterator over the visible pixels of each row
136136
/// in the plane, from top to bottom.
137137
#[inline]
138-
pub fn rows(&self) -> impl Iterator<Item = &[T]> {
138+
#[must_use]
139+
pub fn rows(&self) -> impl DoubleEndedIterator<Item = &[T]> + ExactSizeIterator {
139140
let origin = self.data_origin();
140141
// SAFETY: The plane creation interface ensures the data is large enough
141142
let visible_data = unsafe { self.data.get_unchecked(origin..) };
@@ -151,7 +152,7 @@ where
151152
/// Returns a mutable iterator over the visible pixels of each row
152153
/// in the plane, from top to bottom.
153154
#[inline]
154-
pub fn rows_mut(&mut self) -> impl Iterator<Item = &mut [T]> {
155+
pub fn rows_mut(&mut self) -> impl DoubleEndedIterator<Item = &mut [T]> + ExactSizeIterator {
155156
let origin = self.data_origin();
156157
// SAFETY: The plane creation interface ensures the data is large enough
157158
let visible_data = unsafe { self.data.get_unchecked_mut(origin..) };
@@ -190,42 +191,56 @@ where
190191
/// Returns an iterator over the visible pixels in the plane,
191192
/// in row-major order.
192193
#[inline]
193-
pub fn pixels(&self) -> impl Iterator<Item = T> {
194-
self.rows().flatten().copied()
194+
#[must_use]
195+
pub fn pixels(&self) -> impl DoubleEndedIterator<Item = T> + ExactSizeIterator {
196+
let total = self.width().get() * self.height().get();
197+
ExactSizeWrapper {
198+
iter: self.rows().flatten().copied(),
199+
len: total,
200+
}
195201
}
196202

197203
/// Returns a mutable iterator over the visible pixels in the plane,
198204
/// in row-major order.
199205
#[inline]
200-
pub fn pixels_mut(&mut self) -> impl Iterator<Item = &mut T> {
201-
self.rows_mut().flatten()
206+
pub fn pixels_mut(&mut self) -> impl DoubleEndedIterator<Item = &mut T> + ExactSizeIterator {
207+
let total = self.width().get() * self.height().get();
208+
ExactSizeWrapper {
209+
iter: self.rows_mut().flatten(),
210+
len: total,
211+
}
202212
}
203213

204214
/// Returns an iterator over the visible byte data in the plane,
205215
/// in row-major order. High-bit-depth data is converted to `u8`
206216
/// using low endianness.
207217
#[inline]
208-
pub fn byte_data(&self) -> impl Iterator<Item = u8> {
218+
#[must_use]
219+
pub fn byte_data(&self) -> impl DoubleEndedIterator<Item = u8> + ExactSizeIterator {
209220
let byte_width = size_of::<T>();
210221
assert!(
211222
byte_width <= 2,
212223
"unsupported pixel byte width: {byte_width}"
213224
);
214225

215-
self.pixels().flat_map(move |pix| {
216-
let bytes: [u8; 2] = if byte_width == 1 {
217-
[
218-
pix.to_u8()
219-
.expect("Pixel::byte_data only supports u8 and u16 pixels"),
220-
0,
221-
]
222-
} else {
223-
pix.to_u16()
224-
.expect("Pixel::byte_data only supports u8 and u16 pixels")
225-
.to_le_bytes()
226-
};
227-
bytes.into_iter().take(byte_width)
228-
})
226+
let total = self.width().get() * self.height().get() * byte_width;
227+
ExactSizeWrapper {
228+
iter: self.pixels().flat_map(move |pix| {
229+
let bytes: [u8; 2] = if byte_width == 1 {
230+
[
231+
pix.to_u8()
232+
.expect("Pixel::byte_data only supports u8 and u16 pixels"),
233+
0,
234+
]
235+
} else {
236+
pix.to_u16()
237+
.expect("Pixel::byte_data only supports u8 and u16 pixels")
238+
.to_le_bytes()
239+
};
240+
bytes.into_iter().take(byte_width)
241+
}),
242+
len: total,
243+
}
229244
}
230245

231246
/// Copies the data from `src` into this plane's visible pixels.
@@ -417,3 +432,47 @@ impl PlaneGeometry {
417432
.saturating_add(self.pad_bottom)
418433
}
419434
}
435+
436+
/// Wrapper to add `ExactSizeIterator` implementation to iterators with known length.
437+
struct ExactSizeWrapper<I> {
438+
iter: I,
439+
len: usize,
440+
}
441+
442+
impl<I: Iterator> Iterator for ExactSizeWrapper<I> {
443+
type Item = I::Item;
444+
445+
#[inline]
446+
fn next(&mut self) -> Option<Self::Item> {
447+
if let Some(item) = self.iter.next() {
448+
self.len = self.len.saturating_sub(1);
449+
Some(item)
450+
} else {
451+
None
452+
}
453+
}
454+
455+
#[inline]
456+
fn size_hint(&self) -> (usize, Option<usize>) {
457+
(self.len, Some(self.len))
458+
}
459+
}
460+
461+
impl<I: DoubleEndedIterator> DoubleEndedIterator for ExactSizeWrapper<I> {
462+
#[inline]
463+
fn next_back(&mut self) -> Option<Self::Item> {
464+
if let Some(item) = self.iter.next_back() {
465+
self.len = self.len.saturating_sub(1);
466+
Some(item)
467+
} else {
468+
None
469+
}
470+
}
471+
}
472+
473+
impl<I: Iterator> ExactSizeIterator for ExactSizeWrapper<I> {
474+
#[inline]
475+
fn len(&self) -> usize {
476+
self.len
477+
}
478+
}

‎src/plane/tests.rs‎

Lines changed: 251 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -549,3 +549,254 @@ fn pixels_count() {
549549
let pixel_count = plane.pixels().count();
550550
assert_eq!(pixel_count, 7 * 11);
551551
}
552+
553+
#[test]
554+
fn exact_size_iterator_rows() {
555+
let geometry = simple_geometry(4, 3);
556+
let plane: Plane<u8> = Plane::new(geometry);
557+
558+
let mut rows_iter = plane.rows();
559+
assert_eq!(rows_iter.len(), 3);
560+
561+
rows_iter.next();
562+
assert_eq!(rows_iter.len(), 2);
563+
564+
rows_iter.next();
565+
assert_eq!(rows_iter.len(), 1);
566+
567+
rows_iter.next();
568+
assert_eq!(rows_iter.len(), 0);
569+
570+
assert!(rows_iter.next().is_none());
571+
}
572+
573+
#[test]
574+
fn exact_size_iterator_pixels() {
575+
let geometry = simple_geometry(4, 3);
576+
let plane: Plane<u8> = Plane::new(geometry);
577+
578+
let mut pixels_iter = plane.pixels();
579+
assert_eq!(pixels_iter.len(), 12); // 4 * 3
580+
581+
for i in (0..12).rev() {
582+
assert_eq!(pixels_iter.len(), i + 1);
583+
pixels_iter.next();
584+
}
585+
586+
assert_eq!(pixels_iter.len(), 0);
587+
assert!(pixels_iter.next().is_none());
588+
}
589+
590+
#[test]
591+
fn exact_size_iterator_pixels_mut() {
592+
let geometry = simple_geometry(3, 2);
593+
let mut plane: Plane<u8> = Plane::new(geometry);
594+
595+
let mut pixels_iter = plane.pixels_mut();
596+
assert_eq!(pixels_iter.len(), 6); // 3 * 2
597+
598+
pixels_iter.next();
599+
assert_eq!(pixels_iter.len(), 5);
600+
}
601+
602+
#[test]
603+
fn exact_size_iterator_byte_data_u8() {
604+
let geometry = simple_geometry(4, 3);
605+
let plane: Plane<u8> = Plane::new(geometry);
606+
607+
let mut bytes_iter = plane.byte_data();
608+
assert_eq!(bytes_iter.len(), 12); // 4 * 3 * 1 byte
609+
610+
bytes_iter.next();
611+
assert_eq!(bytes_iter.len(), 11);
612+
}
613+
614+
#[test]
615+
fn exact_size_iterator_byte_data_u16() {
616+
let geometry = simple_geometry(4, 3);
617+
let plane: Plane<u16> = Plane::new(geometry);
618+
619+
let mut bytes_iter = plane.byte_data();
620+
assert_eq!(bytes_iter.len(), 24); // 4 * 3 * 2 bytes
621+
622+
bytes_iter.next();
623+
assert_eq!(bytes_iter.len(), 23);
624+
}
625+
626+
#[test]
627+
fn double_ended_iterator_rows() {
628+
let geometry = simple_geometry(4, 3);
629+
let mut plane: Plane<u8> = Plane::new(geometry);
630+
631+
// Fill each row with its index
632+
for (y, row) in plane.rows_mut().enumerate() {
633+
for pixel in row {
634+
*pixel = y as u8;
635+
}
636+
}
637+
638+
let mut rows_iter = plane.rows();
639+
assert_eq!(rows_iter.len(), 3);
640+
641+
// Get first row
642+
let first_row = rows_iter.next().unwrap();
643+
assert!(first_row.iter().all(|&p| p == 0));
644+
assert_eq!(rows_iter.len(), 2);
645+
646+
// Get last row
647+
let last_row = rows_iter.next_back().unwrap();
648+
assert!(last_row.iter().all(|&p| p == 2));
649+
assert_eq!(rows_iter.len(), 1);
650+
651+
// Get middle row
652+
let middle_row = rows_iter.next().unwrap();
653+
assert!(middle_row.iter().all(|&p| p == 1));
654+
assert_eq!(rows_iter.len(), 0);
655+
}
656+
657+
#[test]
658+
fn double_ended_iterator_rows_mut() {
659+
let geometry = simple_geometry(3, 3);
660+
let mut plane: Plane<u8> = Plane::new(geometry);
661+
662+
let mut rows_iter = plane.rows_mut();
663+
664+
// Modify first row
665+
let first_row = rows_iter.next().unwrap();
666+
for pixel in first_row {
667+
*pixel = 1;
668+
}
669+
670+
// Modify last row
671+
let last_row = rows_iter.next_back().unwrap();
672+
for pixel in last_row {
673+
*pixel = 3;
674+
}
675+
676+
// Middle row should still be 0
677+
let middle_row = rows_iter.next().unwrap();
678+
for pixel in middle_row {
679+
*pixel = 2;
680+
}
681+
682+
drop(rows_iter);
683+
684+
// Verify
685+
assert_eq!(plane.row(0).unwrap(), &[1, 1, 1]);
686+
assert_eq!(plane.row(1).unwrap(), &[2, 2, 2]);
687+
assert_eq!(plane.row(2).unwrap(), &[3, 3, 3]);
688+
}
689+
690+
#[test]
691+
fn double_ended_iterator_pixels() {
692+
let geometry = simple_geometry(3, 2);
693+
let mut plane: Plane<u8> = Plane::new(geometry);
694+
695+
// Fill with sequential values: [0, 1, 2, 3, 4, 5]
696+
for (i, pixel) in plane.pixels_mut().enumerate() {
697+
*pixel = i as u8;
698+
}
699+
700+
let mut pixels_iter = plane.pixels();
701+
assert_eq!(pixels_iter.len(), 6);
702+
703+
// Get first pixel
704+
assert_eq!(pixels_iter.next(), Some(0));
705+
assert_eq!(pixels_iter.len(), 5);
706+
707+
// Get last pixel
708+
assert_eq!(pixels_iter.next_back(), Some(5));
709+
assert_eq!(pixels_iter.len(), 4);
710+
711+
// Get second pixel
712+
assert_eq!(pixels_iter.next(), Some(1));
713+
assert_eq!(pixels_iter.len(), 3);
714+
715+
// Get second-to-last pixel
716+
assert_eq!(pixels_iter.next_back(), Some(4));
717+
assert_eq!(pixels_iter.len(), 2);
718+
719+
// Remaining pixels
720+
assert_eq!(pixels_iter.next(), Some(2));
721+
assert_eq!(pixels_iter.next(), Some(3));
722+
assert_eq!(pixels_iter.len(), 0);
723+
assert!(pixels_iter.next().is_none());
724+
assert!(pixels_iter.next_back().is_none());
725+
}
726+
727+
#[test]
728+
fn double_ended_iterator_pixels_mut() {
729+
let geometry = simple_geometry(2, 2);
730+
let mut plane: Plane<u8> = Plane::new(geometry);
731+
732+
let mut pixels_iter = plane.pixels_mut();
733+
734+
// Set first pixel
735+
*pixels_iter.next().unwrap() = 10;
736+
737+
// Set last pixel
738+
*pixels_iter.next_back().unwrap() = 40;
739+
740+
// Set remaining pixels
741+
*pixels_iter.next().unwrap() = 20;
742+
*pixels_iter.next().unwrap() = 30;
743+
744+
drop(pixels_iter);
745+
746+
// Verify: pixels should be [10, 20, 30, 40]
747+
let result: Vec<u8> = plane.pixels().collect();
748+
assert_eq!(result, vec![10, 20, 30, 40]);
749+
}
750+
751+
#[test]
752+
fn double_ended_iterator_byte_data() {
753+
let geometry = simple_geometry(2, 2);
754+
let mut plane: Plane<u8> = Plane::new(geometry);
755+
756+
// Fill with test data: [1, 2, 3, 4]
757+
for (i, pixel) in plane.pixels_mut().enumerate() {
758+
*pixel = (i + 1) as u8;
759+
}
760+
761+
let mut bytes_iter = plane.byte_data();
762+
assert_eq!(bytes_iter.len(), 4);
763+
764+
// Get first byte
765+
assert_eq!(bytes_iter.next(), Some(1));
766+
assert_eq!(bytes_iter.len(), 3);
767+
768+
// Get last byte
769+
assert_eq!(bytes_iter.next_back(), Some(4));
770+
assert_eq!(bytes_iter.len(), 2);
771+
772+
// Remaining bytes
773+
assert_eq!(bytes_iter.next(), Some(2));
774+
assert_eq!(bytes_iter.next_back(), Some(3));
775+
assert_eq!(bytes_iter.len(), 0);
776+
}
777+
778+
#[test]
779+
fn double_ended_iterator_byte_data_u16() {
780+
let geometry = simple_geometry(2, 1);
781+
let mut plane: Plane<u16> = Plane::new(geometry);
782+
783+
// Set two u16 values: 0x0102 and 0x0304
784+
*plane.pixel_mut(0, 0).unwrap() = 0x0102;
785+
*plane.pixel_mut(1, 0).unwrap() = 0x0304;
786+
787+
let mut bytes_iter = plane.byte_data();
788+
assert_eq!(bytes_iter.len(), 4); // 2 pixels * 2 bytes
789+
790+
// Get first byte (little endian)
791+
assert_eq!(bytes_iter.next(), Some(0x02));
792+
assert_eq!(bytes_iter.len(), 3);
793+
794+
// Get last byte
795+
assert_eq!(bytes_iter.next_back(), Some(0x03));
796+
assert_eq!(bytes_iter.len(), 2);
797+
798+
// Remaining bytes
799+
assert_eq!(bytes_iter.next(), Some(0x01));
800+
assert_eq!(bytes_iter.next_back(), Some(0x04));
801+
assert_eq!(bytes_iter.len(), 0);
802+
}

0 commit comments

Comments
 (0)