OxiPNG + threads PoC

codecs/oxipng/Cargo.lock

@@ -78,6 +78,16 @@ dependencies = [
  "cc",
 ]
 
+[[package]]
+name = "console_log"
+version = "0.2.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "501a375961cef1a0d44767200e66e4a559283097e91d0730b1d75dfb2f8a1494"
+dependencies = [
+ "log",
+ "web-sys",
+]
+
 [[package]]
 name = "crc"
 version = "1.8.1"
@@ -202,6 +212,15 @@ dependencies = [
  "either",
 ]
 
+[[package]]
+name = "js-sys"
+version = "0.3.41"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c4b9172132a62451e56142bff9afc91c8e4a4500aa5b847da36815b63bfda916"
+dependencies = [
+ "wasm-bindgen",
+]
+
 [[package]]
 name = "lazy_static"
 version = "1.4.0"
@@ -323,6 +342,7 @@ dependencies = [
  "libdeflater",
  "log",
  "miniz_oxide",
+ "rayon",
  "rgb",
  "zopfli",
 ]
@@ -401,8 +421,11 @@ checksum = "d29ab0c6d3fc0ee92fe66e2d99f700eab17a8d57d1c1d3b748380fb20baa78cd"
 name = "squoosh-oxipng"
 version = "0.1.0"
 dependencies = [
+ "console_log",
+ "js-sys",
  "log",
  "oxipng",
+ "rayon",
  "wasm-bindgen",
 ]
 
@@ -483,6 +506,16 @@ version = "0.2.64"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "7f7b90ea6c632dd06fd765d44542e234d5e63d9bb917ecd64d79778a13bd79ae"
 
+[[package]]
+name = "web-sys"
+version = "0.3.41"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "863539788676619aac1a23e2df3655e96b32b0e05eb72ca34ba045ad573c625d"
+dependencies = [
+ "js-sys",
+ "wasm-bindgen",
+]
+
 [[package]]
 name = "zopfli"
 version = "0.4.0"

codecs/oxipng/Cargo.toml

@@ -9,9 +9,12 @@ publish = false
 crate-type = ["cdylib"]
 
 [dependencies]
-oxipng = { version = "3.0.0", default-features = false }
+oxipng = { version = "3.0.0", default-features = false, features = ["parallel"] }
 wasm-bindgen = "0.2.64"
 log = { version = "0.4", features = ["release_max_level_off"] }
+rayon = "1.3.0"
+js-sys = "0.3.37"
+console_log = "0.2.0"
 
 [profile.release]
 lto = true

codecs/oxipng/build.sh

@@ -0,0 +1,21 @@
+#!/bin/bash
+
+set -e
+
+echo "============================================="
+echo "Compiling wasm"
+echo "============================================="
+(
+  CC=/opt/wasi-sdk/bin/clang RUSTFLAGS='-C target-feature=+atomics,+bulk-memory' rustup run nightly wasm-pack build -t web -- -Z build-std=panic_abort,std
+  rm pkg/.gitignore
+)
+echo "============================================="
+echo "Compiling wasm  done"
+echo "============================================="
+
+echo "!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!"
+echo "Did you update your docker image?"
+echo "Run \`docker pull ubuntu\`"
+echo "Run \`docker pull rust\`"
+echo "Run \`docker build -t squoosh-hqx .\`"
+echo "!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!"

codecs/oxipng/main.js

@@ -0,0 +1,27 @@
+import initOxiPNG, { start_main_thread, optimise } from './pkg/squoosh_oxipng.js';
+import wasmUrl from "./pkg/squoosh_oxipng_bg.wasm";
+
+async function startMainThread() {
+  await initOxiPNG(fetch(wasmUrl));
+  start_main_thread({
+    length: navigator.hardwareConcurrency,
+    pop: () => ({
+      postMessage: data => {
+        postMessage({ type: 'spawn', data });
+      }
+    })
+  });
+  return {
+    optimise
+  };
+}
+
+const mainThread = startMainThread();
+addEventListener('message', async ({ data: { id, args } }) => {
+  try {
+    let result = (await mainThread).optimise(...args);
+    postMessage({ ok: true, id, result });
+  } catch (result) {
+    postMessage({ ok: false, id, result });
+  }
+});

