Async reduction evaluator

src/evaluate.rs

@@ -0,0 +1,96 @@
+//! Check if a reduction makes file smaller, and keep best reductions.
+//! Works asynchronously when possible
+
+use atomicmin::AtomicMin;
+use deflate;
+use png::PngData;
+use png::PngImage;
+use png::STD_COMPRESSION;
+use png::STD_FILTERS;
+use png::STD_STRATEGY;
+use png::STD_WINDOW;
+#[cfg(feature = "parallel")]
+use rayon::prelude::*;
+use std::sync::mpsc::*;
+use std::sync::Arc;
+use std::sync::Mutex;
+use std::thread;
+
+/// Collect image versions and pick one that compresses best
+pub struct Evaluator {
+    /// images are sent to the thread for evaluation
+    eval_send: Option<SyncSender<(Arc<PngImage>, bool)>>,
+    // the thread helps evaluate images asynchronously
+    eval_thread: thread::JoinHandle<Option<PngData>>,
+}
+
+impl Evaluator {
+    pub fn new() -> Self {
+        // queue size ensures we're not using too much memory for pending reductions
+        let (tx, rx) = sync_channel(4);
+        Self {
+            eval_send: Some(tx),
+            eval_thread: thread::spawn(move || Self::evaluate_images(rx)),
+        }
+    }
+
+    /// Wait for all evaluations to finish and return smallest reduction
+    /// Or `None` if all reductions were worse than baseline.
+    pub fn get_result(mut self) -> Option<PngData> {
+        let _ = self.eval_send.take(); // disconnect the sender, breaking the loop in the thread
+        self.eval_thread.join().expect("eval thread")
+    }
+
+    /// 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)
+    }
+
+    /// Check if the image is smaller than others
+    pub fn try_image(&self, image: Arc<PngImage>) {
+        self.try_image_inner(image, true)
+    }
+
+    fn try_image_inner(&self, image: Arc<PngImage>, is_reduction: bool) {
+        self.eval_send.as_ref().expect("not finished yet").send((image, is_reduction)).expect("send")
+    }
+
+    /// Main loop of evaluation thread
+    fn evaluate_images(from_channel: Receiver<(Arc<PngImage>, bool)>) -> Option<PngData> {
+        let best_candidate_size = AtomicMin::new(None);
+        let best_result = Mutex::new(None);
+        // ends when sender is dropped
+        for (image, is_reduction) in from_channel.iter() {
+            #[cfg(feature = "parallel")]
+            let filters_iter = STD_FILTERS.par_iter().with_max_len(1);
+            #[cfg(not(feature = "parallel"))]
+            let filters_iter = STD_FILTERS.iter();
+
+            filters_iter.for_each(|&f| {
+                if let Ok(idat_data) = deflate::deflate(
+                    &image.filter_image(f),
+                    STD_COMPRESSION,
+                    STD_STRATEGY,
+                    STD_WINDOW,
+                    &best_candidate_size,
+                ) {
+                    let mut res = best_result.lock().unwrap();
+                    if best_candidate_size.get().map_or(true, |len| len >= idat_data.len()) {
+                        best_candidate_size.set_min(idat_data.len());
+                        *res = if is_reduction {
+                            Some((PngData {
+                                idat_data,
+                                raw: Arc::clone(&image),
+                            }, f))
+                        } else {
+                            None
+                        };
+                    }
+                }
+            });
+        }
+        best_result.into_inner().expect("filters should be done")
+            .map(|(img, _)| img)
+    }
+}
+

src/lib.rs

@@ -16,6 +16,7 @@ use reduction::*;
 use atomicmin::AtomicMin;
 use crc::crc32;
 use deflate::inflate;
+use evaluate::Evaluator;
 use image::{DynamicImage, GenericImageView, ImageFormat, Pixel};
 use png::PngImage;
 use png::PngData;
@@ -23,12 +24,12 @@ use png::PngData;
 use rayon::prelude::*;
 use std::collections::{HashMap, HashSet};
 use std::fmt;
-use std::borrow::Cow;
 use std::fs::{copy, File};
 use std::io::{stdin, stdout, BufWriter, Read, Write};
 use std::path::{Path, PathBuf};
 use std::sync::atomic::{AtomicBool, Ordering};
 use std::time::{Duration, Instant};
+use std::sync::Arc;
 
 pub use colors::AlphaOptim;
 pub use deflate::Deflaters;
@@ -39,6 +40,7 @@ mod atomicmin;
 mod colors;
 mod deflate;
 mod error;
+mod evaluate;
 mod filters;
 mod headers;
 mod interlace;
@@ -522,9 +524,16 @@ fn optimize_png(png: &mut PngData, original_data: &[u8], opts: &Options) -> PngR
         }
     }
 
