Separate palette sorting from reduction

benches/reductions.rs

@@ -238,7 +238,7 @@ fn reductions_palette_duplicate_reduction(b: &mut Bencher) {
     ));
     let png = PngData::new(&input, false).unwrap();
 
-    b.iter(|| palette::optimized_palette(&png.raw, false));
+    b.iter(|| palette::reduced_palette(&png.raw, false));
 }
 
 #[bench]
@@ -248,7 +248,7 @@ fn reductions_palette_unused_reduction(b: &mut Bencher) {
     ));
     let png = PngData::new(&input, false).unwrap();
 
-    b.iter(|| palette::optimized_palette(&png.raw, false));
+    b.iter(|| palette::reduced_palette(&png.raw, false));
 }
 
 #[bench]
@@ -258,7 +258,17 @@ fn reductions_palette_full_reduction(b: &mut Bencher) {
     ));
     let png = PngData::new(&input, false).unwrap();
 
-    b.iter(|| palette::optimized_palette(&png.raw, false));
+    b.iter(|| palette::reduced_palette(&png.raw, false));
+}
+
+#[bench]
+fn reductions_palette_sort(b: &mut Bencher) {
+    let input = test::black_box(PathBuf::from(
+        "tests/files/palette_8_should_be_palette_8.png",
+    ));
+    let png = PngData::new(&input, false).unwrap();
+
+    b.iter(|| palette::sorted_palette(&png.raw));
 }
 
 #[bench]

src/png/mod.rs