codecs/oxipng/pkg/README.md

@@ -0,0 +1,5 @@
+# OxiPNG
+
+- Source: <https://github.com/shssoichiro/oxipng>
+- Version: v3.0.0
+- License: MIT

codecs/oxipng/pkg/squoosh_oxipng.d.ts

@@ -1,8 +1,42 @@
 /* tslint:disable */
 /* eslint-disable */
 /**
+* @param {Array<any>} workers
+*/
+export function start_main_thread(workers: Array<any>): void;
+/**
+* @param {number} thread
+*/
+export function start_worker_thread(thread: number): void;
+/**
 * @param {Uint8Array} data
 * @param {number} level
 * @returns {Uint8Array}
 */
 export function optimise(data: Uint8Array, level: number): Uint8Array;
+
+export type InitInput = RequestInfo | URL | Response | BufferSource | WebAssembly.Module;
+
+export interface InitOutput {
+  readonly start_main_thread: (a: number) => void;
+  readonly start_worker_thread: (a: number) => void;
+  readonly optimise: (a: number, b: number, c: number, d: number) => void;
+  readonly malloc: (a: number) => number;
+  readonly free: (a: number) => void;
+  readonly __wbindgen_export_0: WebAssembly.Memory;
+  readonly __wbindgen_malloc: (a: number) => number;
+  readonly __wbindgen_free: (a: number, b: number) => void;
+  readonly __wbindgen_start: () => void;
+}
+
+/**
+* If `module_or_path` is {RequestInfo} or {URL}, makes a request and
+* for everything else, calls `WebAssembly.instantiate` directly.
+*
+* @param {InitInput | Promise<InitInput>} module_or_path
+* @param {WebAssembly.Memory} maybe_memory
+*
+* @returns {Promise<InitOutput>}
+*/
+export default function init (module_or_path?: InitInput | Promise<InitInput>, maybe_memory: WebAssembly.Memory): Promise<InitOutput>;
+        

codecs/oxipng/pkg/squoosh_oxipng.js

