From a330187a515ffce4cb05606022c894d3f53ca968 Mon Sep 17 00:00:00 2001
From: Chris Hennick <hennickc@amazon.com>
Date: Wed, 21 Jun 2023 09:20:00 -0700
Subject: [PATCH] Use deflater on iCCP chunk as well

---
 src/deflate/mod.rs | 17 ++++++++++-------
 src/lib.rs         | 23 ++++++++---------------
 tests/raw.rs       |  6 ++++--
 3 files changed, 22 insertions(+), 24 deletions(-)

diff --git a/src/deflate/mod.rs b/src/deflate/mod.rs
index a678e0e2..ba606689 100644
--- a/src/deflate/mod.rs
+++ b/src/deflate/mod.rs
@@ -107,23 +107,26 @@ impl Deflater for BufferedZopfliDeflater {
             maximum_block_splits: self.max_block_splits,
             ..Default::default() // for forward compatibility
         };
+        let mut out = Vec::with_capacity(self.output_buffer_size);
         let mut buffer = BufWriter::with_capacity(
             self.input_buffer_size,
             DeflateEncoder::new(
                 options,
                 Default::default(),
-                Cursor::new(Vec::with_capacity(self.output_buffer_size)),
+                &mut out,
             ),
         );
-        let result = (|| -> io::Result<Vec<u8>> {
+        let result = (|| -> io::Result<()> {
             buffer.write_all(data)?;
-            Ok(buffer.into_inner()?.finish()?.into_inner())
+            buffer.into_inner()?.finish()?;
+            Ok(())
         })();
-        let result = result.map_err(|e| PngError::new(&e.to_string()))?;
-        if max_size.get().is_some_and(|max| max < result.len()) {
-            Err(PngError::DeflatedDataTooLong(result.len()))
+        result.map_err(|e| PngError::new(&e.to_string()))?;
+        println!("Compressed {} -> {} bytes", data.len(), out.len());
+        if max_size.get().is_some_and(|max| max < out.len()) {
+            Err(PngError::DeflatedDataTooLong(out.len()))
         } else {
-            Ok(result)
+            Ok(out)
         }
     }
 }
diff --git a/src/lib.rs b/src/lib.rs
index e2b32b77..28e5a4fb 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -403,7 +403,7 @@ impl RawImage {
             .filter(|c| opts.strip.keep(&c.name))
             .cloned()
             .collect();
-        postprocess_chunks(&mut png, opts, &self.png.ihdr);
+        postprocess_chunks(&mut png, opts, &self.png.ihdr, deflater);
 
         Ok(png.output())
     }
@@ -579,7 +579,7 @@ fn optimize_png(
         png.idat_data = new_png.idat_data;
     }
 
-    postprocess_chunks(png, opts, &raw.ihdr);
+    postprocess_chunks(png, opts, &raw.ihdr, &opts.deflate);
 
     let output = png.output();
 
@@ -675,17 +675,9 @@ fn optimize_raw<T: Deflater>(
             // We should have a result here - fail if not (e.g. deadline passed)
             let result = eval_result?;
 
-            match opts.deflate {
-                Deflaters::Libdeflater { compression } if compression <= eval_compression => {
-                    // No further compression required
-                    Some((result.filter, result.idat_data))
-                }
-                _ => {
-                    debug!("Trying: {}", result.filter);
-                    let best_size = AtomicMin::new(max_size);
-                    perform_trial(&result.filtered, opts, result.filter, &best_size, deflater)
-                }
-            }
+            debug!("Trying: {}", result.filter);
+            let best_size = AtomicMin::new(max_size);
+            perform_trial(&result.filtered, opts, result.filter, &best_size, deflater)
         } else {
             // Perform full compression trials of selected filters and determine the best
 
@@ -855,7 +847,8 @@ fn report_format(prefix: &str, png: &PngImage) {
 }
 
 /// Perform cleanup of certain chunks from the `PngData` object, after optimization has been completed
-fn postprocess_chunks(png: &mut PngData, opts: &Options, orig_ihdr: &IhdrData) {
+fn postprocess_chunks<T>(png: &mut PngData, opts: &Options, orig_ihdr: &IhdrData, deflater: &T)
+        where T: Deflater {
     if let Some(iccp_idx) = png.aux_chunks.iter().position(|c| &c.name == b"iCCP") {
         // See if we can replace an iCCP chunk with an sRGB chunk
         let may_replace_iccp = opts.strip != StripChunks::None && opts.strip.keep(b"sRGB");
@@ -877,7 +870,7 @@ fn postprocess_chunks(png: &mut PngData, opts: &Options, orig_ihdr: &IhdrData) {
                     name: *b"sRGB",
                     data: vec![intent],
                 };
-            } else if let Ok(iccp) = construct_iccp(&icc, &opts.deflate) {
+            } else if let Ok(iccp) = construct_iccp(&icc, deflater) {
                 let cur_len = png.aux_chunks[iccp_idx].data.len();
                 let new_len = iccp.data.len();
                 if new_len < cur_len {
diff --git a/tests/raw.rs b/tests/raw.rs
index c3e16fb2..8aebc1f8 100644
--- a/tests/raw.rs
+++ b/tests/raw.rs
@@ -14,6 +14,7 @@ fn get_opts() -> Options {
 fn test_it_converts(input: &str) {
     let input = PathBuf::from(input);
     let opts = get_opts();
+    let deflater = BufferedZopfliDeflater::default();
 
     let original_data = PngData::read_file(&PathBuf::from(input)).unwrap();
     let image = PngData::from_slice(&original_data, &opts).unwrap();
@@ -35,7 +36,7 @@ fn test_it_converts(input: &str) {
         raw.add_png_chunk(chunk.name, chunk.data);
     }
 
-    let output = raw.create_optimized_png(&opts).unwrap();
+    let output = raw.create_optimized_png(&opts, &deflater).unwrap();
 
     let new = PngData::from_slice(&output, &opts).unwrap();
     assert!(new.aux_chunks.len() == num_chunks);
@@ -52,6 +53,7 @@ fn from_file() {
 #[test]
 fn custom_indexed() {
     let opts = get_opts();
+    let deflater = BufferedZopfliDeflater::default();
 
     let raw = RawImage::new(
         4,
@@ -69,7 +71,7 @@ fn custom_indexed() {
     )
     .unwrap();
 
-    raw.create_optimized_png(&opts).unwrap();
+    raw.create_optimized_png(&opts, &deflater).unwrap();
 }
 
 #[test]