diff --git a/src/atomicmin.rs b/src/atomicmin.rs new file mode 100644 index 00000000..de87af48 --- /dev/null +++ b/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) -> Self { + Self { + val: AtomicUsize::new(init.unwrap_or(usize::max_value())) + } + } + + pub fn get(&self) -> Option { + 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; + } + } +} diff --git a/src/deflate/miniz_stream.rs b/src/deflate/miniz_stream.rs index f9b44536..2438901d 100644 --- a/src/deflate/miniz_stream.rs +++ b/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 { +pub fn compress_to_vec_oxipng(input: &[u8], level: u8, window_bits: i32, strategy: i32, max_size: Option) -> Result, 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) } diff --git a/src/deflate/mod.rs b/src/deflate/mod.rs index 4d85090b..6a499511 100644 --- a/src/deflate/mod.rs +++ b/src/deflate/mod.rs @@ -13,15 +13,15 @@ pub fn inflate(data: &[u8]) -> Result, PngError> { } /// Compress a data stream using the DEFLATE algorithm -pub fn deflate(data: &[u8], zc: u8, zs: u8, zw: u8) -> Result, PngError> { +pub fn deflate(data: &[u8], zc: u8, zs: u8, zw: u8, max_size: Option) -> Result, 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, PngError> { +pub fn cfzlib_deflate(data: &[u8], level: u8, strategy: u8, window_bits: u8, max_size: Option) -> Result, 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) { - return Err(PngError::new("deflate")); + match deflate(&mut stream, Z_FINISH) { + 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")); diff --git a/src/error.rs b/src/error.rs index 05a1b131..de2c5a7f 100644 --- a/src/error.rs +++ b/src/error.rs @@ -3,6 +3,7 @@ use std::fmt; #[derive(Debug, Clone)] pub enum PngError { + DeflatedDataTooLong(usize), Other(Box), } @@ -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 { - PngError::Other(s) => f.write_str(s), + match *self { + PngError::DeflatedDataTooLong(_) => f.write_str("deflated data too long"), + PngError::Other(ref s) => f.write_str(s), } } } diff --git a/src/lib.rs b/src/lib.rs index a49eaf2a..41736114 100644 --- a/src/lib.rs +++ b/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 = 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 { - return None; + let new_idat = match new_idat { + 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 diff --git a/src/png/mod.rs b/src/png/mod.rs index a213969c..9304bb28 100644 --- a/src/png/mod.rs +++ b/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) { assert!(!alphas.is_empty()); let alphas = alphas.iter().collect::>(); + 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))