@@ -1,2 +1,189 @@
-import * as wasm from "./squoosh_oxipng_bg.wasm";
-export * from "./squoosh_oxipng_bg.js";
+
+let wasm;
+let memory;
+
+const heap = new Array(32).fill(undefined);
+
+heap.push(undefined, null, true, false);
+
+function getObject(idx) { return heap[idx]; }
+
+let heap_next = heap.length;
+
+function dropObject(idx) {
+    if (idx < 36) return;
+    heap[idx] = heap_next;
+    heap_next = idx;
+}
+
+function takeObject(idx) {
+    const ret = getObject(idx);
+    dropObject(idx);
+    return ret;
+}
+
+function addHeapObject(obj) {
+    if (heap_next === heap.length) heap.push(heap.length + 1);
+    const idx = heap_next;
+    heap_next = heap[idx];
+
+    heap[idx] = obj;
+    return idx;
+}
+
+let cachedTextDecoder = new TextDecoder('utf-8', { ignoreBOM: true, fatal: true });
+
+cachedTextDecoder.decode();
+
+let cachegetUint8Memory0 = null;
+function getUint8Memory0() {
+    if (cachegetUint8Memory0 === null || cachegetUint8Memory0.buffer !== wasm.__wbindgen_export_0.buffer) {
+        cachegetUint8Memory0 = new Uint8Array(wasm.__wbindgen_export_0.buffer);
+    }
+    return cachegetUint8Memory0;
+}
+
+function getStringFromWasm0(ptr, len) {
+    return cachedTextDecoder.decode(getUint8Memory0().slice(ptr, ptr + len));
+}
+/**
+* @param {Array<any>} workers
+*/
+export function start_main_thread(workers) {
+    wasm.start_main_thread(addHeapObject(workers));
+}
+
+/**
+* @param {number} thread
+*/
+export function start_worker_thread(thread) {
+    wasm.start_worker_thread(thread);
+}
+
+let WASM_VECTOR_LEN = 0;
+
+function passArray8ToWasm0(arg, malloc) {
+    const ptr = malloc(arg.length * 1);
+    getUint8Memory0().set(arg, ptr / 1);
+    WASM_VECTOR_LEN = arg.length;
+    return ptr;
+}
+
+let cachegetInt32Memory0 = null;
+function getInt32Memory0() {
+    if (cachegetInt32Memory0 === null || cachegetInt32Memory0.buffer !== wasm.__wbindgen_export_0.buffer) {
+        cachegetInt32Memory0 = new Int32Array(wasm.__wbindgen_export_0.buffer);
+    }
+    return cachegetInt32Memory0;
+}
+
+function getArrayU8FromWasm0(ptr, len) {
+    return getUint8Memory0().subarray(ptr / 1, ptr / 1 + len);
+}
+/**
+* @param {Uint8Array} data
+* @param {number} level
+* @returns {Uint8Array}
+*/
+export function optimise(data, level) {
+    var ptr0 = passArray8ToWasm0(data, wasm.__wbindgen_malloc);
+    var len0 = WASM_VECTOR_LEN;
+    wasm.optimise(8, ptr0, len0, level);
+    var r0 = getInt32Memory0()[8 / 4 + 0];
+    var r1 = getInt32Memory0()[8 / 4 + 1];
+    var v1 = getArrayU8FromWasm0(r0, r1).slice();
+    wasm.__wbindgen_free(r0, r1 * 1);
+    return v1;
+}
+
+async function load(module, imports, maybe_memory) {
+    if (typeof Response === 'function' && module instanceof Response) {
+        memory = imports.wbg.memory = new WebAssembly.Memory({initial:17,maximum:16384,shared:true});
+        if (typeof WebAssembly.instantiateStreaming === 'function') {
+            try {
+                return await WebAssembly.instantiateStreaming(module, imports);
+
+            } catch (e) {
+                if (module.headers.get('Content-Type') != 'application/wasm') {
+                    console.warn("`WebAssembly.instantiateStreaming` failed because your server does not serve wasm with `application/wasm` MIME type. Falling back to `WebAssembly.instantiate` which is slower. Original error:\n", e);
+
+                } else {
+                    throw e;
+                }
+            }
+        }
+
+        const bytes = await module.arrayBuffer();
+        return await WebAssembly.instantiate(bytes, imports);
+
+    } else {
+        memory = imports.wbg.memory = maybe_memory;
+        const instance = await WebAssembly.instantiate(module, imports);
+
+        if (instance instanceof WebAssembly.Instance) {
+            return { instance, module };
+
+        } else {
+            return instance;
+        }
+    }
+}
+
+async function init(input, maybe_memory) {
+    if (typeof input === 'undefined') {
+        input = import.meta.url.replace(/\.js$/, '_bg.wasm');
+    }
+    const imports = {};
+    imports.wbg = {};
+    imports.wbg.__wbg_length_8dacd620a01b769a = function(arg0) {
+        var ret = getObject(arg0).length;
+        return ret;
+    };
+    imports.wbg.__wbindgen_object_drop_ref = function(arg0) {
+        takeObject(arg0);
+    };
+    imports.wbg.__wbg_pop_e063afe6cb737ed3 = function(arg0) {
+        var ret = getObject(arg0).pop();
+        return addHeapObject(ret);
+    };
+    imports.wbg.__wbg_new_891c121bc64f5c10 = function() {
+        var ret = new Array();
+        return addHeapObject(ret);
+    };
+    imports.wbg.__wbindgen_module = function() {
+        var ret = init.__wbindgen_wasm_module;
+        return addHeapObject(ret);
+    };
+    imports.wbg.__wbg_push_ffe5167b83871629 = function(arg0, arg1) {
+        var ret = getObject(arg0).push(getObject(arg1));
+        return ret;
+    };
+    imports.wbg.__wbindgen_memory = function() {
+        var ret = wasm.__wbindgen_export_0;
+        return addHeapObject(ret);
+    };
+    imports.wbg.__wbindgen_number_new = function(arg0) {
+        var ret = arg0;
+        return addHeapObject(ret);
+    };
+    imports.wbg.__wbg_postMessage_5ab140e61ca2cb42 = function(arg0, arg1) {
+        getObject(arg0).postMessage(takeObject(arg1));
+    };
+    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);
+    }
+
+    const { instance, module } = await load(await input, imports, maybe_memory);
+
+    wasm = instance.exports;
+    init.__wbindgen_wasm_module = module;
+    wasm.__wbindgen_start();
+    return wasm;
+}
+
+export default init;
+

