Merge pull request #972 from GoogleChromeLabs/png-decoder-colors

[cli] Handle all color types in png decoder

codecs/png/Cargo.lock

@@ -18,6 +18,12 @@ version = "3.4.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "2e8c087f005730276d1096a652e92a8bacee2e2472bcc9715a74d2bec38b5820"
 
+[[package]]
+name = "bytemuck"
+version = "1.5.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "bed57e2090563b83ba8f83366628ce535a7584c9afa4c9fc0612a03925c6df58"
+
 [[package]]
 name = "byteorder"
 version = "1.3.4"
@@ -112,11 +118,21 @@ dependencies = [
  "proc-macro2",
 ]
 
+[[package]]
+name = "rgb"
+version = "0.8.25"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "287f3c3f8236abb92d8b7e36797f19159df4b58f0a658cc3fb6dd3004b1f3bd3"
+dependencies = [
+ "bytemuck",
+]
+
 [[package]]
 name = "squoosh-png"
 version = "0.1.0"
 dependencies = [
  "png",
+ "rgb",
  "wasm-bindgen",
  "web-sys",
 ]

codecs/png/Cargo.toml

@@ -15,6 +15,7 @@ crate-type = ["cdylib"]
 png = "0.16.7"
 wasm-bindgen = "0.2.68"
 web-sys = { version = "0.3.45", features = ["ImageData"] }
+rgb = "0.8.25"
 
 [profile.release]
 lto = true

codecs/png/pkg/squoosh_png.js

@@ -1,6 +1,22 @@
 
 let wasm;
 
+let cachedTextDecoder = new TextDecoder('utf-8', { ignoreBOM: true, fatal: true });
+
+cachedTextDecoder.decode();
+
+let cachegetUint8Memory0 = null;
+function getUint8Memory0() {
+    if (cachegetUint8Memory0 === null || cachegetUint8Memory0.buffer !== wasm.memory.buffer) {
+        cachegetUint8Memory0 = new Uint8Array(wasm.memory.buffer);
+    }
+    return cachegetUint8Memory0;
+}
+
+function getStringFromWasm0(ptr, len) {
+    return cachedTextDecoder.decode(getUint8Memory0().subarray(ptr, ptr + len));
+}
+
 let cachegetUint8ClampedMemory0 = null;
 function getUint8ClampedMemory0() {
     if (cachegetUint8ClampedMemory0 === null || cachegetUint8ClampedMemory0.buffer !== wasm.memory.buffer) {
@@ -28,14 +44,6 @@ function addHeapObject(obj) {
     return idx;
 }
 
-let cachegetUint8Memory0 = null;
-function getUint8Memory0() {
-    if (cachegetUint8Memory0 === null || cachegetUint8Memory0.buffer !== wasm.memory.buffer) {
-        cachegetUint8Memory0 = new Uint8Array(wasm.memory.buffer);
-    }
-    return cachegetUint8Memory0;
-}
-
 let WASM_VECTOR_LEN = 0;
 
 function passArray8ToWasm0(arg, malloc) {
@@ -148,6 +156,9 @@ async function init(input) {
         var ret = new ImageData(v0, arg2 >>> 0, arg3 >>> 0);
         return addHeapObject(ret);
     };
+    imports.wbg.__wbindgen_throw = function(arg0, arg1) {
+        throw new Error(getStringFromWasm0(arg0, arg1));
+    };
 
     if (typeof input === 'string' || (typeof Request === 'function' && input instanceof Request) || (typeof URL === 'function' && input instanceof URL)) {
         input = fetch(input);

codecs/png/src/lib.rs

@@ -1,5 +1,7 @@
-use std::io::Cursor;
-
+use rgb::{
+    alt::{GRAY8, GRAYA8},
+    AsPixels, FromSlice, RGB8, RGBA8,
+};
 use wasm_bindgen::prelude::*;
 use wasm_bindgen::Clamped;
 
@@ -18,38 +20,57 @@ extern "C" {
     ) -> ImageData;
 }
 
-#[wasm_bindgen(catch)]
+#[wasm_bindgen]
 pub fn encode(data: &[u8], width: u32, height: u32) -> Vec<u8> {
-    let mut buffer = Cursor::new(Vec::<u8>::new());
+    let mut buffer = Vec::new();
 
     {
         let mut encoder = png::Encoder::new(&mut buffer, width, height);
         encoder.set_color(png::ColorType::RGBA);
         encoder.set_depth(png::BitDepth::Eight);
-        let mut writer = encoder.write_header().unwrap();
-        writer.write_image_data(data).unwrap();
+        let mut writer = encoder.write_header().unwrap_throw();
+        writer.write_image_data(data).unwrap_throw();
     }
 
-    buffer.into_inner()
+    buffer
 }
 
-#[wasm_bindgen(catch)]
-pub fn decode(data: &[u8]) -> ImageData {
-    let mut decoder = png::Decoder::new(Cursor::new(data));
+// Convert pixels in-place within buffer containing source data but preallocated
+// for entire [num_pixels * sizeof(RGBA)].
+// This works because all the color types are <= RGBA by size.
+fn expand_pixels<Src: Copy>(buf: &mut [u8], to_rgba: impl Fn(Src) -> RGBA8)
+where
+    [u8]: AsPixels<Src> + FromSlice<u8>,
+{
+    assert!(std::mem::size_of::<Src>() <= std::mem::size_of::<RGBA8>());
+    let num_pixels = buf.len() / 4;
+    for i in (0..num_pixels).rev() {
+        let src_pixel = buf.as_pixels()[i];
+        buf.as_rgba_mut()[i] = to_rgba(src_pixel);
+    }
+}
+
+#[wasm_bindgen]
+pub fn decode(mut data: &[u8]) -> ImageData {
+    let mut decoder = png::Decoder::new(&mut data);
     decoder.set_transformations(png::Transformations::EXPAND);
-    let (info, mut reader) = decoder.read_info().unwrap();
+    let (info, mut reader) = decoder.read_info().unwrap_throw();
     let num_pixels = (info.width * info.height) as usize;
     let mut buf = vec![0; num_pixels * 4];
-    reader.next_frame(&mut buf).unwrap();
+    reader.next_frame(&mut buf).unwrap_throw();
 
-    // Transformations::EXPAND will make sure color_type is either
-    // RGBA or RGB. If it’s RGB, we need inject an alpha channel.
-    if info.color_type == png::ColorType::RGB {
-        for i in (0..num_pixels).rev() {
-            buf[i * 4 + 0] = buf[i * 3 + 0];
-            buf[i * 4 + 1] = buf[i * 3 + 1];
-            buf[i * 4 + 2] = buf[i * 3 + 2];
-            buf[i * 4 + 3] = 255;
+    // Transformations::EXPAND will expand indexed palettes and lower-bit
+    // grayscales to higher color types, but we still need to transform
+    // the rest to RGBA.
+    match info.color_type {
+        png::ColorType::RGBA => {}
+        png::ColorType::RGB => expand_pixels(&mut buf, RGB8::into),
+        png::ColorType::GrayscaleAlpha => expand_pixels(&mut buf, GRAYA8::into),
+        png::ColorType::Grayscale => {
+            expand_pixels(&mut buf, |gray: GRAY8| GRAYA8::from(gray).into())
+        }
+        png::ColorType::Indexed => {
+            unreachable!("Found indexed color type, but expected it to be already expanded")
         }
     }