Skip to content

Commit

Permalink
Experimental: allow Zopfli to use any size BufWriter
Browse files Browse the repository at this point in the history
  • Loading branch information
Pr0methean committed Jun 18, 2023
1 parent 96122fa commit 97248f7
Show file tree
Hide file tree
Showing 4 changed files with 96 additions and 33 deletions.
23 changes: 16 additions & 7 deletions benches/zopfli.rs
Original file line number Diff line number Diff line change
Expand Up @@ -9,26 +9,32 @@ 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) };

const DEFAULT_DEFLATER: BufferedZopfliDeflater = BufferedZopfliDeflater::new(
// SAFETY: trivially safe. Stopgap solution until const unwrap is stabilized.
unsafe { NonZeroU8::new_unchecked(15) },
4 * 1024 * 1024
);

#[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();
DEFAULT_DEFLATER.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();
DEFAULT_DEFLATER.deflate(png.raw.data.as_ref(), &max_size).ok();
});
}

Expand All @@ -38,9 +44,10 @@ 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();
DEFAULT_DEFLATER.deflate(png.raw.data.as_ref(), &max_size).ok();
});
}

Expand All @@ -50,9 +57,10 @@ 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();
DEFAULT_DEFLATER.deflate(png.raw.data.as_ref(), &max_size).ok();
});
}

Expand All @@ -62,8 +70,9 @@ 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();
DEFAULT_DEFLATER.deflate(png.raw.data.as_ref(), &max_size).ok();
});
}
55 changes: 50 additions & 5 deletions src/deflate/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -4,10 +4,14 @@ use crate::{PngError, PngResult};
pub use deflater::crc32;
pub use deflater::deflate;
pub use deflater::inflate;
use std::{fmt, fmt::Display};
use std::{fmt, fmt::Display, io};
use std::io::{BufWriter, Cursor, Write};

#[cfg(feature = "zopfli")]
use std::num::NonZeroU8;
#[cfg(feature = "zopfli")]
use zopfli::{DeflateEncoder, Options};

#[cfg(feature = "zopfli")]
mod zopfli_oxipng;
#[cfg(feature = "zopfli")]
Expand All @@ -31,12 +35,16 @@ pub enum Deflaters {
},
}

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 { iterations } => zopfli_deflate(data, *iterations)?,
};
if let Some(max) = max_size.get() {
if compressed.len() > max {
Expand All @@ -47,6 +55,43 @@ impl Deflaters {
}
}

#[cfg(feature = "zopfli")]
#[derive(Copy, Clone, Debug)]
pub struct BufferedZopfliDeflater {
iterations: NonZeroU8,
buffer_size: usize
}

impl BufferedZopfliDeflater {
pub const fn new(iterations: NonZeroU8,
buffer_size: usize) -> Self {
BufferedZopfliDeflater {iterations, buffer_size}
}
}

#[cfg(feature = "zopfli")]
impl Deflater for BufferedZopfliDeflater {
fn deflate(&self, data: &[u8], max_size: &AtomicMin) -> PngResult<Vec<u8>> {
let options = Options {
iteration_count: self.iterations,
..Default::default()
};
let mut buffer = BufWriter::with_capacity(self.buffer_size,
DeflateEncoder::new(
options, Default::default(), Cursor::new(Vec::new())));
let result = (|| -> io::Result<Vec<u8>> {
buffer.write_all(data)?;
Ok(buffer.into_inner()?.finish()?.into_inner())
})();
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()))
} else {
Ok(result)
}
}
}

