Refactor use_heuristics for fast filter evaluation (#463)

src/atomicmin.rs

@@ -27,7 +27,8 @@ impl AtomicMin {
         &self.val
     }
 
-    pub fn set_min(&self, new_val: usize) {
+    /// Try a new value, returning true if it is the new minimum
-        self.val.fetch_min(new_val, SeqCst);
+    pub fn set_min(&self, new_val: usize) -> bool {
+        new_val < self.val.fetch_min(new_val, SeqCst)
     }
 }

src/deflate/mod.rs

@@ -1,5 +1,3 @@
-use indexmap::IndexSet;
-
 mod deflater;
 pub use deflater::crc32;
 pub use deflater::deflate;
@@ -12,13 +10,13 @@ mod zopfli_oxipng;
 #[cfg(feature = "zopfli")]
 pub use zopfli_oxipng::deflate as zopfli_deflate;
 
-#[derive(Clone, Debug, PartialEq, Eq)]
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
 /// DEFLATE algorithms supported by oxipng
 pub enum Deflaters {
     /// Use libdeflater.
     Libdeflater {
-        /// Which compression levels to try on the file (1-12)
+        /// Which compression level to use on the file (1-12)
-        compression: IndexSet<u8>,
+        compression: u8,
     },
     #[cfg(feature = "zopfli")]
     /// Use the better but slower Zopfli implementation

src/evaluate.rs

@@ -11,6 +11,7 @@ use crate::rayon;
 use crate::Deadline;
 #[cfg(feature = "parallel")]
 use crossbeam_channel::{unbounded, Receiver, Sender};
+use indexmap::IndexSet;
 use rayon::prelude::*;
 #[cfg(not(feature = "parallel"))]
 use std::cell::RefCell;
@@ -18,13 +19,10 @@ use std::sync::atomic::AtomicUsize;
 use std::sync::atomic::Ordering::SeqCst;
 use std::sync::Arc;
 
-/// Must use normal (lazy) compression, as faster ones (greedy) are not representative
+pub struct Candidate {
-const STD_COMPRESSION: u8 = 5;
+    pub image: PngData,
-const STD_FILTERS: [RowFilter; 2] = [RowFilter::None, RowFilter::MinSum];
+    pub filter: RowFilter,
-
+    pub is_reduction: bool,
-struct Candidate {
-    image: PngData,
-    filter: RowFilter,
     // first wins tie-breaker
     nth: usize,
 }
@@ -44,6 +42,8 @@ impl Candidate {
 /// Collect image versions and pick one that compresses best
 pub(crate) struct Evaluator {
     deadline: Arc<Deadline>,
+    filters: IndexSet<RowFilter>,
+    compression: u8,
     nth: AtomicUsize,
     best_candidate_size: Arc<AtomicMin>,
     /// images are sent to the caller thread for evaluation
@@ -55,11 +55,13 @@ pub(crate) struct Evaluator {
 }
 
 impl Evaluator {
-    pub fn new(deadline: Arc<Deadline>) -> Self {
+    pub fn new(deadline: Arc<Deadline>, filters: IndexSet<RowFilter>, compression: u8) -> Self {
         #[cfg(feature = "parallel")]
         let eval_channel = unbounded();
         Self {
             deadline,
+            filters,
+            compression,
             best_candidate_size: Arc::new(AtomicMin::new(None)),
             nth: AtomicUsize::new(0),
             #[cfg(feature = "parallel")]
@@ -72,26 +74,27 @@ impl Evaluator {
     /// Wait for all evaluations to finish and return smallest reduction
     /// Or `None` if all reductions were worse than baseline.
     #[cfg(feature = "parallel")]
-    fn get_best_candidate(self) -> Option<Candidate> {
+    pub fn get_best_candidate(self) -> Option<Candidate> {
         let (eval_send, eval_recv) = self.eval_channel;
         drop(eval_send); // disconnect the sender, breaking the loop in the thread
         eval_recv.into_iter().min_by_key(Candidate::cmp_key)
     }
 
     #[cfg(not(feature = "parallel"))]
-    fn get_best_candidate(self) -> Option<Candidate> {
+    pub fn get_best_candidate(self) -> Option<Candidate> {
         self.eval_best_candidate.into_inner()
     }
 
-    pub fn get_result(self) -> Option<PngData> {
-        self.get_best_candidate().map(|candidate| candidate.image)
-    }
-
     /// Set baseline image. It will be used only to measure minimum compression level required
     pub fn set_baseline(&self, image: Arc<PngImage>) {
         self.try_image_inner(image, false)
     }
 
+    /// Set best size, if known in advance
+    pub fn set_best_size(&self, size: usize) {
+        self.best_candidate_size.set_min(size);
+    }
+
     /// Check if the image is smaller than others
     pub fn try_image(&self, image: Arc<PngImage>) {
         self.try_image_inner(image, true)
@@ -101,13 +104,15 @@ impl Evaluator {
         let nth = self.nth.fetch_add(1, SeqCst);
         // These clones are only cheap refcounts
         let deadline = self.deadline.clone();
+        let filters = self.filters.clone();
+        let compression = self.compression;
         let best_candidate_size = self.best_candidate_size.clone();
         // sends it off asynchronously for compression,
         // but results will be collected via the message queue
         #[cfg(feature = "parallel")]
         let eval_send = self.eval_channel.0.clone();
         rayon::spawn(move || {
-            let filters_iter = STD_FILTERS.par_iter().with_max_len(1);
+            let filters_iter = filters.par_iter().with_max_len(1);
 
             // Updating of best result inside the parallel loop would require locks,
             // which are dangerous to do in side Rayon's loop.
@@ -119,21 +124,17 @@ impl Evaluator {
                 }
                 if let Ok(idat_data) = deflate::deflate(
                     &image.filter_image(filter),
-                    STD_COMPRESSION,
+                    compression,
                     &best_candidate_size,
                 ) {
                     best_candidate_size.set_min(idat_data.len());
-                    // ignore baseline images after this point
-                    if !is_reduction {
-                        return;
-                    }
-                    // the rest is shipped to the evavluation/collection thread
                     let new = Candidate {
                         image: PngData {
                             idat_data,
                             raw: Arc::clone(&image),
                         },
                         filter,
+                        is_reduction,
                         nth,
                     };
 

src/lib.rs

@@ -47,7 +47,7 @@ pub use crate::deflate::Deflaters;
 pub use crate::error::PngError;
 pub use crate::filters::RowFilter;
 pub use crate::headers::Headers;
-pub use indexmap::{IndexMap, IndexSet};
+pub use indexmap::{indexset, IndexMap, IndexSet};
 
 mod atomicmin;
 mod colors;
@@ -192,12 +192,12 @@ pub struct Options {
     ///
     /// Default: `Libdeflater`
     pub deflate: Deflaters,
-    /// Whether to use heuristics to pick the best filter and compression
+    /// Whether to use fast evaluation to pick the best filter
     ///
     /// Intended for use with `-o 1` from the CLI interface
     ///
     /// Default: `false`
-    pub use_heuristics: bool,
+    pub fast_evaluation: bool,
 
     /// Maximum amount of time to spend on optimizations.
     /// Further potential optimizations are skipped if the timeout is exceeded.
@@ -232,20 +232,18 @@ impl Options {
         self.idat_recoding = false;
         self.filter.clear();
         if let Deflaters::Libdeflater { compression } = &mut self.deflate {
-            compression.clear();
+            *compression = 5;
-            compression.insert(5);
         }
-        self.use_heuristics = true;
+        self.fast_evaluation = true;
         self
     }
 
     fn apply_preset_1(mut self) -> Self {
         self.filter.clear();
         if let Deflaters::Libdeflater { compression } = &mut self.deflate {
-            compression.clear();
+            *compression = 10;
-            compression.insert(10);
         }
-        self.use_heuristics = true;
+        self.fast_evaluation = true;
         self
     }
 
@@ -263,45 +261,23 @@ impl Options {
 
     fn apply_preset_4(mut self) -> Self {
         if let Deflaters::Libdeflater { compression } = &mut self.deflate {
-            compression.clear();
+            *compression = 12;
-            compression.insert(12);
         }
         self.apply_preset_3()
     }
 
-    fn apply_preset_5(mut self) -> Self {
+    fn apply_preset_5(self) -> Self {
-        if let Deflaters::Libdeflater { compression } = &mut self.deflate {
+        self.apply_preset_4()
-            compression.clear();
-            for i in 9..=12 {
-                compression.insert(i);
-            }
-        }
-        self.apply_preset_3()
     }
 
-    fn apply_preset_6(mut self) -> Self {
+    fn apply_preset_6(self) -> Self {
-        if let Deflaters::Libdeflater { compression } = &mut self.deflate {
+        self.apply_preset_4()
-            compression.clear();
-            for i in 1..=12 {
-                compression.insert(i);
-            }
-        }
-        self.apply_preset_3()
     }
 }
 
 impl Default for Options {
     fn default() -> Options {
         // Default settings based on -o 2 from the CLI interface
-        let mut filter = IndexSet::new();
-        filter.insert(RowFilter::None);
-        filter.insert(RowFilter::MinSum);
-        let mut compression = IndexSet::new();
-        compression.insert(11);
-        // We always need NoOp to be present
-        let mut alphas = IndexSet::new();
-        alphas.insert(AlphaOptim::NoOp);
-
         Options {
             backup: false,
             check: false,
@@ -309,17 +285,17 @@ impl Default for Options {
             fix_errors: false,
             force: false,
             preserve_attrs: false,
-            filter,
+            filter: indexset! {RowFilter::None, RowFilter::MinSum},
             interlace: None,
-            alphas,
+            alphas: IndexSet::new(),
             bit_depth_reduction: true,
             color_type_reduction: true,
             palette_reduction: true,
             grayscale_reduction: true,
             idat_recoding: true,
             strip: Headers::None,
-            deflate: Deflaters::Libdeflater { compression },
+            deflate: Deflaters::Libdeflater { compression: 11 },
-            use_heuristics: false,
+            fast_evaluation: false,
             timeout: None,
         }
     }
@@ -470,6 +446,7 @@ struct TrialOptions {
     pub filter: RowFilter,
     pub compression: u8,
 }
+type TrialWithData = (TrialOptions, Vec<u8>);
 
 /// Perform optimization on the input PNG object using the options provided
 fn optimize_png(
@@ -478,8 +455,6 @@ fn optimize_png(
     opts: &Options,
     deadline: Arc<Deadline>,
 ) -> PngResult<Vec<u8>> {
-    type TrialWithData = (TrialOptions, Vec<u8>);
-
     let original_png = png.clone();
 
     // Print png info
@@ -506,135 +481,106 @@ fn optimize_png(
     info!("    IDAT size = {} bytes", idat_original_size);
     info!("    File size = {} bytes", file_original_size);
 
-    let mut filter = opts.filter.clone();
+    // Must use normal (lazy) compression, as faster ones (greedy) are not representative
-
+    let eval_compression = 5;
-    if opts.use_heuristics {
+    let eval_filters = indexset! {RowFilter::None, RowFilter::MinSum};
-        // Heuristically determine which set of options to use
-        let use_filter = if png.raw.ihdr.bit_depth.as_u8() >= 8
-            && png.raw.ihdr.color_type != colors::ColorType::Indexed
-        {
-            RowFilter::MinSum
-        } else {
-            RowFilter::None
-        };
-        if filter.is_empty() {
-            filter.insert(use_filter);
-        }
-    }
-
     // This will collect all versions of images and pick one that compresses best
-    let eval = Evaluator::new(deadline.clone());
+    let eval = Evaluator::new(deadline.clone(), eval_filters.clone(), eval_compression);
-    // Usually we want transformations that are smaller than the unmodified original,
-    // but if we're interlacing, we have to accept a possible file size increase.
-    if opts.interlace.is_none() {
-        eval.set_baseline(png.raw.clone());
-    }
     perform_reductions(png.raw.clone(), opts, &deadline, &eval);
-    let reduction_occurred = if let Some(result) = eval.get_result() {
+    let (reduction_occurred, mut eval_filter) = if let Some(result) = eval.get_best_candidate() {
-        *png = result;
+        *png = result.image;
-        true
+        (result.is_reduction, Some(result.filter))
     } else {
-        false
+        (false, None)
     };
 
     if opts.idat_recoding || reduction_occurred {
-        // Go through selected permutations and determine the best
+        let mut filters = opts.filter.clone();
-        let combinations = if let Deflaters::Libdeflater { compression } = &opts.deflate {
+        let fast_eval = opts.fast_evaluation && (filters.len() > 1 || eval_filter.is_some());
-            filter.len() * compression.len()
+        let best: Option<TrialWithData> = if fast_eval {
-        } else {
+            // Perform a fast evaluation of selected filters followed by a single main compression trial
-            filter.len()
+            if eval_filter.is_some() {
-        };
+                // Some filters have already been evaluated, we don't need to try them again
-        let mut results: Vec<TrialOptions> = Vec::with_capacity(combinations);
+                filters = filters.difference(&eval_filters).cloned().collect();
+            }
 
-        for f in &filter {
+            if !filters.is_empty() {
-            if let Deflaters::Libdeflater { compression } = &opts.deflate {
+                debug!("Evaluating: {} filters", filters.len());
-                for zc in compression {
+                let eval = Evaluator::new(deadline, filters, eval_compression);
-                    results.push(TrialOptions {
+                if eval_filter.is_some() {
-                        filter: *f,
+                    eval.set_best_size(png.idat_data.len());
-                        compression: *zc,
+                }
-                    });
+                eval.try_image(png.raw.clone());
-                    if deadline.passed() {
+                if let Some(result) = eval.get_best_candidate() {
-                        break;
+                    *png = result.image;
-                    }
+                    eval_filter = Some(result.filter);
+                }
+            }
+
+            let trial = TrialOptions {
+                filter: eval_filter.unwrap(),
+                compression: match opts.deflate {
+                    Deflaters::Libdeflater { compression } => compression,
+                    _ => 0,
+                },
+            };
+            if trial.compression <= eval_compression {
+                // No further compression required
+                if png.idat_data.len() < idat_original_size || opts.force {
+                    Some((trial, png.idat_data.clone()))
+                } else {
+                    None
                 }
             } else {
-                // Zopfli has no additional options.
+                info!("Trying: {}", trial.filter);
+                let original_len = idat_original_size;
+                let best_size = AtomicMin::new(if opts.force { None } else { Some(original_len) });
+                perform_trial(&png.raw, opts, trial, &best_size)
+            }
+        } else {
+            // Perform full compression trials of selected filters and determine the best
+            if filters.is_empty() {
+                // Heuristically determine which filter to use
+                if png.raw.ihdr.bit_depth.as_u8() >= 8
+                    && png.raw.ihdr.color_type != colors::ColorType::Indexed
+                {
+                    filters.insert(RowFilter::MinSum);
+                } else {
+                    filters.insert(RowFilter::None);
+                }
+            }
+
+            let mut results: Vec<TrialOptions> = Vec::with_capacity(filters.len());
+
+            for f in &filters {
                 results.push(TrialOptions {
                     filter: *f,
-                    compression: 0,
+                    compression: match opts.deflate {
+                        Deflaters::Libdeflater { compression } => compression,
+                        _ => 0,
+                    },
                 });
             }
 
-            if deadline.passed() {
+            info!("Trying: {} filters", results.len());
-                break;
-            }
-        }
 
-        info!("Trying: {} combinations", results.len());
+            let original_len = idat_original_size;
-
+            let best_size = AtomicMin::new(if opts.force { None } else { Some(original_len) });
-        let filters: IndexMap<RowFilter, Vec<u8>> = filter
+            let results_iter = results.into_par_iter().with_max_len(1);
-            .par_iter()
+            let best = results_iter.filter_map(|trial| {
-            .with_max_len(1)
+                if deadline.passed() {
-            .map(|f| {
-                let png = png.clone();
-                (*f, png.raw.filter_image(*f))
-            })
-            .collect();
-
-        let original_len = original_png.idat_data.len();
-        let added_interlacing = opts.interlace == Some(1) && original_png.raw.ihdr.interlaced == 0;
-
-        let best_size = AtomicMin::new(if opts.force { None } else { Some(original_len) });
-        let results_iter = results.into_par_iter().with_max_len(1);
-        let best = results_iter.filter_map(|trial| {
-            if deadline.passed() {
-                return None;
-            }
-            let filtered = &filters[&trial.filter];
-            let new_idat = match opts.deflate {
-                Deflaters::Libdeflater { .. } => {
-                    deflate::deflate(filtered, trial.compression, &best_size)
-                }
-                #[cfg(feature = "zopfli")]
-                Deflaters::Zopfli { iterations } => deflate::zopfli_deflate(filtered, iterations),
-            };
-
-            let new_idat = match new_idat {
-                Ok(n) => n,
-                Err(PngError::DeflatedDataTooLong(max)) => {
-                    debug!(
-                        "    zc = {}  f = {} >{} bytes",
-                        trial.compression, trial.filter, max,
-                    );
                     return None;
                 }
-                Err(_) => return None,
+                perform_trial(&png.raw, opts, trial, &best_size)
-            };
+            });
-
+            best.reduce_with(|i, j| {
-            // update best size across all threads
+                if i.1.len() < j.1.len() || (i.1.len() == j.1.len() && i.0 < j.0) {
-            let new_size = new_idat.len();
+                    i
-            best_size.set_min(new_size);
+                } else {
-
+                    j
-            debug!(
+                }
-                "    zc = {}  f = {}  {} bytes",
+            })
-                trial.compression,
+        };
-                trial.filter,
-                new_idat.len()
-            );
-
-            if new_size < original_len || added_interlacing || opts.force {
-                Some((trial, new_idat))
-            } else {
-                None
-            }
-        });
-        let best: Option<TrialWithData> = 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((opts, idat_data)) = best {
             png.idat_data = idat_data;
@@ -645,9 +591,11 @@ fn optimize_png(
                 opts.filter,
                 png.idat_data.len()
             );
-        } else if reduction_occurred {
+        } else if eval_filter.is_some() {
             *png = original_png;
         }
+    } else if png.idat_data.len() >= idat_original_size {
+        *png = original_png;
     }
 
     perform_strip(png, opts);
@@ -713,11 +661,17 @@ fn perform_reductions(
     deadline: &Deadline,
     eval: &Evaluator,
 ) {
+    // The eval baseline will be set from the original png only if we attempt any reductions
+    let mut baseline = Some(png.clone());
+    let mut reduction_occurred = false;
+
     // must be done first to evaluate rest with the correct interlacing
     if let Some(interlacing) = opts.interlace {
         if let Some(reduced) = png.change_interlacing(interlacing) {
             png = Arc::new(reduced);
             eval.try_image(png.clone());
+            // If we're interlacing, we have to accept a possible file size increase
+            baseline = None;
         }
         if deadline.passed() {
             return;
@@ -729,6 +683,7 @@ fn perform_reductions(
             png = Arc::new(reduced);
             eval.try_image(png.clone());
             report_reduction(&png);
+            reduction_occurred = true;
         }
         if deadline.passed() {
             return;
@@ -750,6 +705,7 @@ fn perform_reductions(
                 }
             }
             report_reduction(&png);
+            reduction_occurred = true;
         }
         if deadline.passed() {
             return;
@@ -761,13 +717,62 @@ fn perform_reductions(
             png = Arc::new(reduced);
             eval.try_image(png.clone());
             report_reduction(&png);
+            reduction_occurred = true;
         }
         if deadline.passed() {
             return;
         }
     }
 
-    try_alpha_reductions(png, &opts.alphas, eval);
+    if try_alpha_reductions(png, &opts.alphas, eval) {
+        reduction_occurred = true;
+    }
+
+    if let Some(baseline) = baseline {
+        if reduction_occurred {
+            eval.set_baseline(baseline);
+        }
+    }
+}
+
+/// Execute a compression trial
+fn perform_trial(
+    png: &PngImage,
+    opts: &Options,
+    trial: TrialOptions,
+    best_size: &AtomicMin,
+) -> Option<TrialWithData> {
+    let filtered = &png.filter_image(trial.filter);
+    let new_idat = match opts.deflate {
+        Deflaters::Libdeflater { .. } => deflate::deflate(filtered, trial.compression, best_size),
+        #[cfg(feature = "zopfli")]
+        Deflaters::Zopfli { iterations } => deflate::zopfli_deflate(filtered, iterations),
+    };
+
+    // update best size or convert to error if not smaller
+    let new_idat = match new_idat {
+        Ok(n) if !best_size.set_min(n.len()) => Err(PngError::DeflatedDataTooLong(n.len())),
+        _ => new_idat,
+    };
+
+    match new_idat {
+        Ok(n) => {
+            let bytes = n.len();
+            debug!(
+                "    zc = {}  f = {}  {} bytes",
+                trial.compression, trial.filter, bytes
+            );
+            Some((trial, n))
+        }
+        Err(PngError::DeflatedDataTooLong(bytes)) => {
+            debug!(
+                "    zc = {}  f = {} >{} bytes",
+                trial.compression, trial.filter, bytes,
+            );
+            None
+        }
+        Err(_) => None,
+    }
 }
 
 #[derive(Debug)]

src/main.rs

@@ -184,14 +184,11 @@ fn main() {
         )
         .arg(
             Arg::new("compression")
-                .help("zlib compression levels (1-12) - Default: 12")
+                .help("zlib compression level (1-12) - Default: 11")
                 .long("zc")
                 .takes_value(true)
-                .value_name("levels")
+                .value_name("level")
-                .validator(|x| match parse_numeric_range_opts(x, 1, 12) {
+                .value_parser(1..=12)
-                    Ok(_) => Ok(()),
-                    Err(_) => Err("Invalid option for compression".to_owned()),
-                })
                 .conflicts_with("zopfli"),
         )
         .arg(
@@ -541,8 +538,8 @@ fn parse_opts_into_struct(
             opts.deflate = Deflaters::Zopfli { iterations };
         }
     } else if let Deflaters::Libdeflater { compression } = &mut opts.deflate {
-        if let Some(x) = matches.value_of("compression") {
+        if let Some(x) = matches.get_one::<i64>("compression") {
-            *compression = parse_numeric_range_opts(x, 1, 12).unwrap();
+            *compression = *x as u8;
         }
     }
 

src/reduction/alpha.rs

@@ -14,16 +14,18 @@ pub(crate) fn try_alpha_reductions(
     png: Arc<PngImage>,
     alphas: &IndexSet<AlphaOptim>,
     eval: &Evaluator,
-) {
+) -> bool {
-    if alphas.is_empty() {
+    match png.ihdr.color_type {
-        return;
+        ColorType::RGBA | ColorType::GrayscaleAlpha if !alphas.is_empty() => {
+            alphas
+                .par_iter()
+                .with_max_len(1)
+                .filter_map(|&alpha| filtered_alpha_channel(&png, alpha))
+                .for_each(|image| eval.try_image(Arc::new(image)));
+            true
+        }
+        _ => false,
     }
-
-    alphas
-        .par_iter()
-        .with_max_len(1)
-        .filter_map(|&alpha| filtered_alpha_channel(&png, alpha))
-        .for_each(|image| eval.try_image(Arc::new(image)));
 }
 
 pub fn filtered_alpha_channel(png: &PngImage, optim: AlphaOptim) -> Option<PngImage> {