From d286cf9eaf96912cfbb59b37fce53cc7e314bcef Mon Sep 17 00:00:00 2001
From: Felix Hanau <felix@cloudflare.com>
Date: Tue, 6 Aug 2019 14:58:54 +0100
Subject: [PATCH] Add strategy heuristics

---
 src/lib.rs | 44 +++++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 43 insertions(+), 1 deletion(-)

diff --git a/src/lib.rs b/src/lib.rs
index 84bd58dd..f096aa8b 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -540,6 +540,8 @@ fn optimize_png(
                 (*f, png.raw.filter_image(*f))
             })
             .collect();
+
+        //Use heuristics to choose the best filter strategy.
         if opts.use_heuristics && filter.len() > 1 {
             let best_size_heuristics = AtomicMin::new(None);
             let mut heuristic_results: Vec<TrialOptions> = Vec::with_capacity(filter.len());
@@ -575,7 +577,7 @@ fn optimize_png(
 
                 Some((trial.filter, new_size))
             });
-            let best_filter = best_filter.reduce_with(|i, j| if i.1 <= j.1 {i} else {j});
+            let best_filter = best_filter.reduce_with(|i, j| if i.1 < j.1 || (i.1 == j.1 && i.0 < j.0) {i} else {j});
 
             if let Some(result) = best_filter {
                 filter.clear();
@@ -584,6 +586,46 @@ fn optimize_png(
             }
         }
 
+        //Use heuristics to choose whether to try strategy 0 or 1 at full compression. Compression
+        //level 5 gives a good approximation.
+        if opts.use_heuristics && strategies.len() == 4 {
+            let best_size_heuristics = AtomicMin::new(None);
+            let heuristic_results = vec![0, 1];
+            let strategy_heuristics_iter = heuristic_results.into_par_iter().with_max_len(1);
+            let best_strategy = strategy_heuristics_iter.filter_map(|strategy| {
+                if deadline.passed() {
+                    return None;
+                }
+                let filtered = &filters[filter.iter().last().unwrap()];
+                let new_idat = deflate::deflate(
+                        filtered,
+                        5,
+                        strategy,
+                        opts.window,
+                        &best_size_heuristics,
+                        &deadline,
+                        );
+
+                let new_size = match new_idat {
+                    Ok(n) => n.len(),
+                    _ => return None,
+                };
+
+                // update best size across all threads
+                best_size_heuristics.set_min(new_size);
+
+                Some((strategy, new_size))
+            });
+            let best_strategy = best_strategy.reduce_with(|i, j| if i.1 < j.1 || (i.1 == j.1 && i.0 < j.0) {i} else {j});
+
+            if let Some(result) = best_strategy {
+                strategies.clear();
+                strategies.insert(result.0);
+                strategies.insert(2);
+                strategies.insert(3);
+            }
+        }
+
         // Go through selected permutations and determine the best
         let combinations = if opts.deflate == Deflaters::Zlib && !deadline.passed() {
             filter.len() * compression.len() * strategies.len()