impl Display for Deflaters {
#[inline]
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
Expand Down
5 changes: 2 additions & 3 deletions src/headers.rs
Original file line number Diff line number Diff line change
@@ -1,9 +1,8 @@
use crate::colors::{BitDepth, ColorType};
use crate::deflate::{crc32, inflate};
use crate::deflate::{crc32, Deflater, inflate};
use crate::error::PngError;
use crate::interlace::Interlacing;
use crate::AtomicMin;
use crate::Deflaters;
use crate::PngResult;
use indexmap::IndexSet;
use log::warn;
Expand Down Expand Up @@ -247,7 +246,7 @@ pub fn extract_icc(iccp: &Chunk) -> Option<Vec<u8>> {
}

/// Construct an iCCP chunk by compressing the ICC profile
pub fn construct_iccp(icc: &[u8], deflater: Deflaters) -> PngResult<Chunk> {
pub fn construct_iccp<T: Deflater>(icc: &[u8], deflater: &T) -> PngResult<Chunk> {
let mut compressed = deflater.deflate(icc, &AtomicMin::new(None))?;
let mut data = Vec::with_capacity(compressed.len() + 5);
data.extend(b"icc"); // Profile name - generally unused, can be anything
Expand Down
46 changes: 28 additions & 18 deletions src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,7 @@ use crate::headers::*;
use crate::png::PngData;
use crate::png::PngImage;
use crate::reduction::*;
use log::{debug, info, trace, warn};
use log::{debug, error, info, trace, warn};
use rayon::prelude::*;
use std::fmt;
use std::fs::{copy, File, Metadata};
Expand All @@ -48,6 +48,7 @@ pub use crate::headers::StripChunks;
pub use crate::interlace::Interlacing;
pub use indexmap::{indexset, IndexSet};
pub use rgb::{RGB16, RGBA8};
use crate::deflate::Deflater;

mod atomicmin;
mod colors;
Expand Down Expand Up @@ -380,15 +381,17 @@ impl RawImage {
pub fn add_icc_profile(&mut self, data: &[u8]) {
// Compress with fastest compression level - will be recompressed during optimization
let deflater = Deflaters::Libdeflater { compression: 1 };
if let Ok(iccp) = construct_iccp(data, deflater) {
if let Ok(iccp) = construct_iccp(data, &deflater) {
self.aux_chunks.push(iccp);
}
}

/// Create an optimized png from the raw image data using the options provided
pub fn create_optimized_png(&self, opts: &Options) -> PngResult<Vec<u8>> {
pub fn create_optimized_png<T: Deflater>
(&self, opts: &Options, deflater: &T) -> PngResult<Vec<u8>> {
let deadline = Arc::new(Deadline::new(opts.timeout));
let mut png = optimize_raw(self.png.clone(), opts, deadline, None)
let mut png = optimize_raw(self.png.clone(),
opts, deadline, None, deflater)
.ok_or_else(|| PngError::new("Failed to optimize input data"))?;

// Process aux chunks
Expand Down Expand Up @@ -511,7 +514,7 @@ pub fn optimize(input: &InFile, output: &OutFile, opts: &Options) -> PngResult<(
))
})?;
// force drop and thereby closing of file handle before modifying any timestamp
std::mem::drop(buffer);
drop(buffer);
if let Some(metadata_input) = &opt_metadata_preserved {
copy_times(metadata_input, output_path)?;
}
Expand Down Expand Up @@ -569,7 +572,8 @@ fn optimize_png(
} else {
Some(png.estimated_output_size())
};
if let Some(new_png) = optimize_raw(raw.clone(), opts, deadline, max_size) {
if let Some(new_png) = optimize_raw(raw.clone(), opts, deadline, max_size,
&opts.deflate) {
png.raw = new_png.raw;
png.idat_data = new_png.idat_data;
}
Expand Down Expand Up @@ -614,11 +618,12 @@ fn optimize_png(
}

/// Perform optimization on the input image data using the options provided
fn optimize_raw(
fn optimize_raw<T: Deflater>(
image: Arc<PngImage>,
opts: &Options,
deadline: Arc<Deadline>,
max_size: Option<usize>,
deflater: &T
) -> Option<PngData> {
// Must use normal (lazy) compression, as faster ones (greedy) are not representative
let eval_compression = 5;
Expand Down Expand Up @@ -677,7 +682,7 @@ fn optimize_raw(
_ => {
debug!("Trying: {}", result.filter);
let best_size = AtomicMin::new(max_size);
perform_trial(&result.filtered, opts, result.filter, &best_size)
perform_trial(&result.filtered, opts, result.filter, &best_size, deflater)
}
}
} else {
Expand All @@ -703,7 +708,7 @@ fn optimize_raw(
return None;
}
let filtered = &png.filter_image(filter, opts.optimize_alpha);
perform_trial(filtered, opts, filter, &best_size)
perform_trial(filtered, opts, filter, &best_size, deflater)
});
best.reduce_with(|i, j| {
if i.1.len() < j.1.len() || (i.1.len() == j.1.len() && i.0 < j.0) {
Expand Down Expand Up @@ -755,13 +760,15 @@ fn optimize_raw(
}

/// Execute a compression trial
fn perform_trial(
fn perform_trial<T: Deflater>(
filtered: &[u8],
opts: &Options,
filter: RowFilter,
best_size: &AtomicMin,
deflater: &T
) -> Option<TrialResult> {
match opts.deflate.deflate(filtered, best_size) {
let result = deflater.deflate(filtered, best_size);
match result {
Ok(new_idat) => {
let bytes = new_idat.len();
best_size.set_min(bytes);
Expand All @@ -775,14 +782,17 @@ fn perform_trial(
}
Err(PngError::DeflatedDataTooLong(bytes)) => {
trace!(
" zc = {} f = {:8} >{} bytes",
opts.deflate,
filter,
bytes,
);
" zc = {} f = {:8} >{} bytes",
opts.deflate,
filter,
bytes,
);
None
}
Err(e) => {
error!("I/O error: {}", e);
None
}
Err(_) => None,
}
}

Expand Down Expand Up @@ -866,7 +876,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, &opts.deflate) {
let cur_len = png.aux_chunks[iccp_idx].data.len();
let new_len = iccp.data.len();
if new_len < cur_len {
Expand Down

0 comments on commit 97248f7

Please sign in to comment.