#![warn(trivial_casts, trivial_numeric_casts, unused_import_braces)]
#![deny(missing_debug_implementations, missing_copy_implementations)]
#![warn(clippy::expl_impl_clone_on_copy)]
#![warn(clippy::float_cmp_const)]
#![warn(clippy::linkedlist)]
#![warn(clippy::map_flatten)]
#![warn(clippy::match_same_arms)]
#![warn(clippy::mem_forget)]
#![warn(clippy::mut_mut)]
#![warn(clippy::mutex_integer)]
#![warn(clippy::needless_continue)]
#![warn(clippy::path_buf_push_overwrite)]
#![warn(clippy::range_plus_one)]
#![allow(clippy::cognitive_complexity)]
#![allow(clippy::upper_case_acronyms)]
#![cfg_attr(
not(feature = "zopfli"),
allow(irrefutable_let_patterns),
allow(unreachable_patterns)
)]
#[cfg(feature = "parallel")]
extern crate rayon;
#[cfg(not(feature = "parallel"))]
mod rayon;
use std::{
borrow::Cow,
fs::{File, Metadata},
io::{stdin, stdout, BufWriter, Read, Write},
path::Path,
sync::{
atomic::{AtomicBool, Ordering},
Arc,
},
time::{Duration, Instant},
};
pub use indexmap::{indexset, IndexSet};
use log::{debug, info, trace, warn};
use rayon::prelude::*;
pub use rgb::{RGB16, RGBA8};
use crate::{
atomicmin::AtomicMin,
evaluate::Evaluator,
headers::*,
png::{PngData, PngImage},
reduction::*,
};
pub use crate::{
colors::{BitDepth, ColorType},
deflate::Deflaters,
error::PngError,
filters::RowFilter,
headers::StripChunks,
interlace::Interlacing,
options::{InFile, Options, OutFile},
};
mod atomicmin;
mod colors;
mod deflate;
mod display_chunks;
mod error;
mod evaluate;
mod filters;
mod headers;
mod interlace;
mod options;
mod png;
mod reduction;
#[cfg(feature = "sanity-checks")]
mod sanity_checks;
#[doc(hidden)]
pub mod internal_tests {
#[cfg(feature = "sanity-checks")]
pub use crate::sanity_checks::*;
pub use crate::{atomicmin::*, deflate::*, png::*, reduction::*};
}
pub type PngResult<T> = Result<T, PngError>;
#[derive(Debug)]
pub struct RawImage {
png: Arc<PngImage>,
aux_chunks: Vec<Chunk>,
}
impl RawImage {
pub fn new(
width: u32,
height: u32,
color_type: ColorType,
bit_depth: BitDepth,
data: Vec<u8>,
) -> Result<Self, PngError> {
let valid_depth = match color_type {
ColorType::Grayscale { .. } => true,
ColorType::Indexed { .. } => (bit_depth as u8) <= 8,
_ => (bit_depth as u8) >= 8,
};
if !valid_depth {
return Err(PngError::InvalidDepthForType(bit_depth, color_type));
}
let bpp = bit_depth as usize * color_type.channels_per_pixel() as usize;
let row_bytes = (bpp * width as usize + 7) / 8;
let expected_len = row_bytes * height as usize;
if data.len() != expected_len {
return Err(PngError::IncorrectDataLength(data.len(), expected_len));
}
Ok(Self {
png: Arc::new(PngImage {
ihdr: IhdrData {
width,
height,
color_type,
bit_depth,
interlaced: Interlacing::None,
},
data,
}),
aux_chunks: Vec::new(),
})
}
pub fn add_png_chunk(&mut self, name: [u8; 4], data: Vec<u8>) {
self.aux_chunks.push(Chunk { name, data });
}
pub fn add_icc_profile(&mut self, data: &[u8]) {
let deflater = Deflaters::Libdeflater { compression: 1 };
if let Ok(iccp) = construct_iccp(data, deflater) {
self.aux_chunks.push(iccp);
}
}
pub fn create_optimized_png(&self, opts: &Options) -> PngResult<Vec<u8>> {
let deadline = Arc::new(Deadline::new(opts.timeout));
let mut png = optimize_raw(self.png.clone(), opts, deadline.clone(), None)
.ok_or_else(|| PngError::new("Failed to optimize input data"))?;
png.aux_chunks = self
.aux_chunks
.iter()
.filter(|c| opts.strip.keep(&c.name))
.cloned()
.collect();
postprocess_chunks(&mut png, opts, deadline, &self.png.ihdr);
Ok(png.output())
}
}
pub fn optimize(input: &InFile, output: &OutFile, opts: &Options) -> PngResult<()> {
info!("Processing: {}", input);
let deadline = Arc::new(Deadline::new(opts.timeout));
let opt_metadata_preserved;
let in_data = match *input {
InFile::Path(ref input_path) => {
if matches!(
output,
OutFile::Path {
preserve_attrs: true,
..
}
) {
opt_metadata_preserved = input_path
.metadata()
.map_err(|err| {
PngError::new(&format!(
"Unable to read metadata from input file {:?}: {}",
input_path, err
))
})
.map(Some)?;
trace!("preserving metadata: {:?}", opt_metadata_preserved);
} else {
opt_metadata_preserved = None;
}
PngData::read_file(input_path)?
}
InFile::StdIn => {
opt_metadata_preserved = None;
let mut data = Vec::new();
stdin()
.read_to_end(&mut data)
.map_err(|e| PngError::new(&format!("Error reading stdin: {}", e)))?;
data
}
};
let mut png = PngData::from_slice(&in_data, opts)?;
let mut optimized_output = optimize_png(&mut png, &in_data, opts, deadline)?;
let in_length = in_data.len();
if is_fully_optimized(in_data.len(), optimized_output.len(), opts) {
match (output, input) {
(OutFile::Path { path, .. }, InFile::Path(ref input_path))
if path.as_ref().map_or(true, |p| p == input_path) =>
{
info!("{}: Could not optimize further, no change written", input);
return Ok(());
}
_ => {
optimized_output = in_data;
}
}
}
let savings = if in_length >= optimized_output.len() {
format!(
"{} bytes ({:.2}% smaller)",
optimized_output.len(),
(in_length - optimized_output.len()) as f64 / in_length as f64 * 100_f64
)
} else {
format!(
"{} bytes ({:.2}% larger)",
optimized_output.len(),
(optimized_output.len() - in_length) as f64 / in_length as f64 * 100_f64
)
};
match (output, input) {
(OutFile::None, _) => {
info!("{}: Running in pretend mode, no output", savings);
}
(&OutFile::StdOut, _) | (&OutFile::Path { path: None, .. }, &InFile::StdIn) => {
let mut buffer = BufWriter::new(stdout());
buffer
.write_all(&optimized_output)
.map_err(|e| PngError::new(&format!("Unable to write to stdout: {}", e)))?;
}
(OutFile::Path { path, .. }, _) => {
let output_path = path
.as_ref()
.map(|p| p.as_path())
.unwrap_or_else(|| input.path().unwrap());
let out_file = File::create(output_path).map_err(|err| {
PngError::new(&format!(
"Unable to write to file {}: {}",
output_path.display(),
err
))
})?;
if let Some(metadata_input) = &opt_metadata_preserved {
copy_permissions(metadata_input, &out_file)?;
}
let mut buffer = BufWriter::new(out_file);
buffer
.write_all(&optimized_output)
.and_then(|()| buffer.flush())
.map_err(|e| {
PngError::new(&format!(
"Unable to write to {}: {}",
output_path.display(),
e
))
})?;
std::mem::drop(buffer);
if let Some(metadata_input) = &opt_metadata_preserved {
copy_times(metadata_input, output_path)?;
}
info!("{}: {}", savings, output_path.display());
}
}
Ok(())
}
pub fn optimize_from_memory(data: &[u8], opts: &Options) -> PngResult<Vec<u8>> {
info!("Processing from memory");
let deadline = Arc::new(Deadline::new(opts.timeout));
let original_size = data.len();
let mut png = PngData::from_slice(data, opts)?;
let optimized_output = optimize_png(&mut png, data, opts, deadline)?;
if is_fully_optimized(original_size, optimized_output.len(), opts) {
info!("Image already optimized");
Ok(data.to_vec())
} else {
Ok(optimized_output)
}
}
type TrialResult = (RowFilter, Vec<u8>);
fn optimize_png(
png: &mut PngData,
original_data: &[u8],
opts: &Options,
deadline: Arc<Deadline>,
) -> PngResult<Vec<u8>> {
let file_original_size = original_data.len();
let idat_original_size = png.idat_data.len();
let raw = png.raw.clone();
debug!(
" {}x{} pixels, PNG format",
raw.ihdr.width, raw.ihdr.height
);
report_format(" ", &raw);
debug!(" IDAT size = {} bytes", idat_original_size);
debug!(" File size = {} bytes", file_original_size);
let opts = if png.aux_chunks.iter().any(|c| &c.name == b"acTL") {
warn!("APNG detected, disabling all reductions");
let mut opts = opts.to_owned();
opts.interlace = None;
opts.bit_depth_reduction = false;
opts.color_type_reduction = false;
opts.palette_reduction = false;
opts.grayscale_reduction = false;
Cow::Owned(opts)
} else {
Cow::Borrowed(opts)
};
let max_size = if opts.force {
None
} else {
Some(png.estimated_output_size())
};
if let Some(new_png) = optimize_raw(raw.clone(), &opts, deadline.clone(), max_size) {
png.raw = new_png.raw;
png.idat_data = new_png.idat_data;
}
postprocess_chunks(png, &opts, deadline, &raw.ihdr);
let output = png.output();
if idat_original_size >= png.idat_data.len() {
debug!(
" IDAT size = {} bytes ({} bytes decrease)",
png.idat_data.len(),
idat_original_size - png.idat_data.len()
);
} else {
debug!(
" IDAT size = {} bytes ({} bytes increase)",
png.idat_data.len(),
png.idat_data.len() - idat_original_size
);
}
if file_original_size >= output.len() {
debug!(
" file size = {} bytes ({} bytes = {:.2}% decrease)",
output.len(),
file_original_size - output.len(),
(file_original_size - output.len()) as f64 / file_original_size as f64 * 100_f64
);
} else {
debug!(
" file size = {} bytes ({} bytes = {:.2}% increase)",
output.len(),
output.len() - file_original_size,
(output.len() - file_original_size) as f64 / file_original_size as f64 * 100_f64
);
}
#[cfg(feature = "sanity-checks")]
assert!(sanity_checks::validate_output(&output, original_data));
Ok(output)
}
fn optimize_raw(
image: Arc<PngImage>,
opts: &Options,
deadline: Arc<Deadline>,
max_size: Option<usize>,
) -> Option<PngData> {
let eval_compression = match opts.deflate {
Deflaters::Libdeflater { compression } => {
if opts.fast_evaluation { 7 } else { 8 }.min(compression)
}
_ => 8,
};
let eval_filters = if opts.filter.len() == 1 {
opts.filter.clone()
} else {
indexset! {RowFilter::None, RowFilter::Bigrams}
};
let eval = Evaluator::new(
deadline.clone(),
eval_filters.clone(),
eval_compression,
false,
);
let mut png = perform_reductions(image.clone(), opts, &deadline, &eval);
let mut eval_result = eval.get_best_candidate();
if let Some(ref result) = eval_result {
png = result.image.clone();
}
let reduction_occurred = png.ihdr.color_type != image.ihdr.color_type
|| png.ihdr.bit_depth != image.ihdr.bit_depth
|| png.ihdr.interlaced != image.ihdr.interlaced;
if reduction_occurred {
report_format("Transformed image to ", &png);
}
if opts.idat_recoding || reduction_occurred {
let mut filters = opts.filter.clone();
let fast_eval = opts.fast_evaluation && (filters.len() > 1 || eval_result.is_some());
let best: Option<TrialResult> = if fast_eval {
if eval_result.is_some() {
filters = filters.difference(&eval_filters).cloned().collect();
}
if !filters.is_empty() {
trace!("Evaluating: {} filters", filters.len());
let eval = Evaluator::new(deadline, filters, eval_compression, opts.optimize_alpha);
if let Some(ref result) = eval_result {
eval.set_best_size(result.idat_data.len());
}
eval.try_image(png.clone());
if let Some(result) = eval.get_best_candidate() {
eval_result = Some(result);
}
}
let result = eval_result?;
match opts.deflate {
Deflaters::Libdeflater { compression } if compression <= eval_compression => {
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)
}
}
} else {
if filters.is_empty() {
if png.ihdr.bit_depth as u8 >= 8 {
filters.insert(RowFilter::Bigrams);
} else {
filters.insert(RowFilter::None);
}
}
debug!("Trying: {} filters", filters.len());
let best_size = AtomicMin::new(max_size);
let results_iter = filters.into_par_iter().with_max_len(1);
let best = results_iter.filter_map(|filter| {
if deadline.passed() {
return None;
}
let filtered = &png.filter_image(filter, opts.optimize_alpha);
perform_trial(filtered, opts, filter, &best_size)
});
best.reduce_with(|i, j| {
if i.1.len() < j.1.len() || (i.1.len() == j.1.len() && i.0 < j.0) {
i
} else {
j
}
})
};
if let Some((filter, idat_data)) = best {
let image = PngData {
raw: png,
idat_data,
aux_chunks: Vec::new(),
};
if image.estimated_output_size() < max_size.unwrap_or(usize::MAX) {
debug!("Found better combination:");
debug!(
" zc = {} f = {:8} {} bytes",
opts.deflate,
filter,
image.idat_data.len()
);
return Some(image);
}
}
} else if let Some(result) = eval_result {
let image = PngData {
raw: result.image,
idat_data: result.idat_data,
aux_chunks: Vec::new(),
};
if image.estimated_output_size() < max_size.unwrap_or(usize::MAX) {
debug!("Found better combination:");
debug!(
" zc = {} f = {:8} {} bytes",
eval_compression,
result.filter,
image.idat_data.len()
);
return Some(image);
}
}
None
}
fn perform_trial(
filtered: &[u8],
opts: &Options,
filter: RowFilter,
best_size: &AtomicMin,
) -> Option<TrialResult> {
match opts.deflate.deflate(filtered, best_size) {
Ok(new_idat) => {
let bytes = new_idat.len();
best_size.set_min(bytes);
trace!(
" zc = {} f = {:8} {} bytes",
opts.deflate,
filter,
bytes
);
Some((filter, new_idat))
}
Err(PngError::DeflatedDataTooLong(bytes)) => {
trace!(
" zc = {} f = {:8} >{} bytes",
opts.deflate,
filter,
bytes,
);
None
}
Err(_) => None,
}
}
#[derive(Debug)]
struct DeadlineImp {
start: Instant,
timeout: Duration,
print_message: AtomicBool,
}
#[doc(hidden)]
#[derive(Debug)]
pub struct Deadline {
imp: Option<DeadlineImp>,
}
impl Deadline {
pub fn new(timeout: Option<Duration>) -> Self {
Self {
imp: timeout.map(|timeout| DeadlineImp {
start: Instant::now(),
timeout,
print_message: AtomicBool::new(true),
}),
}
}
pub fn passed(&self) -> bool {
if let Some(imp) = &self.imp {
let elapsed = imp.start.elapsed();
if elapsed > imp.timeout {
if match imp.print_message.compare_exchange(
true,
false,
Ordering::SeqCst,
Ordering::SeqCst,
) {
Ok(x) | Err(x) => x,
} {
warn!("Timed out after {} second(s)", elapsed.as_secs());
}
return true;
}
}
false
}
}
fn report_format(prefix: &str, png: &PngImage) {
debug!(
"{}{}-bit {}, {}",
prefix, png.ihdr.bit_depth, png.ihdr.color_type, png.ihdr.interlaced
);
}
fn postprocess_chunks(
png: &mut PngData,
opts: &Options,
deadline: Arc<Deadline>,
orig_ihdr: &IhdrData,
) {
if let Some(iccp_idx) = png.aux_chunks.iter().position(|c| &c.name == b"iCCP") {
let may_replace_iccp = opts.strip != StripChunks::None && opts.strip.keep(b"sRGB");
if may_replace_iccp && png.aux_chunks.iter().any(|c| &c.name == b"sRGB") {
trace!("Removing iCCP chunk due to conflict with sRGB chunk");
png.aux_chunks.remove(iccp_idx);
} else if let Some(icc) = extract_icc(&png.aux_chunks[iccp_idx]) {
let intent = if may_replace_iccp {
srgb_rendering_intent(&icc)
} else {
None
};
if let Some(intent) = intent {
trace!("Replacing iCCP chunk with equivalent sRGB chunk");
png.aux_chunks[iccp_idx] = Chunk {
name: *b"sRGB",
data: vec![intent],
};
} else if opts.idat_recoding {
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 {
debug!(
"Recompressed iCCP chunk: {} ({} bytes decrease)",
new_len,
cur_len - new_len
);
png.aux_chunks[iccp_idx] = iccp;
}
}
}
}
}
let ihdr = &png.raw.ihdr;
if orig_ihdr.bit_depth != ihdr.bit_depth || orig_ihdr.color_type != ihdr.color_type {
png.aux_chunks.retain(|c| {
let invalid = &c.name == b"bKGD" || &c.name == b"sBIT" || &c.name == b"hIST";
if invalid {
warn!(
"Removing {} chunk as it no longer matches the image data",
std::str::from_utf8(&c.name).unwrap()
);
}
!invalid
});
}
let mut fdat: Vec<_> = png
.aux_chunks
.iter_mut()
.filter(|c| &c.name == b"fdAT")
.collect();
if opts.idat_recoding && !fdat.is_empty() {
let buffer_size = orig_ihdr.raw_data_size();
fdat.par_iter_mut()
.with_max_len(1)
.enumerate()
.for_each(|(i, c)| {
if deadline.passed() || c.data.len() <= 4 {
return;
}
if let Ok(mut data) = deflate::inflate(&c.data[4..], buffer_size).and_then(|data| {
let max_size = AtomicMin::new(Some(c.data.len() - 5));
opts.deflate.deflate(&data, &max_size)
}) {
debug!(
"Recompressed fdAT #{:<2}: {} ({} bytes decrease)",
i,
c.data.len(),
c.data.len() - 4 - data.len()
);
c.data.truncate(4);
c.data.append(&mut data);
}
})
}
}
fn is_fully_optimized(original_size: usize, optimized_size: usize, opts: &Options) -> bool {
original_size <= optimized_size && !opts.force
}
fn copy_permissions(metadata_input: &Metadata, out_file: &File) -> PngResult<()> {
out_file
.set_permissions(metadata_input.permissions())
.map_err(|err_io| {
PngError::new(&format!(
"unable to set permissions for output file: {}",
err_io
))
})
}
#[cfg(not(feature = "filetime"))]
fn copy_times(_: &Metadata, _: &Path) -> PngResult<()> {
Ok(())
}
#[cfg(feature = "filetime")]
fn copy_times(input_path_meta: &Metadata, out_path: &Path) -> PngResult<()> {
let atime = filetime::FileTime::from_last_access_time(input_path_meta);
let mtime = filetime::FileTime::from_last_modification_time(input_path_meta);
trace!(
"attempting to set file times: atime: {:?}, mtime: {:?}",
atime,
mtime
);
filetime::set_file_times(out_path, atime, mtime).map_err(|err_io| {
PngError::new(&format!(
"unable to set file times on {:?}: {}",
out_path, err_io
))
})
}