From f31776fed4a6e7466599eb70e00ac60fe953f874 Mon Sep 17 00:00:00 2001
From: "Christopher N. Hesse" <raymanfx@gmail.com>
Date: Sun, 5 Feb 2023 20:59:04 +0100
Subject: [PATCH] iter: Add write() method

write() allows callers to directly write pixels to an output stream
(aka Iterator). The parameter P2 allows for direct conversion from
one pixel format to another.

Signed-off-by: Christopher N. Hesse <raymanfx@gmail.com>
---
 ffimage/src/iter.rs | 41 ++++++++++++++++++++++++++++++++++++-----
 1 file changed, 36 insertions(+), 5 deletions(-)

diff --git a/ffimage/src/iter.rs b/ffimage/src/iter.rs
index 915a605..5f78abe 100644
--- a/ffimage/src/iter.rs
+++ b/ffimage/src/iter.rs
@@ -1,11 +1,11 @@
-use core::marker::PhantomData;
+use core::{marker::PhantomData, ops::Deref};
 
 /// Adapter which converts a bytestream into a typed pixel stream.
 ///
 /// The trait is automatically implemented for all pixel types which implement the `From<[T; C]>`
 /// trait where T: Copy and C means the number of channels (e.g. 3 for RGB).
-pub trait PixelsExt: Iterator {
-    fn pixels<P, const C: usize>(self) -> Pixels<Self::Item, Self, P, C>
+pub trait PixelsExt<const C: usize>: Iterator {
+    fn pixels<P>(self) -> Pixels<Self::Item, Self, P, C>
     where
         Self: Sized,
     {
@@ -13,7 +13,7 @@ pub trait PixelsExt: Iterator {
     }
 }
 
-impl<I> PixelsExt for I where I: Iterator {}
+impl<I, const C: usize> PixelsExt<C> for I where I: Iterator {}
 
 pub struct Pixels<T, I, P, const C: usize> {
     _marker: PhantomData<(T, P)>,
@@ -29,6 +29,26 @@ impl<T, I, P, const C: usize> Pixels<T, I, P, C> {
     }
 }
 
+impl<'a, T, I, P, const C: usize> Pixels<T, I, P, C>
+where
+    T: Copy + 'a,
+    I: Iterator<Item = T>,
+    P: From<[T; C]>,
+{
+    pub fn write<P2>(self, mut out: impl Iterator<Item = &'a mut T>)
+    where
+        P2: From<P> + Deref,
+        <P2 as Deref>::Target: AsRef<[T]>,
+    {
+        self.map(|p| P2::from(p)).for_each(|p2| {
+            p2.as_ref().iter().for_each(|t| {
+                let _out = out.next().unwrap();
+                *_out = *t;
+            });
+        });
+    }
+}
+
 impl<T, I, P, const C: usize> Iterator for Pixels<T, I, P, C>
 where
     T: Copy,
@@ -49,7 +69,7 @@ where
 #[cfg(test)]
 mod tests {
     use super::*;
-    use crate::color::Rgb;
+    use crate::color::{Bgr, Rgb};
 
     #[test]
     fn pixels() {
@@ -61,4 +81,15 @@ mod tests {
         assert_eq!(pixels.next(), Some(Rgb::<u8>([7, 8, 9])));
         assert_eq!(pixels.next(), None);
     }
+
+    #[test]
+    fn write() {
+        let buf = [1, 2, 3, 4, 5, 6, 7, 8, 9];
+        let mut out = [0; 9];
+        buf.iter()
+            .copied()
+            .pixels::<Rgb<u8>>()
+            .write::<Bgr<u8>>(out.iter_mut());
+        assert_eq!(out, [3, 2, 1, 6, 5, 4, 9, 8, 7]);
+    }
 }