Provide BufferedZopfliDeflater and allow user to pass in a custom Deflater

.whitesource

@@ -0,0 +1,14 @@
+{
+  "scanSettings": {
+    "baseBranches": []
+  },
+  "checkRunSettings": {
+    "vulnerableCheckRunConclusionLevel": "failure",
+    "displayMode": "diff",
+    "useMendCheckNames": true
+  },
+  "issueSettings": {
+    "minSeverityLevel": "LOW",
+    "issueType": "DEPENDENCY"
+  }
+}

Cargo.toml

@@ -28,6 +28,7 @@ required-features = ["zopfli"]
 
 [dependencies]
 zopfli = { version = "0.7.4", optional = true, default-features = false, features = ["std", "zlib"] }
+simd-adler32 = { version = "0.3.5", optional = true, default-features = false }
 rgb = "0.8.36"
 indexmap = "2.0.0"
 libdeflater = "0.14.0"
@@ -66,6 +67,7 @@ version = "0.24.6"
 rustc_version = "0.4.0"
 
 [features]
+zopfli = ["zopfli/std", "zopfli/zlib", "simd-adler32"]
 binary = ["clap", "wild", "stderrlog"]
 default = ["binary", "filetime", "parallel", "zopfli"]
 parallel = ["rayon", "indexmap/rayon", "crossbeam-channel"]

benches/zopfli.rs

@@ -5,30 +5,32 @@ extern crate test;
 
 use oxipng::internal_tests::*;
 use oxipng::*;
-use std::num::NonZeroU8;
 use std::path::PathBuf;
 use test::Bencher;
 
-// SAFETY: trivially safe. Stopgap solution until const unwrap is stabilized.
-const DEFAULT_ZOPFLI_ITERATIONS: NonZeroU8 = unsafe { NonZeroU8::new_unchecked(15) };
-
 #[bench]
 fn zopfli_16_bits_strategy_0(b: &mut Bencher) {
     let input = test::black_box(PathBuf::from("tests/files/rgb_16_should_be_rgb_16.png"));
     let png = PngData::new(&input, &Options::default()).unwrap();
+    let max_size = AtomicMin::new(Some(png.idat_data.len()));
 
     b.iter(|| {
-        zopfli_deflate(png.raw.data.as_ref(), DEFAULT_ZOPFLI_ITERATIONS).ok();
+        BufferedZopfliDeflater::default()
+            .deflate(png.raw.data.as_ref(), &max_size)
+            .ok();
     });
 }
 
 #[bench]
 fn zopfli_8_bits_strategy_0(b: &mut Bencher) {
     let input = test::black_box(PathBuf::from("tests/files/rgb_8_should_be_rgb_8.png"));
     let png = PngData::new(&input, &Options::default()).unwrap();
+    let max_size = AtomicMin::new(Some(png.idat_data.len()));
 
     b.iter(|| {
-        zopfli_deflate(png.raw.data.as_ref(), DEFAULT_ZOPFLI_ITERATIONS).ok();
+        BufferedZopfliDeflater::default()
+            .deflate(png.raw.data.as_ref(), &max_size)
+            .ok();
     });
 }
 
@@ -38,9 +40,12 @@ fn zopfli_4_bits_strategy_0(b: &mut Bencher) {
         "tests/files/palette_4_should_be_palette_4.png",
     ));
     let png = PngData::new(&input, &Options::default()).unwrap();
+    let max_size = AtomicMin::new(Some(png.idat_data.len()));
 
     b.iter(|| {
-        zopfli_deflate(png.raw.data.as_ref(), DEFAULT_ZOPFLI_ITERATIONS).ok();
+        BufferedZopfliDeflater::default()
+            .deflate(png.raw.data.as_ref(), &max_size)
+            .ok();
     });
 }
 
@@ -50,9 +55,12 @@ fn zopfli_2_bits_strategy_0(b: &mut Bencher) {
         "tests/files/palette_2_should_be_palette_2.png",
     ));
     let png = PngData::new(&input, &Options::default()).unwrap();
+    let max_size = AtomicMin::new(Some(png.idat_data.len()));
 
     b.iter(|| {
-        zopfli_deflate(png.raw.data.as_ref(), DEFAULT_ZOPFLI_ITERATIONS).ok();
+        BufferedZopfliDeflater::default()
+            .deflate(png.raw.data.as_ref(), &max_size)
+            .ok();
     });
 }
 
@@ -62,8 +70,11 @@ fn zopfli_1_bits_strategy_0(b: &mut Bencher) {
         "tests/files/palette_1_should_be_palette_1.png",
     ));
     let png = PngData::new(&input, &Options::default()).unwrap();
+    let max_size = AtomicMin::new(Some(png.idat_data.len()));
 
     b.iter(|| {
-        zopfli_deflate(png.raw.data.as_ref(), DEFAULT_ZOPFLI_ITERATIONS).ok();
+        BufferedZopfliDeflater::default()
+            .deflate(png.raw.data.as_ref(), &max_size)
+            .ok();
     });
 }

