Try to reduce alpha to tRNS

src/reduction/alpha.rs

@@ -1,4 +1,4 @@
-use crate::colors::ColorType;
+use crate::colors::{BitDepth, ColorType};
 use crate::headers::IhdrData;
 use crate::png::PngImage;
 
@@ -37,37 +37,61 @@ pub fn cleaned_alpha_channel(png: &PngImage) -> Option<PngImage> {
 }
 
 #[must_use]
-pub fn reduced_alpha_channel(png: &PngImage) -> Option<PngImage> {
+pub fn reduced_alpha_channel(png: &PngImage, optimize_alpha: bool) -> Option<PngImage> {
     let target_color_type = match png.ihdr.color_type {
         ColorType::GrayscaleAlpha => ColorType::Grayscale,
         ColorType::RGBA => ColorType::RGB,
         _ => return None,
     };
-    let byte_depth = png.ihdr.bit_depth.as_u8() >> 3;
-    let channels = png.channels_per_pixel();
+    let byte_depth = (png.ihdr.bit_depth.as_u8() >> 3) as usize;
+    let channels = png.channels_per_pixel() as usize;
     let bpp = channels * byte_depth;
-    let bpp_mask = bpp - 1;
-    if 0 != bpp & bpp_mask {
-        return None;
-    }
     let colored_bytes = bpp - byte_depth;
+
+    // If alpha optimisation is enabled, see if the image contains only fully opaque and fully transparent pixels.
+    // In case this occurs, we want to try and find an unused color we can use for the tRNS chunk.
+    // Rather than an exhaustive search, we will just keep track of 256 shades of gray, which should cover many cases.
+    let mut has_transparency = false;
+    let mut used_colors = vec![false; 256];
+
     for line in png.scan_lines() {
-        for (i, &byte) in line.data.iter().enumerate() {
-            if i as u8 & bpp_mask >= colored_bytes && byte != 255 {
+        for pixel in line.data.chunks(bpp) {
+            if optimize_alpha && pixel.iter().skip(colored_bytes).all(|b| *b == 0) {
+                // Fully transparent, we may be able to reduce with tRNS
+                has_transparency = true;
+            } else if pixel.iter().skip(colored_bytes).any(|b| *b != 255) {
+                // Partially transparent, the image is not reducible
                 return None;
+            } else if optimize_alpha && pixel.iter().take(colored_bytes).all(|b| *b == pixel[0]) {
+                // Opaque shade of gray, we can't use this color for tRNS
+                used_colors[pixel[0] as usize] = true;
             }
         }
     }
 
+    let transparency_pixel = if has_transparency {
+        // If no unused color was found we will have to fail here
+        // Otherwise, proceed to construct the tRNS chunk
+        let unused_color = used_colors.iter().position(|b| !*b)? as u8;
+        Some(match png.ihdr.bit_depth {
+            BitDepth::Sixteen => vec![unused_color; colored_bytes],
+            // 8-bit is still stored as 16-bit, with the high byte set to 0
+            _ => [0, unused_color].repeat(colored_bytes),
+        })
+    } else {
+        None
+    };
+
     let mut raw_data = Vec::with_capacity(png.data.len());
     for line in png.scan_lines() {
         raw_data.push(line.filter);
-        for (i, &byte) in line.data.iter().enumerate() {
-            if i as u8 & bpp_mask >= colored_bytes {
-                continue;
-            }
-
-            raw_data.push(byte);
+        for pixel in line.data.chunks(bpp) {
+            match transparency_pixel {
+                Some(ref trns) if pixel.iter().skip(colored_bytes).all(|b| *b == 0) => {
+                    raw_data.resize(raw_data.len() + colored_bytes, trns[1]);
+                }
+                _ => raw_data.extend_from_slice(&pixel[0..colored_bytes]),
+            };
         }
     }
 
@@ -86,7 +110,7 @@ pub fn reduced_alpha_channel(png: &PngImage) -> Option<PngImage> {
             ..png.ihdr
         },
         aux_headers,
-        transparency_pixel: None,
+        transparency_pixel,
         palette: None,
     })
 }

src/reduction/mod.rs

@@ -213,7 +213,7 @@ pub fn reduce_color_type(
         } else {
             None
         }
-        .or_else(|| reduced_alpha_channel(&reduced))
+        .or_else(|| reduced_alpha_channel(&reduced, optimize_alpha))
         {
             reduced = Cow::Owned(r);
         } else if let Some(r) = reduce_to_palette(&reduced) {
@@ -223,7 +223,9 @@ pub fn reduce_color_type(
     }
 
     if reduced.ihdr.color_type == ColorType::GrayscaleAlpha {
-        if let Some(r) = reduced_alpha_channel(&reduced).or_else(|| reduce_to_palette(&reduced)) {
+        if let Some(r) =
+            reduced_alpha_channel(&reduced, optimize_alpha).or_else(|| reduce_to_palette(&reduced))
+        {
             reduced = Cow::Owned(r);
             should_reduce_bit_depth = true;
         }