codecs/oxipng/pkg/squoosh_oxipng_bg.d.ts

@@ -1,8 +1,11 @@
 /* tslint:disable */
 /* eslint-disable */
-export const memory: WebAssembly.Memory;
+export function start_main_thread(a: number): void;
+export function start_worker_thread(a: number): void;
 export function optimise(a: number, b: number, c: number, d: number): void;
 export function malloc(a: number): number;
 export function free(a: number): void;
+export const __wbindgen_export_0: WebAssembly.Memory;
 export function __wbindgen_malloc(a: number): number;
 export function __wbindgen_free(a: number, b: number): void;
+export function __wbindgen_start(): void;

codecs/oxipng/pkg/squoosh_oxipng_bg.js

@@ -1,60 +0,0 @@
-import * as wasm from './squoosh_oxipng_bg.wasm';
-
-const lTextDecoder = typeof TextDecoder === 'undefined' ? (0, module.require)('util').TextDecoder : TextDecoder;
-
-let cachedTextDecoder = new lTextDecoder('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 WASM_VECTOR_LEN = 0;
-
-function passArray8ToWasm0(arg, malloc) {
-    const ptr = malloc(arg.length * 1);
-    getUint8Memory0().set(arg, ptr / 1);
-    WASM_VECTOR_LEN = arg.length;
-    return ptr;
-}
-
-let cachegetInt32Memory0 = null;
-function getInt32Memory0() {
-    if (cachegetInt32Memory0 === null || cachegetInt32Memory0.buffer !== wasm.memory.buffer) {
-        cachegetInt32Memory0 = new Int32Array(wasm.memory.buffer);
-    }
-    return cachegetInt32Memory0;
-}
-
-function getArrayU8FromWasm0(ptr, len) {
-    return getUint8Memory0().subarray(ptr / 1, ptr / 1 + len);
-}
-/**
-* @param {Uint8Array} data
-* @param {number} level
-* @returns {Uint8Array}
-*/
-export function optimise(data, level) {
-    var ptr0 = passArray8ToWasm0(data, wasm.__wbindgen_malloc);
-    var len0 = WASM_VECTOR_LEN;
-    wasm.optimise(8, ptr0, len0, level);
-    var r0 = getInt32Memory0()[8 / 4 + 0];
-    var r1 = getInt32Memory0()[8 / 4 + 1];
-    var v1 = getArrayU8FromWasm0(r0, r1).slice();
-    wasm.__wbindgen_free(r0, r1 * 1);
-    return v1;
-}
-
-export const __wbindgen_throw = function(arg0, arg1) {
-    throw new Error(getStringFromWasm0(arg0, arg1));
-};
-

codecs/oxipng/spawn.js

@@ -0,0 +1,32 @@
+let main = new Worker("./main.js", { type: "module" });
+
+let workers = Array.from(
+  { length: navigator.hardwareConcurrency },
+  () => new Worker("./worker.js", { type: "module" })
+);
+
+main.addEventListener('message', ({ data: { type, data } }) => {
+  if (type === 'spawn') {
+    workers.pop().postMessage(data);
+  }
+});
+
+let ID = 0;
+
+export function optimise(...args) {
+  return new Promise((resolve, reject) => {
+    let sendId = ID++;
+
+    main.addEventListener('message', function onMessage({ data: { ok, id, result } }) {
+      if (id !== sendId) return;
+      main.removeEventListener('message', onMessage);
+      if (ok) {
+        resolve(result);
+      } else {
+        reject(result);
+      }
+    });
+
+    main.postMessage({ id: sendId, args });
+  });
+}

codecs/oxipng/src/lib.rs

@@ -1,8 +1,44 @@
 mod malloc_shim;
 
+use js_sys::Array;
 use wasm_bindgen::prelude::*;
+use wasm_bindgen::{JsCast, JsValue};
 use oxipng::AlphaOptim;
 
+#[wasm_bindgen]
+extern "C" {
+    type Worker;
+
+    #[wasm_bindgen(method, js_name = postMessage)]
+    fn post_message(worker: &Worker, msg: JsValue);
+}
+
+#[wasm_bindgen]
+pub fn start_main_thread(workers: Array) {
+    // console_log::init_with_level(log::Level::Trace);
+
+    rayon::ThreadPoolBuilder::new()
+        .num_threads(workers.length() as _)
+        .spawn_handler(move |thread| {
+            Ok(workers.pop().unchecked_into::<Worker>().post_message({
+                let arr = Array::new();
+                arr.push(&wasm_bindgen::module());
+                arr.push(&wasm_bindgen::memory());
+                arr.push(&JsValue::from(Box::into_raw(Box::new(thread)) as u32));
+                arr.into()
+            }))
+        })
+        .build_global()
+        .unwrap_throw()
+}
+
+#[wasm_bindgen]
+pub fn start_worker_thread(thread: *mut rayon::ThreadBuilder) {
+  // console_log::init_with_level(log::Level::Trace);
+
+  unsafe { Box::from_raw(thread) }.run()
+}
+
 #[wasm_bindgen(catch)]
 pub fn optimise(data: &[u8], level: u8) -> Vec<u8> {
   let mut options = oxipng::Options::from_preset(level);

codecs/oxipng/worker.js

@@ -0,0 +1,8 @@
+import initOxiPNG, { start_worker_thread } from './pkg/squoosh_oxipng.js';
+
+addEventListener('message', async ({ data: [module, memory, threadPtr] }) => {
+  // console.log([module, memory, threadPtr]);
+  await initOxiPNG(module, memory);
+  // console.log('Starting', threadPtr);
+  start_worker_thread(threadPtr);
+}, { once: true });

src/codecs/oxipng/encoder.ts

@@ -1,6 +1,7 @@
-import { optimise } from '../../../codecs/oxipng/pkg';
+// @ts-ignore
+import { optimise } from '../../../codecs/oxipng/spawn.js';
 import { EncodeOptions } from './encoder-meta';
 
 export async function compress(data: ArrayBuffer, options: EncodeOptions): Promise<ArrayBuffer> {
-  return optimise(new Uint8Array(data), options.level).buffer;
+  return (await optimise(new Uint8Array(data), options.level)).buffer;
 }

webpack.config.js

@@ -149,7 +149,7 @@ module.exports = async function(_, env) {
         {
           // Emscripten modules don't work with Webpack's Wasm loader.
           test: /\.wasm$/,
-          exclude: /_bg\.wasm$/,
+          exclude: /(?<!oxipng)_bg\.wasm$/,
           // This is needed to make webpack NOT process wasm files.
           // See https://github.com/webpack/webpack/issues/6725
           type: 'javascript/auto',
@@ -160,7 +160,7 @@ module.exports = async function(_, env) {
         },
         {
           // Wasm modules generated by Rust + wasm-pack work great with Webpack.
-          test: /_bg\.wasm$/,
+          test: /(?<!oxipng)_bg\.wasm$/,
           type: 'webassembly/experimental',
         },
         {