@@ -186,28 +186,12 @@ impl PngData {
         match &self.raw.ihdr.color_type {
             ColorType::Indexed { palette } => {
                 let mut palette_data = Vec::with_capacity(palette.len() * 3);
-                let mut max_palette_size = 1 << (self.raw.ihdr.bit_depth as u8);
+                for px in palette {
-                // 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());
                 }
                 write_png_block(b"PLTE", &palette_data, &mut output);
-                let num_transparent = palette.iter().take(max_palette_size).enumerate().fold(
+                if let Some(last_trns) = palette.iter().rposition(|px| px.a != 255) {
-                    0,
+                    let trns_data: Vec<_> = palette[0..=last_trns].iter().map(|px| px.a).collect();
-                    |prev, (index, px)| {
-                        if px.a == 255 {
-                            prev
-                        } else {
-                            index + 1
-                        }
-                    },
-                );
-                if num_transparent > 0 {
-                    let trns_data: Vec<_> =
-                        palette[0..num_transparent].iter().map(|px| px.a).collect();
                     write_png_block(b"tRNS", &trns_data, &mut output);
                 }
             }

src/reduction/mod.rs

@@ -57,6 +57,15 @@ pub(crate) fn perform_reductions(
         }
     }
 
+    // Attempt to reduce the palette
+    // This may change bytes but should always be beneficial
+    if opts.palette_reduction && !deadline.passed() {
+        if let Some(reduced) = reduced_palette(&png, opts.optimize_alpha) {
+            png = Arc::new(reduced);
+            reduction_occurred = true;
+        }
+    }
+
     // Now retain the current png for the evaluator baseline
     // It will only be entered into the evaluator if there are also others to evaluate
     let mut baseline = png.clone();
@@ -77,9 +86,9 @@ pub(crate) fn perform_reductions(
         }
     }
 
-    // Attempt to reduce the palette size
+    // Attempt to sort the palette
     if opts.palette_reduction && !deadline.passed() {
-        if let Some(reduced) = optimized_palette(&png, opts.optimize_alpha) {
+        if let Some(reduced) = sorted_palette(&png) {
             png = Arc::new(reduced);
             eval.try_image(png.clone());
             evaluation_added = true;
@@ -90,8 +99,8 @@ pub(crate) fn perform_reductions(
     if opts.color_type_reduction && !deadline.passed() {
         if let Some(reduced) = reduce_to_palette(&png) {
             png = Arc::new(reduced);
-            // Make sure the palette gets sorted (ideally, this should be done within reduce_to_palette)
+            // Make sure the palette gets sorted (but don't bother evaluating both results)
-            if let Some(reduced) = optimized_palette(&png, opts.optimize_alpha) {
+            if let Some(reduced) = sorted_palette(&png) {
                 png = Arc::new(reduced);
             }
             eval.try_image(png.clone());

src/reduction/palette.rs

@@ -1,27 +1,81 @@
 use crate::colors::{BitDepth, ColorType};
 use crate::headers::IhdrData;
 use crate::png::PngImage;
-use indexmap::map::{Entry::*, IndexMap};
+use indexmap::IndexSet;
 use rgb::RGBA8;
 
-/// Attempt to shrink and sort the palette, returning the optimized image if successful
+/// Attempt to reduce the number of colors in the palette, returning the reduced image if successful
 #[must_use]
-pub fn optimized_palette(png: &PngImage, optimize_alpha: bool) -> Option<PngImage> {
+pub fn reduced_palette(png: &PngImage, optimize_alpha: bool) -> Option<PngImage> {
     let palette = match &png.ihdr.color_type {
-        ColorType::Indexed { palette } => palette,
+        ColorType::Indexed { palette } if palette.len() > 1 => palette,
-        // Can't reduce if there is no palette
         _ => return None,
     };
-    if png.ihdr.bit_depth == BitDepth::One {
+
-        // Gains from 1-bit images will be at most 1 byte
+    let used = get_used_entries(png);
-        // Not worth the CPU time
+
-        return None;
+    let black = RGBA8::new(0, 0, 0, 255);
+    let mut condensed = IndexSet::with_capacity(palette.len());
+    let mut palette_map = [0; 256];
+    let mut did_change = false;
+    for (i, used) in used.iter().enumerate() {
+        if !used {
+            continue;
+        }
+        // There are invalid files that use pixel indices beyond palette size
+        let color = *palette.get(i).unwrap_or(&black);
+        palette_map[i] = add_color_to_set(color, &mut condensed, optimize_alpha);
+        if palette_map[i] as usize != i {
+            did_change = true;
+        }
     }
 
-    let mut palette_map = [None; 256];
+    // Update bKGD if it exists, ensuring it comes last in the palette if otherwise unused
+    let mut aux_headers = png.aux_headers.clone();
+    if let Some(idx) = aux_headers.remove(b"bKGD").and_then(|b| b.first().cloned()) {
+        if let Some(&color) = palette.get(idx as usize) {
+            let idx = add_color_to_set(color, &mut condensed, optimize_alpha);
+            aux_headers.insert(*b"bKGD", vec![idx]);
+        }
+    }
+
+    let data = if did_change {
+        // Reassign data bytes to new indices
+        let byte_map = palette_map_to_byte_map(png.ihdr.bit_depth, &palette_map);
+        png.data.iter().map(|b| byte_map[*b as usize]).collect()
+    } else if condensed.len() < palette.len() {
+        // Data is unchanged but palette will be truncated
+        png.data.clone()
+    } else {
+        // Nothing has changed
+        return None;
+    };
+
+    let palette: Vec<_> = condensed.into_iter().collect();
+
+    Some(PngImage {
+        ihdr: IhdrData {
+            color_type: ColorType::Indexed { palette },
+            ..png.ihdr
+        },
+        data,
+        aux_headers,
+    })
+}
+
+fn add_color_to_set(mut color: RGBA8, set: &mut IndexSet<RGBA8>, optimize_alpha: bool) -> u8 {
+    // If there are multiple fully transparent entries, reduce them into one
+    if optimize_alpha && color.a == 0 {
+        color.r = 0;
+        color.g = 0;
+        color.b = 0;
+    }
+    let (idx, _) = set.insert_full(color);
+    idx as u8
+}
+
+fn get_used_entries(png: &PngImage) -> [bool; 256] {
     let mut used = [false; 256];
-    {
-        // Find palette entries that are never used
     match png.ihdr.bit_depth {
         BitDepth::Eight => {
             for &byte in &png.data {
@@ -42,139 +96,106 @@ pub fn optimized_palette(png: &PngImage, optimize_alpha: bool) -> Option<PngImag
                 used[(byte >> 6) as usize] = true;
             }
         }
+        BitDepth::One => {
+            // Only two options, don't bother checking which are actually used
+            used[0] = true;
+            used[1] = true;
+        }
+        _ => unreachable!(),
+    };
+    used
+}
+
+fn palette_map_to_byte_map(bit_depth: BitDepth, palette_map: &[u8; 256]) -> [u8; 256] {
+    // Low bit-depths can be pre-computed for every byte value
+    match bit_depth {
+        BitDepth::Eight => *palette_map,
+        BitDepth::Four => {
+            let mut byte_map = [0_u8; 256];
+            for byte in 0..256 {
+                byte_map[byte] = palette_map[byte & 0x0F] | (palette_map[byte >> 4] << 4);
+            }
+            byte_map
+        }
+        BitDepth::Two => {
+            let mut byte_map = [0_u8; 256];
+            for byte in 0..256 {
+                byte_map[byte] = palette_map[byte & 0x03]
+                    | (palette_map[(byte >> 2) & 0x03] << 2)
+                    | (palette_map[(byte >> 4) & 0x03] << 4)
+                    | (palette_map[byte >> 6] << 6);
+            }
+            byte_map
+        }
         _ => unreachable!(),
     }
+}
 
-        let mut used_enumerated: Vec<(usize, &bool)> = used.iter().enumerate().collect();
+/// Attempt to sort the colors in the palette, returning the sorted image if successful
-        used_enumerated.sort_by(|a, b| {
+#[must_use]
-            //Sort by ascending alpha and descending luma.
+pub fn sorted_palette(png: &PngImage) -> Option<PngImage> {
-            let color_val = |i| {
+    if png.ihdr.bit_depth == BitDepth::One {
-                let color = palette
+        // Don't bother trying to sort a 1-bit image
-                    .get(i)
+        return None;
-                    .copied()
+    }
-                    .unwrap_or_else(|| RGBA8::new(0, 0, 0, 255));
+    let palette = match &png.ihdr.color_type {
+        ColorType::Indexed { palette } => palette,
+        _ => return None,
+    };
+
+    let mut enumerated: Vec<_> = palette.iter().enumerate().collect();
+
+    // If the background is the last entry in the palette we should make sure it stays last
+    // Otherwise an entry that's unused by the idat could prevent reduction to a lower depth
+    let mut aux_headers = png.aux_headers.clone();
+    let bkgd_idx = aux_headers.remove(b"bKGD").and_then(|b| b.first().cloned());
+    let bkgd_last = match bkgd_idx {
+        Some(idx) if idx as usize + 1 == palette.len() => enumerated.pop(),
+        _ => None,
+    };
+
+    // Sort the palette
+    enumerated.sort_by(|a, b| {
+        // Sort by ascending alpha and descending luma
+        let color_val = |color: &RGBA8| {
             ((color.a as i32) << 18)
             // These are coefficients for standard sRGB to luma conversion
             - i32::from(color.r) * 299
             - i32::from(color.g) * 587
             - i32::from(color.b) * 114
         };
-            color_val(a.0).cmp(&color_val(b.0))
+        color_val(a.1).cmp(&color_val(b.1))
     });
 
-        // Make sure the background is also included, but only after sorting since it may not be used in idat
+    if let Some(bkgd) = bkgd_last {
-        if let Some(&idx) = png.aux_headers.get(b"bKGD").and_then(|b| b.first()) {
+        enumerated.push(bkgd);
-            if !used[idx as usize] {
-                used_enumerated.push((idx as usize, &true));
-            }
     }
 
-        let mut next_index = 0_u16;
+    // Extract the new palette and determine if anything changed
-        let mut seen = IndexMap::with_capacity(palette.len());
+    let (old_map, palette): (Vec<_>, Vec<RGBA8>) = enumerated.into_iter().unzip();
-        for (i, used) in used_enumerated.iter().cloned() {
+    if old_map.iter().enumerate().all(|(a, b)| a == *b) {
-            if !used {
+        return None;
-                continue;
-            }
-            // There are invalid files that use pixel indices beyond palette size
-            let mut color = palette
-                .get(i)
-                .cloned()
-                .unwrap_or_else(|| RGBA8::new(0, 0, 0, 255));
-            // If there are multiple fully transparent entries, reduce them into one
-            if optimize_alpha && color.a == 0 {
-                color.r = 0;
-                color.g = 0;
-                color.b = 0;
-            }
-            match seen.entry(color) {
-                Vacant(new) => {
-                    palette_map[i] = Some(next_index as u8);
-                    new.insert(next_index as u8);
-                    next_index += 1;
-                }
-                Occupied(remap_to) => palette_map[i] = Some(*remap_to.get()),
-            }
-        }
     }
 
-    do_palette_reduction(png, palette, &palette_map)
+    // Construct the palette and byte maps and convert the data
-}
+    let mut new_map = [0; 256];
-
+    for (i, &v) in old_map.iter().enumerate() {
-#[must_use]
+        new_map[v] = i as u8;
-fn do_palette_reduction(
-    png: &PngImage,
-    palette: &[RGBA8],
-    palette_map: &[Option<u8>; 256],
-) -> Option<PngImage> {
-    let byte_map = palette_map_to_byte_map(png, palette_map)?;
-
-    // Reassign data bytes to new indices
-    let raw_data = png.data.iter().map(|b| byte_map[*b as usize]).collect();
-
-    let mut aux_headers = png.aux_headers.clone();
-    if let Some(bkgd_header) = png.aux_headers.get(b"bKGD") {
-        if let Some(Some(map_to)) = bkgd_header
-            .first()
-            .and_then(|&idx| palette_map.get(idx as usize))
-        {
-            aux_headers.insert(*b"bKGD", vec![*map_to]);
     }
+    let byte_map = palette_map_to_byte_map(png.ihdr.bit_depth, &new_map);
+    let data = png.data.iter().map(|&b| byte_map[b as usize]).collect();
+
+    // Update bKGD if it exists
+    if let Some(idx) = bkgd_idx.map(|idx| new_map[idx as usize]) {
+        aux_headers.insert(*b"bKGD", vec![idx]);
     }
 
     Some(PngImage {
         ihdr: IhdrData {
-            color_type: ColorType::Indexed {
+            color_type: ColorType::Indexed { palette },
-                palette: reordered_palette(palette, palette_map),
-            },
             ..png.ihdr
         },
-        data: raw_data,
+        data,
         aux_headers,
     })
 }
-
-fn palette_map_to_byte_map(png: &PngImage, palette_map: &[Option<u8>; 256]) -> Option<[u8; 256]> {
-    if (0..256).all(|i| palette_map[i].map_or(true, |to| to == i as u8)) {
-        // No reduction necessary
-        return None;
-    }
-
-    let mut byte_map = [0_u8; 256];
-
-    // low bit-depths can be pre-computed for every byte value
-    match png.ihdr.bit_depth {
-        BitDepth::Eight => {
-            for byte in 0..=255usize {
-                byte_map[byte] = palette_map[byte].unwrap_or(0)
-            }
-        }
-        BitDepth::Four => {
-            for byte in 0..=255usize {
-                byte_map[byte] = palette_map[byte & 0x0F].unwrap_or(0)
-                    | (palette_map[byte >> 4].unwrap_or(0) << 4);
-            }
-        }
-        BitDepth::Two => {
-            for byte in 0..=255usize {
-                byte_map[byte] = palette_map[byte & 0x03].unwrap_or(0)
-                    | (palette_map[(byte >> 2) & 0x03].unwrap_or(0) << 2)
-                    | (palette_map[(byte >> 4) & 0x03].unwrap_or(0) << 4)
-                    | (palette_map[byte >> 6].unwrap_or(0) << 6);
-            }
-        }
-        _ => {}
-    }
-
-    Some(byte_map)
-}
-
-fn reordered_palette(palette: &[RGBA8], palette_map: &[Option<u8>; 256]) -> Vec<RGBA8> {
-    let max_index = palette_map.iter().cloned().flatten().max().unwrap_or(0) as usize;
-    let mut new_palette = vec![RGBA8::new(0, 0, 0, 255); max_index + 1];
-    for (&color, &map_to) in palette.iter().zip(palette_map.iter()) {
-        if let Some(map_to) = map_to {
-            new_palette[map_to as usize] = color;
-        }
-    }
-    new_palette
-}