Palette reduction fixes (#466)

* Verify bKGD in tests

* Don't truncate bKGD on output

* Make sure correct bKGD is retained

* Update safe chunk list

* Perform strip first

* Check entire palette map

src/headers.rs

@@ -71,7 +71,7 @@ pub enum Headers {
     None,
     /// Remove specific chunks
     Strip(Vec<String>),
-    /// Headers that won't affect rendering (all but cHRM, gAMA, iCCP, sBIT, sRGB, bKGD, hIST, pHYs, sPLT)
+    /// Headers that won't affect rendering (all but cICP, iCCP, sBIT, sRGB, pHYs)
     Safe,
     /// Remove all non-critical chunks except these
     Keep(IndexSet<String>),

src/lib.rs

@@ -455,8 +455,6 @@ fn optimize_png(
     opts: &Options,
     deadline: Arc<Deadline>,
 ) -> PngResult<Vec<u8>> {
-    let original_png = png.clone();
-
     // Print png info
     let file_original_size = original_data.len();
     let idat_original_size = png.idat_data.len();
@@ -481,6 +479,10 @@ fn optimize_png(
     info!("    IDAT size = {} bytes", idat_original_size);
     info!("    File size = {} bytes", file_original_size);
 
+    // Do this first so that reductions can ignore certain chunks such as bKGD
+    perform_strip(png, opts);
+    let stripped_png = png.clone();
+
     // Must use normal (lazy) compression, as faster ones (greedy) are not representative
     let eval_compression = 5;
     let eval_filters = indexset! {RowFilter::None, RowFilter::MinSum};
@@ -592,14 +594,12 @@ fn optimize_png(
                 png.idat_data.len()
             );
         } else if eval_filter.is_some() {
-            *png = original_png;
+            *png = stripped_png;
         }
     } else if png.idat_data.len() >= idat_original_size {
-        *png = original_png;
+        *png = stripped_png;
     }
 
-    perform_strip(png, opts);
-
     let output = png.output();
 
     if idat_original_size >= png.idat_data.len() {
@@ -857,10 +857,8 @@ fn perform_strip(png: &mut PngData, opts: &Options) {
             }
         }
         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",
-            ];
+            const PRESERVED_HEADERS: [[u8; 4]; 5] =
+                [*b"cICP", *b"iCCP", *b"sBIT", *b"sRGB", *b"pHYs"];
             let keys: Vec<[u8; 4]> = raw.aux_headers.keys().cloned().collect();
             for hdr in &keys {
                 if !PRESERVED_HEADERS.contains(hdr) {

src/png/mod.rs

@@ -198,7 +198,11 @@ impl PngData {
         // Palette
         if let Some(ref palette) = self.raw.palette {
             let mut palette_data = Vec::with_capacity(palette.len() * 3);
-            let max_palette_size = 1 << (self.raw.ihdr.bit_depth.as_u8() as usize);
+            let mut max_palette_size = 1 << (self.raw.ihdr.bit_depth.as_u8() as usize);
+            // Ensure bKGD color doesn't get truncated from palette
+            if let Some(&idx) = self.raw.aux_headers.get(b"bKGD").and_then(|b| b.first()) {
+                max_palette_size = max_palette_size.max(idx as usize + 1);
+            }
             for px in palette.iter().take(max_palette_size) {
                 palette_data.extend_from_slice(px.rgb().as_slice());
             }

src/reduction/mod.rs

@@ -76,6 +76,13 @@ pub fn reduced_palette(png: &PngImage) -> Option<PngImage> {
             color_val(a.0).cmp(&color_val(b.0))
         });
 
+        // Make sure the background is also included, but only after sorting since it may not be used in idat
+        if let Some(&idx) = png.aux_headers.get(b"bKGD").and_then(|b| b.first()) {
+            if !used[idx as usize] {
+                used_enumerated.push((idx as usize, &true));
+            }
+        }
+
         let mut next_index = 0_u16;
         let mut seen = IndexMap::with_capacity(palette.len());
         for (i, used) in used_enumerated.iter().cloned() {
@@ -137,8 +144,7 @@ fn do_palette_reduction(png: &PngImage, palette_map: &[Option<u8>; 256]) -> Opti
 }
 
 fn palette_map_to_byte_map(png: &PngImage, palette_map: &[Option<u8>; 256]) -> Option<[u8; 256]> {
-    let len = png.palette.as_ref().map_or(0, |p| p.len());
-    if (0..len).all(|i| palette_map[i].map_or(true, |to| to == i as u8)) {
+    if (0..256).all(|i| palette_map[i].map_or(true, |to| to == i as u8)) {
         // No reduction necessary
         return None;
     }

tests/regression.rs

@@ -62,7 +62,13 @@ fn test_it_converts(
         "optimized to wrong bit depth"
     );
     if let Some(palette) = png.raw.palette.as_ref() {
-        assert!(palette.len() <= 1 << (png.raw.ihdr.bit_depth.as_u8() as usize));
+        let mut max_palette_size = 1 << (png.raw.ihdr.bit_depth.as_u8() as usize);
+        // Ensure bKGD color is valid
+        if let Some(&idx) = png.raw.aux_headers.get(b"bKGD").and_then(|b| b.first()) {
+            assert!(palette.len() > idx as usize);
+            max_palette_size = max_palette_size.max(idx as usize + 1);
+        }
+        assert!(palette.len() <= max_palette_size);
     } else {
         assert_ne!(png.raw.ihdr.color_type, ColorType::Indexed);
     }