Merge pull request #438 from GoogleChromeLabs/rust-rotate

Rotate implementation in Rust

codecs/rotate/Dockerfile

@@ -0,0 +1,17 @@
+FROM ubuntu
+RUN apt-get update && \
+  apt-get install -qqy git build-essential cmake python2.7
+RUN git clone --recursive https://github.com/WebAssembly/wabt /usr/src/wabt
+RUN mkdir -p /usr/src/wabt/build
+WORKDIR /usr/src/wabt/build
+RUN cmake .. -DCMAKE_INSTALL_PREFIX=/opt/wabt && \
+  make && \
+  make install
+
+FROM rust
+RUN rustup install nightly && \
+  rustup target add --toolchain nightly wasm32-unknown-unknown
+
+COPY --from=0 /opt/wabt /opt/wabt
+ENV PATH="/opt/wabt/bin:${PATH}"
+WORKDIR /src

codecs/rotate/build.sh

@@ -0,0 +1,26 @@
+#!/bin/bash
+
+set -e
+
+echo "============================================="
+echo "Compiling wasm"
+echo "============================================="
+(
+  rustup run nightly \
+    rustc \
+    --target=wasm32-unknown-unknown \
+    -C opt-level=3 \
+    -o rotate.wasm \
+    rotate.rs
+    wasm-strip rotate.wasm
+)
+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-rotate .\`"
+echo "!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!"

codecs/rotate/package.json

@@ -0,0 +1,7 @@
+{
+  "name": "rotate",
+  "scripts": {
+    "build:image": "docker build -t squoosh-rotate .",
+    "build": "docker run --rm -v $(pwd):/src squoosh-rotate ./build.sh"
+  }
+}

codecs/rotate/rotate.rs

@@ -0,0 +1,81 @@
+#![no_std]
+#![no_main]
+
+use core::panic::PanicInfo;
+use core::slice::from_raw_parts_mut;
+
+#[no_mangle]
+fn rotate(input_width: isize, input_height: isize, rotate: isize) {
+  let mut i = 0isize;
+
+  // In the straight-copy case, d1 is x, d2 is y.
+  // x starts at 0 and increases.
+  // y starts at 0 and increases.
+  let mut d1_start: isize = 0;
+  let mut d1_limit: isize = input_width;
+  let mut d1_advance: isize = 1;
+  let mut d1_multiplier: isize = 1;
+  let mut d2_start: isize = 0;
+  let mut d2_limit: isize = input_height;
+  let mut d2_advance: isize = 1;
+  let mut d2_multiplier: isize = input_width;
+
+  if rotate == 90 {
+    // d1 is y, d2 is x.
+    // y starts at its max value and decreases.
+    // x starts at 0 and increases.
+    d1_start = input_height - 1;
+    d1_limit = input_height;
+    d1_advance = -1;
+    d1_multiplier = input_width;
+    d2_start = 0;
+    d2_limit = input_width;
+    d2_advance = 1;
+    d2_multiplier = 1;
+  } else if rotate == 180 {
+    // d1 is x, d2 is y.
+    // x starts at its max and decreases.
+    // y starts at its max and decreases.
+    d1_start = input_width - 1;
+    d1_limit = input_width;
+    d1_advance = -1;
+    d1_multiplier = 1;
+    d2_start = input_height - 1;
+    d2_limit = input_height;
+    d2_advance = -1;
+    d2_multiplier = input_width;
+  } else if rotate == 270 {
+    // d1 is y, d2 is x.
+    // y starts at 0 and increases.
+    // x starts at its max and decreases.
+    d1_start = 0;
+    d1_limit = input_height;
+    d1_advance = 1;
+    d1_multiplier = input_width;
+    d2_start = input_width - 1;
+    d2_limit = input_width;
+    d2_advance = -1;
+    d2_multiplier = 1;
+  }
+
+  let num_pixels = (input_width * input_height) as usize;
+  let in_b: &mut [u32];
+  let out_b: &mut [u32];
+  unsafe {
+    in_b = from_raw_parts_mut::<u32>(4 as *mut u32, num_pixels);
+    out_b = from_raw_parts_mut::<u32>((input_width * input_height * 4 + 4) as *mut u32, num_pixels);
+  }
+
+  for d2 in 0..d2_limit {
+    for d1 in 0..d1_limit {
+      let in_idx = (d1_start + d1 * d1_advance) * d1_multiplier + (d2_start + d2 * d2_advance) * d2_multiplier;
+      out_b[i as usize] = in_b[in_idx as usize];
+      i += 1;
+    }
+  }
+}
+
+#[panic_handler]
+fn panic(_info: &PanicInfo) -> ! {
+  loop {}
+}

src/codecs/rotate/processor-meta.ts

@@ -3,3 +3,10 @@ export interface RotateOptions {
 }
 
 export const defaultOptions: RotateOptions = { rotate: 0 };
+
+export interface RotateModuleInstance {
+  exports: {
+    memory: WebAssembly.Memory;
+    rotate(width: number, height: number, rotate: 0 | 90 | 180 | 270): void;
+  };
+}

src/codecs/rotate/processor.ts

@@ -1,73 +1,33 @@
-import { RotateOptions } from './processor-meta';
+import wasmUrl from '../../../codecs/rotate/rotate.wasm';
+import { RotateOptions, RotateModuleInstance } from './processor-meta';
 
-export function rotate(data: ImageData, opts: RotateOptions): ImageData {
-  const { rotate } = opts;
-  const flipDimensions = rotate % 180 !== 0;
-  const { width: inputWidth, height: inputHeight } = data;
-  const outputWidth = flipDimensions ? inputHeight : inputWidth;
-  const outputHeight = flipDimensions ? inputWidth : inputHeight;
-  const out = new ImageData(outputWidth, outputHeight);
-  let i = 0;
+const instancePromise = (WebAssembly as any).instantiateStreaming(fetch(wasmUrl));
 
-  // In the straight-copy case, d1 is x, d2 is y.
-  // x starts at 0 and increases.
-  // y starts at 0 and increases.
-  let d1Start = 0;
-  let d1Limit = inputWidth;
-  let d1Advance = 1;
-  let d1Multiplier = 1;
-  let d2Start = 0;
-  let d2Limit = inputHeight;
-  let d2Advance = 1;
-  let d2Multiplier = inputWidth;
+export async function rotate(
+  data: ImageData,
+  opts: RotateOptions,
+): Promise<ImageData> {
+  const { instance } = (await instancePromise) as {instance: RotateModuleInstance};
 
-  if (rotate === 90) {
-    // d1 is y, d2 is x.
-    // y starts at its max value and decreases.
-    // x starts at 0 and increases.
-    d1Start = inputHeight - 1;
-    d1Limit = inputHeight;
-    d1Advance = -1;
-    d1Multiplier = inputWidth;
-    d2Start = 0;
-    d2Limit = inputWidth;
-    d2Advance = 1;
-    d2Multiplier = 1;
-  } else if (rotate === 180) {
-    // d1 is x, d2 is y.
-    // x starts at its max and decreases.
-    // y starts at its max and decreases.
-    d1Start = inputWidth - 1;
-    d1Limit = inputWidth;
-    d1Advance = -1;
-    d1Multiplier = 1;
-    d2Start = inputHeight - 1;
-    d2Limit = inputHeight;
-    d2Advance = -1;
-    d2Multiplier = inputWidth;
-  } else if (rotate === 270) {
-    // d1 is y, d2 is x.
-    // y starts at 0 and increases.
-    // x starts at its max and decreases.
-    d1Start = 0;
-    d1Limit = inputHeight;
-    d1Advance = 1;
-    d1Multiplier = inputWidth;
-    d2Start = inputWidth - 1;
-    d2Limit = inputWidth;
-    d2Advance = -1;
-    d2Multiplier = 1;
+  // Number of wasm memory pages (á 64KiB) needed to store the image twice.
+  const bytesPerImage = data.width * data.height * 4;
+  const numPagesNeeded = Math.ceil(bytesPerImage * 2 / (64 * 1024) + 4);
+  // Only count full pages, just to be safe.
+  const numPagesAvailable = Math.floor(instance.exports.memory.buffer.byteLength / (64 * 1024));
+  const additionalPagesToAllocate = numPagesNeeded - numPagesAvailable;
+
+  if (additionalPagesToAllocate > 0) {
+    instance.exports.memory.grow(additionalPagesToAllocate);
   }
+  const view = new Uint8ClampedArray(instance.exports.memory.buffer);
+  view.set(data.data);
 
-  const inB = new Uint32Array(data.data.buffer);
-  const outB = new Uint32Array(out.data.buffer);
-  for (let d2 = d2Start; d2 >= 0 && d2 < d2Limit; d2 += d2Advance) {
-    for (let d1 = d1Start; d1 >= 0 && d1 < d1Limit; d1 += d1Advance) {
-      const start = ((d1 * d1Multiplier) + (d2 * d2Multiplier));
-      outB[i] = inB[start];
-      i += 1;
-    }
-  }
+  instance.exports.rotate(data.width, data.height, opts.rotate);
 
-  return out;
+  const flipDimensions = opts.rotate % 180 !== 0;
+  return new ImageData(
+    view.slice(bytesPerImage, bytesPerImage * 2),
+    flipDimensions ? data.height : data.width,
+    flipDimensions ? data.width : data.height,
+  );
 }