Test encoder by dogfooding decoded image

Author: friendlymatthew

src/png/chunk.rs

@@ -1,4 +1,5 @@
 use crate::png::grammar::ImageHeader;
+use crate::png::scanline_writer::ScanlineWriter;
 use anyhow::Result;
 use flate2::write::ZlibEncoder;
 use flate2::Compression;
@@ -35,11 +36,11 @@ pub trait PngChunk {
 }
 
 #[derive(Debug)]
-pub struct IHDRChunk {
-    pub image_header: ImageHeader,
+pub struct IHDRChunk<'a> {
+    pub image_header: &'a ImageHeader,
 }
 
-impl PngChunk for IHDRChunk {
+impl PngChunk for IHDRChunk<'_> {
     const NAME: [u8; 4] = *b"IHDR";
 
     fn data(&self) -> Result<Vec<u8>> {
@@ -58,10 +59,10 @@ impl PngChunk for IHDRChunk {
         buffer.extend_from_slice(&width.to_be_bytes());
         buffer.extend_from_slice(&height.to_be_bytes());
         buffer.extend_from_slice(&bit_depth.to_be_bytes());
-        buffer.extend_from_slice(&(color_type as u8).to_be_bytes());
+        buffer.extend_from_slice(&(*color_type as u8).to_be_bytes());
         buffer.extend_from_slice(&compression_method.to_be_bytes());
         buffer.extend_from_slice(&filter_method.to_be_bytes());
-        buffer.extend_from_slice(&(interlace_method as u8).to_be_bytes());
+        buffer.extend_from_slice(&(*interlace_method as u8).to_be_bytes());
 
         Ok(buffer)
     }
@@ -75,18 +76,23 @@ impl PngChunk for PLTEChunk {
 }
 
 #[derive(Debug)]
-pub struct IDATChunk {
-    pub data: Vec<u8>,
+pub struct IDATChunk<'a> {
+    pub image_header: &'a ImageHeader,
+    pub data: &'a [u8],
 }
 
-impl PngChunk for IDATChunk {
+impl PngChunk for IDATChunk<'_> {
     const NAME: [u8; 4] = *b"IDAT";
 
     fn data(&self) -> Result<Vec<u8>> {
+        let scanline_writer = ScanlineWriter::new(self.image_header);
+        let mut scanned_pixels = Vec::new();
+        scanline_writer.write(&mut scanned_pixels, self.data)?;
+
         let compressed_data = Vec::new();
 
         let mut encoder = ZlibEncoder::new(compressed_data, Compression::fast());
-        encoder.write_all(&self.data)?;
+        encoder.write_all(&scanned_pixels)?;
 
         Ok(encoder.finish()?)
     }

src/png/encoder.rs

@@ -12,7 +12,7 @@ impl<W: Write> PngEncoder<W> {
         Self { writer }
     }
 
-    pub fn encode(&mut self, png: Png) -> Result<()> {
+    pub fn encode(&mut self, png: &Png) -> Result<()> {
         self.writer.write_all(b"\x89PNG\r\n\x1A\n")?;
 
         let Png {
@@ -27,7 +27,10 @@ impl<W: Write> PngEncoder<W> {
         // let palette_chunk = PLTEChunk;
         // palette_chunk.write(&mut self.writer)?;
 
-        let image_data_chunk = IDATChunk { data: pixel_buffer };
+        let image_data_chunk = IDATChunk {
+            image_header,
+            data: pixel_buffer,
+        };
         image_data_chunk.write(&mut self.writer)?;
 
         let image_end = IENDChunk;
@@ -51,7 +54,24 @@ mod tests {
         let file = File::create("./tests/obama_encoded.png")?;
         let mut encoder = PngEncoder::new(file);
 
-        encoder.encode(png)?;
+        encoder.encode(&png)?;
+
+        Ok(())
+    }
+
+    #[test]
+    fn test_encode_round_trip() -> Result<()> {
+        let data = std::fs::read("./tests/obama.png")?;
+        let png = PngDecoder::new(&data).decode()?;
+
+        let file = File::create("./tests/obama_encoded.png")?;
+        let mut encoder = PngEncoder::new(file);
+        encoder.encode(&png)?;
+
+        let data = std::fs::read("./tests/obama_encoded.png")?;
+        let from_encoded_png = PngDecoder::new(&data).decode()?;
+
+        assert_eq!(png, from_encoded_png);
 
         Ok(())
     }

src/png/mod.rs

@@ -9,3 +9,4 @@ mod decoder;
 mod encoder;
 mod interlace;
 mod scanline_reader;
+mod scanline_writer;

src/png/scanline_writer.rs

@@ -0,0 +1,32 @@
+use crate::png::grammar::{Filter, ImageHeader};
+use anyhow::Result;
+use std::io::Write;
+
+#[derive(Debug)]
+pub struct ScanlineWriter<'a> {
+    image_header: &'a ImageHeader,
+}
+
+impl<'a> ScanlineWriter<'a> {
+    pub const fn new(image_header: &'a ImageHeader) -> Self {
+        Self { image_header }
+    }
+
+    pub fn write<W: Write>(&self, writer: &mut W, pixel_buffer: &'a [u8]) -> Result<()> {
+        assert_eq!(
+            self.image_header.width as usize
+                * self.image_header.height as usize
+                * self.image_header.num_bytes_per_pixel(),
+            pixel_buffer.len()
+        );
+
+        for chunk in pixel_buffer.chunks_exact(
+            self.image_header.width as usize * self.image_header.num_bytes_per_pixel(),
+        ) {
+            writer.write_all(&[Filter::None as u8])?;
+            writer.write_all(chunk)?;
+        }
+
+        Ok(())
+    }
+}