src/atomicmin.rs

@@ -23,7 +23,7 @@ impl AtomicMin {
     }
 
     /// Unset value is usize_max
-    pub fn as_atomic_usize(&self) -> &AtomicUsize {
+    pub const fn as_atomic_usize(&self) -> &AtomicUsize {
         &self.val
     }
 

src/colors.rs

@@ -45,7 +45,7 @@ impl Display for ColorType {
 impl ColorType {
     /// Get the code used by the PNG specification to denote this color type
     #[inline]
-    pub fn png_header_code(&self) -> u8 {
+    pub const fn png_header_code(&self) -> u8 {
         match self {
             ColorType::Grayscale { .. } => 0,
             ColorType::RGB { .. } => 2,
@@ -56,7 +56,7 @@ impl ColorType {
     }
 
     #[inline]
-    pub(crate) fn channels_per_pixel(&self) -> u8 {
+    pub(crate) const fn channels_per_pixel(&self) -> u8 {
         match self {
             ColorType::Grayscale { .. } | ColorType::Indexed { .. } => 1,
             ColorType::GrayscaleAlpha => 2,
@@ -66,25 +66,25 @@ impl ColorType {
     }
 
     #[inline]
-    pub(crate) fn is_rgb(&self) -> bool {
+    pub(crate) const fn is_rgb(&self) -> bool {
         matches!(self, ColorType::RGB { .. } | ColorType::RGBA)
     }
 
     #[inline]
-    pub(crate) fn is_gray(&self) -> bool {
+    pub(crate) const fn is_gray(&self) -> bool {
         matches!(
             self,
             ColorType::Grayscale { .. } | ColorType::GrayscaleAlpha
         )
     }
 
     #[inline]
-    pub(crate) fn has_alpha(&self) -> bool {
+    pub(crate) const fn has_alpha(&self) -> bool {
         matches!(self, ColorType::GrayscaleAlpha | ColorType::RGBA)
     }
 
     #[inline]
-    pub(crate) fn has_trns(&self) -> bool {
+    pub(crate) const fn has_trns(&self) -> bool {
         match self {
             ColorType::Grayscale { transparent_shade } => transparent_shade.is_some(),
             ColorType::RGB { transparent_color } => transparent_color.is_some(),

src/deflate/mod.rs

@@ -7,10 +7,15 @@ pub use deflater::inflate;
 use std::{fmt, fmt::Display};
 
 #[cfg(feature = "zopfli")]
-use std::num::NonZeroU8;
+use std::io::{self, copy, BufWriter, Cursor, Write};
+
+#[cfg(feature = "zopfli")]
+use zopfli::{DeflateEncoder, Options};
 #[cfg(feature = "zopfli")]
 mod zopfli_oxipng;
 #[cfg(feature = "zopfli")]
+use simd_adler32::Adler32;
+#[cfg(feature = "zopfli")]
 pub use zopfli_oxipng::deflate as zopfli_deflate;
 
 #[derive(Clone, Copy, Debug, PartialEq, Eq)]
@@ -24,19 +29,21 @@ pub enum Deflaters {
     #[cfg(feature = "zopfli")]
     /// Use the better but slower Zopfli implementation
     Zopfli {
-        /// The number of compression iterations to do. 15 iterations are fine
-        /// for small files, but bigger files will need to be compressed with
-        /// less iterations, or else they will be too slow.
-        iterations: NonZeroU8,
+        /// Zopfli compression options
+        options: Options,
     },
 }
 
-impl Deflaters {
-    pub(crate) fn deflate(self, data: &[u8], max_size: &AtomicMin) -> PngResult<Vec<u8>> {
+pub trait Deflater: Sync + Send {
+    fn deflate(&self, data: &[u8], max_size: &AtomicMin) -> PngResult<Vec<u8>>;
+}
+
+impl Deflater for Deflaters {
+    fn deflate(&self, data: &[u8], max_size: &AtomicMin) -> PngResult<Vec<u8>> {
         let compressed = match self {
-            Self::Libdeflater { compression } => deflate(data, compression, max_size)?,
+            Self::Libdeflater { compression } => deflate(data, *compression, max_size)?,
             #[cfg(feature = "zopfli")]
-            Self::Zopfli { iterations } => zopfli_deflate(data, iterations)?,
+            Self::Zopfli { options } => zopfli_deflate(data, options)?,
         };
         if let Some(max) = max_size.get() {
             if compressed.len() > max {
@@ -47,6 +54,75 @@ impl Deflaters {
     }
 }
 
+#[cfg(feature = "zopfli")]
+#[derive(Copy, Clone, Debug)]
+pub struct BufferedZopfliDeflater {
+    input_buffer_size: usize,
+    output_buffer_size: usize,
+    options: Options,
+}
+
+#[cfg(feature = "zopfli")]
+impl BufferedZopfliDeflater {
+    pub const fn new(
+        input_buffer_size: usize,
+        output_buffer_size: usize,
+        options: Options,
+    ) -> Self {
+        BufferedZopfliDeflater {
+            input_buffer_size,
+            output_buffer_size,
+            options,
+        }
+    }
+}
+
+#[cfg(feature = "zopfli")]
+impl Default for BufferedZopfliDeflater {
+    fn default() -> Self {
+        BufferedZopfliDeflater {
+            input_buffer_size: 1024 * 1024,
+            output_buffer_size: 64 * 1024,
+            options: Options::default(),
+        }
+    }
+}
+
+#[cfg(feature = "zopfli")]
+impl Deflater for BufferedZopfliDeflater {
+    /// Fork of the zlib_compress function in Zopfli.
+    fn deflate(&self, data: &[u8], max_size: &AtomicMin) -> PngResult<Vec<u8>> {
+        let mut out = Cursor::new(Vec::with_capacity(self.output_buffer_size));
+        let cmf = 120; /* CM 8, CINFO 7. See zlib spec.*/
+        let flevel = 3;
+        let fdict = 0;
+        let mut cmfflg: u16 = 256 * cmf + fdict * 32 + flevel * 64;
+        let fcheck = 31 - cmfflg % 31;
+        cmfflg += fcheck;
+
+        let out = (|| -> io::Result<Vec<u8>> {
+            let mut rolling_adler = Adler32::new();
+            let mut in_data =
+                zopfli_oxipng::HashingAndCountingRead::new(data, &mut rolling_adler, None);
+            out.write_all(&cmfflg.to_be_bytes())?;
+            let mut buffer = BufWriter::with_capacity(
+                self.input_buffer_size,
+                DeflateEncoder::new(self.options, Default::default(), &mut out),
+            );
+            copy(&mut in_data, &mut buffer)?;
+            buffer.into_inner()?.finish()?;
+            out.write_all(&rolling_adler.finish().to_be_bytes())?;
+            Ok(out.into_inner())
+        })();
+        let out = out.map_err(|e| PngError::new(&e.to_string()))?;
+        if max_size.get().map(|max| max < out.len()).unwrap_or(false) {
+            Err(PngError::DeflatedDataTooLong(out.len()))
+        } else {
+            Ok(out)
+        }
+    }
+}
+
 impl Display for Deflaters {
     #[inline]
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {

src/deflate/zopfli_oxipng.rs

@@ -1,18 +1,64 @@
 use crate::{PngError, PngResult};
-use std::num::NonZeroU8;
+use simd_adler32::Adler32;
+use std::io::{Error, ErrorKind, Read};
 
-pub fn deflate(data: &[u8], iterations: NonZeroU8) -> PngResult<Vec<u8>> {
+pub fn deflate(data: &[u8], options: &zopfli::Options) -> PngResult<Vec<u8>> {
     use std::cmp::max;
 
     let mut output = Vec::with_capacity(max(1024, data.len() / 20));
-    let options = zopfli::Options {
-        iteration_count: iterations,
-        ..Default::default()
-    };
-    match zopfli::compress(&options, &zopfli::Format::Zlib, data, &mut output) {
+    match zopfli::compress(options, &zopfli::Format::Zlib, data, &mut output) {
         Ok(_) => (),
         Err(_) => return Err(PngError::new("Failed to compress in zopfli")),
     };
     output.shrink_to_fit();
     Ok(output)
 }
+
+/// Forked from zopfli crate
+pub trait Hasher {
+    fn update(&mut self, data: &[u8]);
+}
+
+impl Hasher for &mut Adler32 {
+    fn update(&mut self, data: &[u8]) {
+        Adler32::write(self, data)
+    }
+}
+
+/// A reader that wraps another reader, a hasher and an optional counter,
+/// updating the hasher state and incrementing a counter of bytes read so
+/// far for each block of data read.
+pub struct HashingAndCountingRead<'counter, R: Read, H: Hasher> {
+    inner: R,
+    hasher: H,
+    bytes_read: Option<&'counter mut u32>,
+}
+
+impl<'counter, R: Read, H: Hasher> HashingAndCountingRead<'counter, R, H> {
+    pub fn new(inner: R, hasher: H, bytes_read: Option<&'counter mut u32>) -> Self {
+        Self {
+            inner,
+            hasher,
+            bytes_read,
+        }
+    }
+}
+
+impl<R: Read, H: Hasher> Read for HashingAndCountingRead<'_, R, H> {
+    fn read(&mut self, buf: &mut [u8]) -> Result<usize, Error> {
+        match self.inner.read(buf) {
+            Ok(bytes_read) => {
+                self.hasher.update(&buf[..bytes_read]);
+
+                if let Some(total_bytes_read) = &mut self.bytes_read {
+                    **total_bytes_read = total_bytes_read
+                        .checked_add(bytes_read.try_into().map_err(|_| ErrorKind::Other)?)
+                        .ok_or(ErrorKind::Other)?;
+                }
+
+                Ok(bytes_read)
+            }
+            Err(err) => Err(err),
+        }
+    }
+}