Integrate WASI module with Squoosh UI

codecs/mozjpeg/enc/mozjpeg_enc.d.ts

@@ -0,0 +1,51 @@
+export const enum MozJpegColorSpace {
+  GRAYSCALE = 1,
+  RGB,
+  YCbCr,
+}
+
+export interface EncodeOptions {
+  quality: number;
+  baseline: boolean;
+  arithmetic: boolean;
+  progressive: boolean;
+  optimize_coding: boolean;
+  smoothing: number;
+  color_space: MozJpegColorSpace;
+  quant_table: number;
+  trellis_multipass: boolean;
+  trellis_opt_zero: boolean;
+  trellis_opt_table: boolean;
+  trellis_loops: number;
+  auto_subsample: boolean;
+  chroma_subsample: number;
+  separate_chroma_quality: boolean;
+  chroma_quality: number;
+}
+
+export interface MozJPEGModuleExports {
+  memory: WebAssembly.Memory;
+  alloc(size: number): number;
+  dealloc(ptr: number): void;
+  encode(
+    data: number,
+    width: number,
+    height: number
+  ): number;
+  set_opts_quality( quality: number): void;
+  set_opts_baseline( baseline: boolean): void;
+  set_opts_arithmetic( arithmetic: boolean): void;
+  set_opts_progressive( progressive: boolean): void;
+  set_opts_optimize_coding( optimize_coding: boolean): void;
+  set_opts_smoothing( smoothing: number): void;
+  set_opts_color_space( color_space: number): void;
+  set_opts_quant_table( quant_table: number): void;
+  set_opts_trellis_multipass( trellis_multipass: boolean): void;
+  set_opts_trellis_opt_zero( trellis_opt_zero: boolean): void;
+  set_opts_trellis_opt_table( trellis_opt_table: boolean): void;
+  set_opts_trellis_loops( trellis_loops: number): void;
+  set_opts_auto_subsample( auto_subsample: boolean): void;
+  set_opts_chroma_subsample( chroma_subsample: number): void;
+  set_opts_separate_chroma_quality( separate_chroma_quality: boolean): void;
+  set_opts_chroma_quality( chroma_quality: number): void;
+}

src/features/encoders/mozJPEG/shared/meta.ts

@@ -13,8 +13,9 @@
 import {
   EncodeOptions,
   MozJpegColorSpace,
+  MozJPEGModuleExports,
 } from 'codecs/mozjpeg/enc/mozjpeg_enc';
-export { EncodeOptions, MozJpegColorSpace };
+export { EncodeOptions, MozJpegColorSpace, MozJPEGModuleExports };
 
 export const label = 'MozJPEG';
 export const mimeType = 'image/jpeg';

src/features/encoders/mozJPEG/worker/mozjpegEncode.ts

@@ -10,23 +10,44 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-import mozjpeg_enc, { MozJPEGModule } from 'codecs/mozjpeg/enc/mozjpeg_enc';
-import { EncodeOptions } from '../shared/meta';
+import { MozJPEGModuleExports, EncodeOptions } from '../shared/meta';
 import wasmUrl from 'url:codecs/mozjpeg/enc/mozjpeg_enc.wasm';
-import { initEmscriptenModule } from 'features/worker-utils';
+import { instantiateStreaming } from 'features/worker-utils';
+import {
+  makeEverythingElseThrow,
+  makeWasiEnv,
+} from 'features/worker-utils/wasi-utils';
 
-let emscriptenModule: Promise<MozJPEGModule>;
+const instancePromise: Promise<WebAssembly.Instance> = instantiateStreaming(
+  fetch(wasmUrl),
+  {
+    wasi_snapshot_preview1: makeEverythingElseThrow(makeWasiEnv()),
+  },
+).then(({ instance }) => instance);
 
 export default async function encode(
   data: ImageData,
   options: EncodeOptions,
 ): Promise<ArrayBuffer> {
-  if (!emscriptenModule) {
-    emscriptenModule = initEmscriptenModule(mozjpeg_enc, wasmUrl);
-  }
+  const instance = await instancePromise;
+  const exports: MozJPEGModuleExports = (instance.exports as unknown) as MozJPEGModuleExports;
 
-  const module = await emscriptenModule;
-  const resultView = module.encode(data.data, data.width, data.height, options);
+  for (const [opt, value] of Object.entries(options)) {
+    // @ts-ignore Can’t be bothered to make these typings works
+    exports[`set_opts_${opt}`](value);
+  }
+  const inPtr = exports.alloc(data.data.byteLength);
+  new Uint8ClampedArray(exports.memory.buffer, inPtr, data.data.length).set(
+    data.data,
+  );
+  const resultPtr = exports.encode(inPtr, data.width, data.height);
+  const dv = new DataView(exports.memory.buffer);
+  const length = dv.getUint32(resultPtr, true);
+  const outPtr = dv.getUint32(resultPtr + 4, true);
+  const result = new Uint8Array(exports.memory.buffer, outPtr, length).slice();
+  exports.dealloc(inPtr);
+  exports.dealloc(outPtr);
+  exports.dealloc(resultPtr);
   // wasm can’t run on SharedArrayBuffers, so we hard-cast to ArrayBuffer.
-  return resultView.buffer as ArrayBuffer;
+  return result.buffer;
 }

src/features/worker-utils/index.ts

@@ -29,3 +29,16 @@ export function initEmscriptenModule<T extends EmscriptenWasm.Module>(
 export function blobToArrayBuffer(blob: Blob): Promise<ArrayBuffer> {
   return new Response(blob).arrayBuffer();
 }
+
+export async function instantiateStreaming(
+  resp: Response | PromiseLike<Response>,
+  importObj?: WebAssembly.Imports,
+): Promise<WebAssembly.WebAssemblyInstantiatedSource> {
+  if (WebAssembly.instantiateStreaming) {
+    return WebAssembly.instantiateStreaming(resp, importObj);
+  }
+  return WebAssembly.instantiate(
+    await Promise.resolve(resp).then((r) => r.arrayBuffer()),
+    importObj,
+  );
+}

src/features/worker-utils/wasi-utils.ts

@@ -0,0 +1,32 @@
+/**
+ * Copyright 2021 Google Inc. All Rights Reserved.
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+export function makeEverythingElseThrow(obj: { [x: string]: {} }): {} {
+  return new Proxy(obj, {
+    get(target, prop: string) {
+      if (prop in target) {
+        return target[prop];
+      }
+      return () => {
+        throw Error(`${prop} not implemented`);
+      };
+    },
+  });
+}
+
+export function makeWasiEnv() {
+  return {
+    environ_sizes_get: () => 0,
+    environ_get: () => 0,
+  };
+}