From 798a1209269c1991da9cf86f504a5a0d92ca28e2 Mon Sep 17 00:00:00 2001
From: Andrew <andrew@digerati.co.nz>
Date: Sun, 15 Jan 2023 20:55:38 +1300
Subject: [PATCH] Include PLTE/tRNS size in evaluations

---
 src/evaluate.rs     |  4 ++--
 src/png/mod.rs      | 21 +++++++++++++++++++++
 tests/regression.rs |  8 ++++----
 3 files changed, 27 insertions(+), 6 deletions(-)

diff --git a/src/evaluate.rs b/src/evaluate.rs
index 23cc9832..ee92b17a 100644
--- a/src/evaluate.rs
+++ b/src/evaluate.rs
@@ -30,7 +30,7 @@ pub struct Candidate {
 impl Candidate {
     fn cmp_key(&self) -> impl Ord {
         (
-            self.image.idat_data.len(),
+            self.image.estimated_output_size(),
             self.image.raw.data.len(),
             self.image.raw.ihdr.bit_depth,
             self.filter,
@@ -134,7 +134,6 @@ impl Evaluator {
                 if let Ok(idat_data) =
                     deflate::deflate(&filtered, compression, &best_candidate_size)
                 {
-                    best_candidate_size.set_min(idat_data.len());
                     let new = Candidate {
                         image: PngData {
                             idat_data,
@@ -145,6 +144,7 @@ impl Evaluator {
                         is_reduction,
                         nth,
                     };
+                    best_candidate_size.set_min(new.image.estimated_output_size());
 
                     #[cfg(feature = "parallel")]
                     {
diff --git a/src/png/mod.rs b/src/png/mod.rs
index 32870581..dd37f69c 100644
--- a/src/png/mod.rs
+++ b/src/png/mod.rs
@@ -134,6 +134,27 @@ impl PngData {
         })
     }
 
+    /// Return an estimate of the output size
+    pub fn estimated_output_size(&self) -> usize {
+        // Add the size of the PLTE and tRNS chunks to the compressed idat size
+        // This can help with evaluation of very small data
+        let size = self.idat_data.len();
+        size + match &self.raw.ihdr.color_type {
+            ColorType::Indexed { palette } => {
+                let plte = 12 + palette.len() * 3;
+                let trns = palette.iter().filter(|p| p.a != 255).count();
+                if trns != 0 {
+                    plte + 12 + trns
+                } else {
+                    plte
+                }
+            }
+            ColorType::Grayscale { transparent_shade } if transparent_shade.is_some() => 12 + 2,
+            ColorType::RGB { transparent_color } if transparent_color.is_some() => 12 + 6,
+            _ => 0,
+        }
+    }
+
     /// Format the `PngData` struct into a valid PNG bytestream
     pub fn output(&self) -> Vec<u8> {
         // PNG header
diff --git a/tests/regression.rs b/tests/regression.rs
index 06833b25..a161a3bf 100644
--- a/tests/regression.rs
+++ b/tests/regression.rs
@@ -147,7 +147,7 @@ fn issue_52_02() {
         None,
         RGBA,
         BitDepth::Eight,
-        INDEXED,
+        RGBA,
         BitDepth::Eight,
     );
 }
@@ -159,7 +159,7 @@ fn issue_52_03() {
         None,
         RGBA,
         BitDepth::Eight,
-        INDEXED,
+        RGBA,
         BitDepth::Eight,
     );
 }
@@ -195,8 +195,8 @@ fn issue_52_06() {
         None,
         RGBA,
         BitDepth::Eight,
-        INDEXED,
-        BitDepth::Two,
+        RGBA,
+        BitDepth::Eight,
     );
 }