Abort failing trials early

src/atomicmin.rs

@@ -0,0 +1,32 @@
+use std::sync::atomic::AtomicUsize;
+use std::sync::atomic::Ordering::{SeqCst, Relaxed};
+
+pub struct AtomicMin {
+    val: AtomicUsize,
+}
+
+impl AtomicMin {
+    pub fn new(init: Option<usize>) -> Self {
+        Self {
+            val: AtomicUsize::new(init.unwrap_or(usize::max_value()))
+        }
+    }
+
+    pub fn get(&self) -> Option<usize> {
+        let val = self.val.load(SeqCst);
+        if val == usize::max_value() {None} else {Some(val)}
+    }
+
+    pub fn set_min(&self, new_val: usize) {
+        let mut current_val = self.val.load(Relaxed);
+        loop {
+            if new_val < current_val {
+                if let Err(v) = self.val.compare_exchange(current_val, new_val, SeqCst, Relaxed) {
+                    current_val = v;
+                    continue;
+                }
+            }
+            break;
+        }
+    }
+}

src/deflate/miniz_stream.rs

@@ -1,10 +1,12 @@
+use error::PngError;
 use miniz_oxide::deflate::core::*;
 
-pub fn compress_to_vec_oxipng(input: &[u8], level: u8, window_bits: i32, strategy: i32) -> Vec<u8> {
+pub fn compress_to_vec_oxipng(input: &[u8], level: u8, window_bits: i32, strategy: i32, max_size: Option<usize>) -> Result<Vec<u8>, PngError> {
     // The comp flags function sets the zlib flag if the window_bits parameter is > 0.
     let flags = create_comp_flags_from_zip_params(level.into(), window_bits, strategy);
     let mut compressor = CompressorOxide::new(flags);
-    let mut output = Vec::with_capacity(input.len() / 2);
+    // if max size is known, then expect that much data (but no more than input.len())
+    let mut output = Vec::with_capacity(max_size.unwrap_or(input.len() / 2).min(input.len()));
     // # Unsafe
     // We trust compress to not read the uninitialized bytes.
     unsafe {
@@ -30,6 +32,11 @@ pub fn compress_to_vec_oxipng(input: &[u8], level: u8, window_bits: i32, strateg
                 break;
             }
             TDEFLStatus::Okay => {
+                if let Some(max) = max_size {
+                    if output.len() > max {
+                        return Err(PngError::DeflatedDataTooLong(output.len()))
+                    }
+                }
                 // We need more space, so extend the vector.
                 if output.len().saturating_sub(out_pos) < 30 {
                     let current_len = output.len();
@@ -48,5 +55,5 @@ pub fn compress_to_vec_oxipng(input: &[u8], level: u8, window_bits: i32, strateg
         }
     }
 
-    output
+    Ok(output)
 }

src/deflate/mod.rs

@@ -13,15 +13,15 @@ pub fn inflate(data: &[u8]) -> Result<Vec<u8>, PngError> {
 }
 
 /// Compress a data stream using the DEFLATE algorithm
-pub fn deflate(data: &[u8], zc: u8, zs: u8, zw: u8) -> Result<Vec<u8>, PngError> {
+pub fn deflate(data: &[u8], zc: u8, zs: u8, zw: u8, max_size: Option<usize>) -> Result<Vec<u8>, PngError> {
     #[cfg(feature = "cfzlib")]
     {
         if is_cfzlib_supported() {
-            return cfzlib_deflate(data, zc, zs, zw)
+            return cfzlib_deflate(data, zc, zs, zw, max_size)
         }
     }
 
-    Ok(miniz_stream::compress_to_vec_oxipng(data, zc, zw.into(), zs.into()))
+    miniz_stream::compress_to_vec_oxipng(data, zc, zw.into(), zs.into(), max_size)
 }
 
 #[cfg(feature = "cfzlib")]
@@ -40,7 +40,7 @@ fn is_cfzlib_supported() -> bool {
 }
 
 #[cfg(feature = "cfzlib")]
-pub fn cfzlib_deflate(data: &[u8], level: u8, strategy: u8, window_bits: u8) -> Result<Vec<u8>, PngError> {
+pub fn cfzlib_deflate(data: &[u8], level: u8, strategy: u8, window_bits: u8, max_size: Option<usize>) -> Result<Vec<u8>, PngError> {
     use std::mem;
     use cloudflare_zlib_sys::*;
 
@@ -57,7 +57,8 @@ pub fn cfzlib_deflate(data: &[u8], level: u8, strategy: u8, window_bits: u8) ->
             return Err(PngError::new("deflateInit2"));
         }
 
-        let max_size = deflateBound(&mut stream, data.len() as uLong) as usize;
+        let upper_bound = deflateBound(&mut stream, data.len() as uLong) as usize;
+        let max_size = max_size.unwrap_or(upper_bound).min(upper_bound);
         // it's important to have the capacity pre-allocated,
         // as unsafe set_len is called later
         let mut out = Vec::with_capacity(max_size);
@@ -67,8 +68,10 @@ pub fn cfzlib_deflate(data: &[u8], level: u8, strategy: u8, window_bits: u8) ->
         stream.avail_in = data.len() as uInt;
         stream.next_out = out.as_mut_ptr();
         stream.avail_out = out.capacity() as uInt;
-        if Z_STREAM_END != deflate(&mut stream, Z_FINISH) {
+        match deflate(&mut stream, Z_FINISH) {
-            return Err(PngError::new("deflate"));
+            Z_STREAM_END => {},
+            Z_OK => return Err(PngError::DeflatedDataTooLong(max_size)),
+            _ => return Err(PngError::new("deflate")),
         }
         if Z_OK != deflateEnd(&mut stream) {
             return Err(PngError::new("deflateEnd"));

src/error.rs

@@ -3,6 +3,7 @@ use std::fmt;
 
 #[derive(Debug, Clone)]
 pub enum PngError {
+    DeflatedDataTooLong(usize),
     Other(Box<str>),
 }
 
@@ -16,8 +17,9 @@ impl Error for PngError {
 impl fmt::Display for PngError {
     #[inline]
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
-        match self {
+        match *self {
-            PngError::Other(s) => f.write_str(s),
+            PngError::DeflatedDataTooLong(_) => f.write_str("deflated data too long"),
+            PngError::Other(ref s) => f.write_str(s),
         }
     }
 }

src/lib.rs

@@ -20,6 +20,7 @@ use std::collections::{HashMap, HashSet};
 use std::fs::{copy, File};
 use std::io::{stdout, BufWriter, Write};
 use std::path::{Path, PathBuf};
+use atomicmin::AtomicMin;
 
 pub use colors::AlphaOptim;
 pub use deflate::Deflaters;
@@ -39,6 +40,7 @@ mod interlace;
 #[doc(hidden)]
 pub mod png;
 mod reduction;
+mod atomicmin;
 
 #[derive(Clone, Debug)]
 /// Options controlling the output of the `optimize` function
@@ -493,20 +495,36 @@ fn optimize_png(
         let original_len = original_png.idat_data.len();
         let added_interlacing = opts.interlace == Some(1) && original_png.ihdr_data.interlaced == 0;
 
+        let best_size = AtomicMin::new(if opts.force {None} else {Some(original_len)});
         let best: Option<TrialWithData> = results
             .into_par_iter()
             .with_max_len(1)
             .filter_map(|trial| {
                 let filtered = &filters[&trial.filter];
                 let new_idat = if opts.deflate == Deflaters::Zlib {
-                    deflate::deflate(filtered, trial.compression, trial.strategy, opts.window)
+                    deflate::deflate(filtered, trial.compression, trial.strategy, opts.window, best_size.get())
                 } else {
                     deflate::zopfli_deflate(filtered)
                 };
-                let new_idat = if let Ok(n) = new_idat {n} else {
+                let new_idat = match new_idat {
-                    return None;
+                    Ok(n) => n,
+                    Err(PngError::DeflatedDataTooLong(max)) if opts.verbosity == Some(1) => {
+                        eprintln!(
+                            "    zc = {}  zs = {}  f = {}       >{} bytes",
+                            trial.compression,
+                            trial.strategy,
+                            trial.filter,
+                            max,
+                        );
+                        return None;
+                    },
+                    _ => return None,
                 };
 
+                // update best size across all threads
+                let new_size = new_idat.len();
+                best_size.set_min(new_size);
+
                 if opts.verbosity == Some(1) {
                     eprintln!(
                         "    zc = {}  zs = {}  f = {}        {} bytes",
@@ -517,7 +535,7 @@ fn optimize_png(
                     );
                 }
 
-                if new_idat.len() < original_len || added_interlacing || opts.force {
+                if new_size < original_len || added_interlacing || opts.force {
                     Some((trial, new_idat))
                 } else {
                     None

src/png/mod.rs

@@ -16,6 +16,7 @@ use std::io::{Read, Seek, SeekFrom};
 use std::iter::Iterator;
 use std::path::Path;
 use rayon::prelude::*;
+use atomicmin::AtomicMin;
 
 const STD_COMPRESSION: u8 = 8;
 const STD_STRATEGY: u8 = 2; // Huffman only
@@ -603,6 +604,7 @@ impl PngData {
     pub fn try_alpha_reduction(&mut self, alphas: &HashSet<AlphaOptim>) {
         assert!(!alphas.is_empty());
         let alphas = alphas.iter().collect::<Vec<_>>();
+        let best_size = AtomicMin::new(None);
         let best = alphas
             .par_iter()
             .with_max_len(1)
@@ -618,8 +620,12 @@ impl PngData {
                             STD_COMPRESSION,
                             STD_STRATEGY,
                             STD_WINDOW,
+                            best_size.get(),
                         ).ok()
-                        .as_ref().map(|l| l.len())
+                        .as_ref().map(|l| {
+                            best_size.set_min(l.len());
+                            l.len()
+                        })
                     })
                     .min()
                     .map(|size| (size, image))