iter: Add BytesExt trait for converting pixels into chunks

Signed-off-by: Christopher N. Hesse <[email protected]>

README.md

@@ -24,7 +24,7 @@ Below you can find a quick example usage of this crate. It introduces the basics
 
 ```rust
 use ffimage::color::{Gray, Rgb};
-use ffimage::iter::{ColorConvertExt, PixelsExt, WriteExt};
+use ffimage::iter::{BytesExt, ColorConvertExt, PixelsExt};
 
 fn main() {
     // This is our RGB image memory (2x2 pixels).
@@ -42,6 +42,7 @@ fn main() {
         .copied()
         .pixels::<Rgb<u8>>()
         .colorconvert::<Gray<u8>>()
+        .bytes()
         .write(&mut gray);
 }
 ```

ffimage-yuv/benches/convert.rs

@@ -2,7 +2,7 @@ use criterion::{black_box, criterion_group, Criterion};
 
 use ffimage::{
     color::Rgb,
-    iter::{ColorConvertExt, PixelsExt, WriteExt},
+    iter::{BytesExt, ColorConvertExt, PixelsExt},
 };
 use ffimage_yuv::{yuv::Yuv, yuv420::Yuv420p, yuv422::Yuv422};
 
@@ -19,6 +19,7 @@ pub fn yuv_to_rgb(c: &mut Criterion) {
                     .copied()
                     .pixels::<Yuv<u8>>()
                     .colorconvert::<Rgb<u8>>()
+                    .bytes()
                     .write(black_box(&mut rgb))
             })
         });
@@ -43,6 +44,7 @@ pub fn yuv422_to_rgb(c: &mut Criterion) {
                         .colorconvert::<[Yuv<u8>; 2]>()
                         .flatten()
                         .colorconvert::<Rgb<u8>>()
+                        .bytes()
                         .write(black_box(&mut rgb));
                 })
             },
@@ -64,6 +66,7 @@ pub fn yuv420p_to_rgb(c: &mut Criterion) {
                     Yuv420p::pack(&yuv420p)
                         .into_iter()
                         .colorconvert::<Rgb<u8>>()
+                        .bytes()
                         .write(black_box(&mut rgb));
                 })
             },

ffimage/benches/convert.rs

@@ -2,7 +2,7 @@ use criterion::{black_box, criterion_group, Criterion};
 
 use ffimage::{
     color::{Bgr, Gray, Rgb},
-    iter::{ColorConvertExt, PixelsExt, WriteExt},
+    iter::{BytesExt, ColorConvertExt, PixelsExt},
 };
 
 pub fn rgb_to_bgr(c: &mut Criterion) {
@@ -18,6 +18,7 @@ pub fn rgb_to_bgr(c: &mut Criterion) {
                     .copied()
                     .pixels::<Rgb<u8>>()
                     .colorconvert::<Bgr<u8>>()
+                    .bytes()
                     .write(black_box(&mut bgr))
             })
         });
@@ -37,6 +38,7 @@ pub fn rgb_to_gray(c: &mut Criterion) {
                     .copied()
                     .pixels::<Rgb<u8>>()
                     .colorconvert::<Gray<u8>>()
+                    .bytes()
                     .write(black_box(&mut gray))
             })
         });

ffimage/src/iter.rs

@@ -93,25 +93,63 @@ where
 /// The trait is automatically implemented for all pixel types which implement the
 /// `Deref<Target = [T; C]>` trait where T: Copy and C means the number of channels
 /// (e.g. 3 for RGB).
-pub trait WriteExt<'a, T, O, const C: usize>: Iterator {
-    fn write(self, out: O)
+pub trait BytesExt<T, const C: usize>: Iterator {
+    fn bytes(self) -> Bytes<T, Self, C>
     where
         Self: Sized,
-        Self::Item: Deref<Target = [T; C]>,
-        T: 'a + Copy,
+    {
+        Bytes::new(self)
+    }
+}
+
+impl<T, I, const C: usize> BytesExt<T, C> for I where I: Iterator {}
+
+pub struct Bytes<T, I, const C: usize> {
+    _marker: PhantomData<T>,
+    iter: I,
+}
+
+impl<T, I, const C: usize> Bytes<T, I, C> {
+    pub fn new(iter: I) -> Self {
+        Bytes {
+            _marker: PhantomData,
+            iter,
+        }
+    }
+}
+
+impl<T, I, const C: usize> Iterator for Bytes<T, I, C>
+where
+    T: Copy,
+    I: Iterator,
+    I::Item: Deref<Target = [T; C]>,
+{
+    type Item = [T; C];
+
+    fn next(&mut self) -> Option<Self::Item> {
+        Some(*self.iter.next()?)
+    }
+}
+
+impl<'a, T, I, const C: usize> Bytes<T, I, C>
+where
+    T: 'a + Copy,
+    I: Iterator,
+    I::Item: Deref<Target = [T; C]>,
+{
+    pub fn write<O>(self, out: O)
+    where
         O: IntoIterator<Item = &'a mut T>,
     {
         let mut out = out.into_iter();
         self.for_each(|chunk| {
-            for i in 0..C {
-                *(out.next().unwrap()) = chunk[i];
-            }
+            chunk.iter().for_each(|channel| {
+                *(out.next().expect("output iterator ended prematurely")) = *channel
+            })
         });
     }
 }
 
-impl<'a, T, I, O, const C: usize> WriteExt<'a, T, O, C> for I where I: Iterator {}
-
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -129,13 +167,14 @@ mod tests {
     }
 
     #[test]
-    fn write() {
+    fn bytes() {
         let buf = [1, 2, 3, 4, 5, 6, 7, 8, 9];
         let mut out = [0; 9];
         buf.iter()
             .copied()
             .pixels::<Rgb<u8>>()
             .colorconvert::<Bgr<u8>>()
+            .bytes()
             .write(&mut out);
         assert_eq!(out, [3, 2, 1, 6, 5, 4, 9, 8, 7]);
     }