-    let reduction_occurred = if let Some(reduced) = perform_reductions(&png.raw, opts, &deadline) {
+    // This will collect all versions of images and pick one that compresses best
-        png.raw = reduced;
+    let eval = Evaluator::new();
-        png.idat_data.clear(); // this field is out of date and needs to be replaced
+    // 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() {
+        *png = result;
         true
     } else {
         false
@@ -731,66 +740,59 @@ fn optimize_png(png: &mut PngData, original_data: &[u8], opts: &Options) -> PngR
     Err(PngError::new("The resulting image is corrupted"))
 }
 
-fn if_owned(cow: Cow<PngImage>) -> Option<PngImage> {
+fn perform_reductions(mut png: Arc<PngImage>, opts: &Options, deadline: &Deadline, eval: &Evaluator) {
-    match cow {
-        Cow::Owned(png) => Some(png),
-        _ => None,
-    }
-}
 
-fn perform_reductions(png: &PngImage, opts: &Options, deadline: &Deadline) -> Option<PngImage> {
+    // must be done first to evaluate rest with the correct interlacing
-    let mut reduced = Cow::Borrowed(png);
+    if let Some(interlacing) = opts.interlace {
+        if let Some(reduced) = png.change_interlacing(interlacing) {
+            png = Arc::new(reduced);
+            eval.try_image(png.clone());
+        }
+        if deadline.passed() {
+            return;
+        }
+    }
 
     if opts.palette_reduction {
-        if let Some(r) = reduced_palette(&reduced) {
+        if let Some(reduced) = reduced_palette(&png) {
-            reduced = Cow::Owned(r);
+            png = Arc::new(reduced);
+            eval.try_image(png.clone());
             if opts.verbosity == Some(1) {
-                report_reduction(&reduced);
+                report_reduction(&png);
             }
         }
-    }
-
         if deadline.passed() {
-        return if_owned(reduced);
+            return;
+        }
     }
 
     if opts.bit_depth_reduction {
-        if let Some(r) = reduce_bit_depth(&reduced) {
+        if let Some(reduced) = reduce_bit_depth(&png) {
-            reduced = Cow::Owned(r);
+            png = Arc::new(reduced);
+            eval.try_image(png.clone());
             if opts.verbosity == Some(1) {
-                report_reduction(&reduced);
+                report_reduction(&png);
             }
         }
-    }
-
         if deadline.passed() {
-        return if_owned(reduced);
+            return;
+        }
     }
 
     if opts.color_type_reduction {
-        if let Some(r) = reduce_color_type(&reduced) {
+        if let Some(reduced) = reduce_color_type(&png) {
-            reduced = Cow::Owned(r);
+            png = Arc::new(reduced);
+            eval.try_image(png.clone());
             if opts.verbosity == Some(1) {
-                report_reduction(&reduced);
+                report_reduction(&png);
             }
         }
-    }
-
-    if let Some(interlacing) = opts.interlace {
-        if let Some(r) = reduced.change_interlacing(interlacing) {
-            reduced = Cow::Owned(r);
-        }
-    }
-
         if deadline.passed() {
-        return if_owned(reduced);
+            return;
+        }
     }
 
-    if let Some(r) = try_alpha_reduction(&reduced, &opts.alphas) {
+    try_alpha_reductions(png, &opts.alphas, eval);
-        reduced = Cow::Owned(r);
-    }
-
-    if_owned(reduced)
 }
 
 /// Keep track of processing timeout
@@ -846,29 +848,30 @@ fn report_reduction(png: &PngImage) {
 
 /// Strip headers from the `PngData` object, as requested by the passed `Options`
 fn perform_strip(png: &mut PngData, opts: &Options) {
+    let raw = Arc::make_mut(&mut png.raw);
     match opts.strip {
         // Strip headers
         Headers::None => (),
         Headers::Keep(ref hdrs) => {
-            png.raw.aux_headers.retain(|chunk, _| {
+            raw.aux_headers.retain(|chunk, _| {
                 std::str::from_utf8(chunk)
                     .ok()
                     .map_or(false, |name| hdrs.contains(name))
             });
         }
         Headers::Strip(ref hdrs) => for hdr in hdrs {
-            png.raw.aux_headers.remove(hdr.as_bytes());
+            raw.aux_headers.remove(hdr.as_bytes());
         },
         Headers::Safe => {
             const PRESERVED_HEADERS: [[u8; 4]; 9] = [
                 *b"cHRM", *b"gAMA", *b"iCCP", *b"sBIT", *b"sRGB", *b"bKGD", *b"hIST", *b"pHYs",
                 *b"sPLT",
             ];
-            png.raw.aux_headers
+            raw.aux_headers
                 .retain(|hdr, _| PRESERVED_HEADERS.contains(hdr));
         }
         Headers::All => {
-            png.raw.aux_headers = HashMap::new();
+            raw.aux_headers = HashMap::new();
         }
     }
 
@@ -881,18 +884,18 @@ fn perform_strip(png: &mut PngData, opts: &Options) {
     };
 
     if may_replace_iccp {
-        if png.raw.aux_headers.get(b"sRGB").is_some() {
+        if raw.aux_headers.get(b"sRGB").is_some() {
             // Files aren't supposed to have both chunks, so we chose to honor sRGB
-            png.raw.aux_headers.remove(b"iCCP");
+            raw.aux_headers.remove(b"iCCP");
-        } else if let Some(intent) = png.raw
+        } else if let Some(intent) = raw
             .aux_headers
             .get(b"iCCP")
             .and_then(|iccp| srgb_rendering_intent(iccp))
         {
             // sRGB-like profile can be safely replaced with
             // an sRGB chunk with the same rendering intent
-            png.raw.aux_headers.remove(b"iCCP");
+            raw.aux_headers.remove(b"iCCP");
-            png.raw.aux_headers.insert(*b"sRGB", vec![intent]);
+            raw.aux_headers.insert(*b"sRGB", vec![intent]);
         }
     }
 }

src/png/mod.rs

@@ -13,11 +13,13 @@ use std::fs::File;
 use std::io::{Read, Seek, SeekFrom};
 use std::iter::Iterator;
 use std::path::Path;
+use std::sync::Arc;
 
-pub(crate) const STD_COMPRESSION: u8 = 8;
+pub(crate) const STD_COMPRESSION: u8 = 6;
 /// Must use normal compression, as faster ones (Huffman/RLE-only) are not representative
 pub(crate) const STD_STRATEGY: u8 = 0;
-pub(crate) const STD_WINDOW: u8 = 15;
+/// OK to use a bit smalller window for evaluation
+pub(crate) const STD_WINDOW: u8 = 13;
 pub(crate) const STD_FILTERS: [u8; 2] = [0, 5];
 
 pub(crate) mod scan_lines;
@@ -42,7 +44,8 @@ pub struct PngImage {
 /// Contains all data relevant to a PNG image
 #[derive(Debug, Clone)]
 pub struct PngData {
-    pub raw: PngImage,
+    /// Uncompressed image data
+    pub raw: Arc<PngImage>,
     /// The filtered and compressed data of the IDAT chunk
     pub idat_data: Vec<u8>,
 }
@@ -122,19 +125,19 @@ impl PngData {
             aux_headers.remove(b"tRNS"),
         )?;
 
-        let mut png_data = Self {
+        let mut raw = PngImage {
-            idat_data: idat_headers,
-            raw: PngImage {
             ihdr: ihdr_header,
             data: raw_data,
             palette,
             transparency_pixel,
             aux_headers,
-            }
         };
-        png_data.raw.data = png_data.raw.unfilter_image();
+        raw.data = raw.unfilter_image();
         // Return the PngData
-        Ok(png_data)
+        Ok(Self {
+            idat_data: idat_headers,
+            raw: Arc::new(raw),
+        })
     }
 
     /// Handle transparency header
@@ -162,7 +165,6 @@ impl PngData {
         }
     }
 
-
     /// Format the `PngData` struct into a valid PNG bytestream
     pub fn output(&self) -> Vec<u8> {
         // PNG header

src/reduction/alpha.rs

@@ -1,56 +1,25 @@
-use png::STD_STRATEGY;
+use evaluate::Evaluator;
-use png::STD_WINDOW;
 use itertools::flatten;
 use png::scan_lines::ScanLine;
 use std::collections::HashSet;
-use png::STD_COMPRESSION;
+use std::sync::Arc;
-use png::STD_FILTERS;
 use colors::AlphaOptim;
 use headers::IhdrData;
 use png::PngImage;
 use colors::ColorType;
-use atomicmin::AtomicMin;
 #[cfg(feature = "parallel")]
 use rayon::prelude::*;
-use deflate;
 
-pub fn try_alpha_reduction(png: &PngImage, alphas: &HashSet<AlphaOptim>) -> Option<PngImage> {
+pub fn try_alpha_reductions(png: Arc<PngImage>, alphas: &HashSet<AlphaOptim>, eval: &Evaluator) {
     assert!(!alphas.is_empty());
     let alphas = alphas.iter().collect::<Vec<_>>();
-    let best_size = AtomicMin::new(None);
     #[cfg(feature = "parallel")]
     let alphas_iter = alphas.par_iter().with_max_len(1);
     #[cfg(not(feature = "parallel"))]
     let alphas_iter = alphas.iter();
-    let best = alphas_iter
+    alphas_iter
-        .filter_map(|&alpha| {
+        .filter_map(|&alpha| filtered_alpha_channel(&png, *alpha))
-            let image = match filtered_alpha_channel(png, *alpha) {
+        .for_each(|image| eval.try_image(Arc::new(image)));
-                Some(image) => image,
-                None => return None,
-            };
-            #[cfg(feature = "parallel")]
-            let filters_iter = STD_FILTERS.par_iter().with_max_len(1);
-            #[cfg(not(feature = "parallel"))]
-            let filters_iter = STD_FILTERS.iter();
-            filters_iter
-                .filter_map(|f| {
-                    deflate::deflate(
-                        &image.filter_image(*f),
-                        STD_COMPRESSION,
-                        STD_STRATEGY,
-                        STD_WINDOW,
-                        &best_size,
-                    ).ok()
-                    .as_ref()
-                    .map(|l| {
-                        best_size.set_min(l.len());
-                        l.len()
-                    })
-                }).min()
-                .map(|size| (size, image))
-        }).min_by_key(|&(size, _)| size);
-
-    best.map(|(_, image)| image)
 }
 
 pub fn filtered_alpha_channel(png: &PngImage, optim: AlphaOptim) -> Option<PngImage> {

src/reduction/mod.rs

@@ -14,7 +14,7 @@ pub mod color;
 use color::*;
 
 pub use bit_depth::reduce_bit_depth;
-pub use alpha::try_alpha_reduction;
+pub use alpha::try_alpha_reductions;
 
 /// Attempt to reduce the number of colors in the palette
 /// Returns `None` if palette hasn't changed

tests/reduction.rs

@@ -718,7 +718,7 @@ fn palette_should_be_reduced_with_dupes() {
 
     assert_eq!(png.raw.ihdr.color_type, ColorType::Indexed);
     assert_eq!(png.raw.ihdr.bit_depth, BitDepth::Eight);
-    assert_eq!(png.raw.palette.unwrap().len(), 43);
+    assert_eq!(png.raw.palette.as_ref().unwrap().len(), 43);
 
     match oxipng::optimize(&InFile::Path(input), &output, &opts) {
         Ok(_) => (),
@@ -737,7 +737,7 @@ fn palette_should_be_reduced_with_dupes() {
 
     assert_eq!(png.raw.ihdr.color_type, ColorType::Indexed);
     assert_eq!(png.raw.ihdr.bit_depth, BitDepth::Eight);
-    assert_eq!(png.raw.palette.unwrap().len(), 35);
+    assert_eq!(png.raw.palette.as_ref().unwrap().len(), 35);
 
     remove_file(output).ok();
 }
@@ -751,7 +751,7 @@ fn palette_should_be_reduced_with_unused() {
 
     assert_eq!(png.raw.ihdr.color_type, ColorType::Indexed);
     assert_eq!(png.raw.ihdr.bit_depth, BitDepth::Eight);
-    assert_eq!(png.raw.palette.unwrap().len(), 35);
+    assert_eq!(png.raw.palette.as_ref().unwrap().len(), 35);
 
     match oxipng::optimize(&InFile::Path(input), &output, &opts) {
         Ok(_) => (),
@@ -770,7 +770,7 @@ fn palette_should_be_reduced_with_unused() {
 
     assert_eq!(png.raw.ihdr.color_type, ColorType::Indexed);
     assert_eq!(png.raw.ihdr.bit_depth, BitDepth::Eight);
-    assert_eq!(png.raw.palette.unwrap().len(), 33);
+    assert_eq!(png.raw.palette.as_ref().unwrap().len(), 33);
 
     remove_file(output).ok();
 }
@@ -784,7 +784,7 @@ fn palette_should_be_reduced_with_both() {
 
     assert_eq!(png.raw.ihdr.color_type, ColorType::Indexed);
     assert_eq!(png.raw.ihdr.bit_depth, BitDepth::Eight);
-    assert_eq!(png.raw.palette.unwrap().len(), 43);
+    assert_eq!(png.raw.palette.as_ref().unwrap().len(), 43);
 
     match oxipng::optimize(&InFile::Path(input), &output, &opts) {
         Ok(_) => (),
@@ -803,7 +803,7 @@ fn palette_should_be_reduced_with_both() {
 
     assert_eq!(png.raw.ihdr.color_type, ColorType::Indexed);
     assert_eq!(png.raw.ihdr.bit_depth, BitDepth::Eight);
-    assert_eq!(png.raw.palette.unwrap().len(), 33);
+    assert_eq!(png.raw.palette.as_ref().unwrap().len(), 33);
 
     remove_file(output).ok();
 }