Webassembly in Remix

Intro

Today is a momentous day for me, a culmination of nearly seven months of waiting. My pull request adding Webassembly(often abbreviated as WASM) support to Remix has finally been merged! While not in an official release yet, I thought it would be interesting to revisit some of the neat things that can be done with WebAssembly in Remix that I talked at RemixConf way back in May this year.

But first, let's review how Webassembly in Remix can be used. A typical WASM bundle typically consists of two things, a binary.wasm file that houses the acutal compiled code, and a JS file that contains a number of autogenerated helper functions. These functions will convert function inputs and outputs from other types into i32,i64,f32, and f64 or v128, do other administrative tasks, and an init() function that is used to load the WASM file into memory. Above that, the rest is packaging.

The trouble comes getting the binary .wasm files through the Remix compiler and into the output bundle. The PR I submitted essentially tells esbuild that if it finds an import for a .wasm file, it should be treated as a static file, and dropped into the public/build/_assets folder as part of the bundle output. This works great for WebAssembly files that you control, and for npm packages that offer you flexibility, but becomes a bit dicier when dealing with packages that make assumptions about the end use environment. As is often the case, the JS ecosystem is squarely to blame here, making things harder then they need to be.

Loading Webassembly in JS Environments

In the end, the trouble comes down to how Webassembly files are instantiated. Before a function in webassembly can be called, it needs to be loaded into memory and analyzed by the JS environment running it. Because there is no common standard, and each environment was developed separately by different people at different times, implementations vary.

My favorite implementation comes from Node, which added support for Webassembly in Node v8. A sample from their docs can be seen below. The wasmFile is loaded into memory, instantiated, and then we can call into it!

// Assume add.wasm file exists that contains a single function adding 2 provided arguments
const fs = require('fs');

const wasmBuffer = fs.readFileSync('/path/to/add.wasm');
WebAssembly.instantiate(wasmBuffer).then(wasmModule => {
  // Exported function live under instance.exports
  const { add } = wasmModule.instance.exports;
  const sum = add(5, 6);
  console.log(sum); // Outputs: 11
});

Here's a sample from my remix-rust repo, in the rust_functions package. It gets generated when wasm-pack build --target node is run.

/* @param {number} a
* @param {number} b
* @returns {number}
*/
module.exports.add = function(a, b) {
    const ret = wasm.add(a, b);
    return ret >>> 0;
};

...
const bytes = require('fs').readFileSync(path);

const wasmModule = new WebAssembly.Module(bytes);
const wasmInstance = new WebAssembly.Instance(wasmModule, imports);
wasm = wasmInstance.exports;
module.exports.__wasm = wasm;

This also makes it pretty easy to call it, as I do in the example repo, from within Node in one of Remix's Action or Loader functions.

import { add } from "rust_functions";

export const action: ActionFunction = async ({ request }) => {
    const result = add(2,2);
      console.log("result", result);
      return json({
        result,
      });
};

Seems pretty straightforward right? Import and use.

Unforunately the Webassembly story inside the browser isn't quite as rosy. While supported, there's no integration in the browser between script type or import statements, so we need to manually load the .wasm file. Check out this code generated by wasm-pack --target web, which is designed to be loaded in the browser.

async function load(module, imports) {
    if (typeof Response === 'function' && module instanceof Response) {
        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 {
        const instance = await WebAssembly.instantiate(module, imports);

        if (instance instanceof WebAssembly.Instance) {
            return { instance, module };

        } else {
            return instance;
        }
    }
}

async function init(input) {
    if (typeof input === 'undefined') {
        input = new URL('rust_functions_bg.wasm', import.meta.url);
    }
    const imports = getImports();

    if (typeof input === 'string' || (typeof Request === 'function' && input instanceof Request) || (typeof URL === 'function' && input instanceof URL)) {
        input = fetch(input);
    }

    initMemory(imports);

    const { instance, module } = await load(await input, imports);

    return finalizeInit(instance, module);
} 

export { initSync }
export default init;

This is a bit more verbose than the previous example. The init() function gets the location of the Webassembly file, initialize memory for it, and then calls load(). load() streams the binary data into memory, instantiates it, and then returns the wasm object to be used with functions. Unlike the Node version, the Browser version gives the developer the responsibility to call init() manually before calling any of the Webassembly functions.

Depending on the complexity of the build pipeline, this can cause some pain. Forgetting to call init before using the add() function results in an undefined error that isn't always the easiest to track down and understand.

In addition, this code can ONLY run in the browser. If you're inside an SSR framework like Remix, trying to load your .wasm file in the server will error out, as Node does not have instantiateStreaming(), and it needs to be loaded into memory on the client, not the server.

Understanding when files get loaded and how to control that in your framework is crucial to using Webassembly bundled this way.

Remix gives you a few choices in regard to how to use it. First, it must be initialized. This can be done alongside the function call inside of your component

import init, {add} from "../../../rust_functions/build/browser/rust_functions";
import wasm from "../../rust_functions/build/browser/rust_functions_bg.wasm";

export default function Index() {
    // This useEffect prevents the add function from being called on the server
    useEffect(() => {
        init(wasm).then()  => {
            add(2,2)
        }, [])
    }
}

One of the neat tricks you can do in Webassembly is to load your WASM before hydration! This should decrease the time to interactivity, and make your UX more seamless. If I wasn't using the WASM across the entire site, I'd consider putting the below init() call inside of a conditional to prevent loading it for all pages.

//entry.client.server.tsx
...
import init from "../../rust_functions/build/browser/rust_functions";
import wasm from "../../rust_functions/build/browser/rust_functions_bg.wasm"

function hydrate() {
    React.startTransition(() => {
      init(wasm).then(() => hydrateRoot(
        document,
        <React.StrictMode>
          <RemixBrowser />
        </React.StrictMode>
      ))
    });
  }
  
  if (window.requestIdleCallback) {
    window.requestIdleCallback(hydrate);
  } else {
    window.setTimeout(hydrate, 1);
  }

// app/routes/index.tsx
export default function Index() {
    // Prevent WASM from running on the server
    useEffect(() => {
        add(2,2)
    }

In many ways, this is the simplest bundling method for static sites and sites that don't need to run in Node. The init() function can be called in a regular script tag, and the JS called as expected.

Bundlers!

The third way I've seen Webassembly handled, and one that is outside the scope of my PR, is to use wasm-pack --target bundler. This generates a third kind of javascript wrapper, designed to be bundled by Webpack. Here we can see what the main JS file looks like when bundled this way.

import * as wasm from "./rust_functions_bg.wasm";
export * from "./rust_functions_bg.js";

In this mode, it offloads the initialization of the Webassembly to the bundler itself, whcih has pluses and minuses. The bundler needs to understand and handle instantiation for the file. The esbuild webassembly plugin talk about this process.

A stub module is created that wraps the wasm import from the file above, like so:

import wasm from '/path/to/example.wasm'
export default (imports) =>
  WebAssembly.instantiate(wasm, imports).then(
    result => result.instance.exports)

Then esbuild creates another virtual module to hold the .wasm file's binary code, so that it can be called. Note that one would still need to do something like this in your Remix code:

import load from './example.wasm'
load(imports).then(exports => { ... })

Much like the previous method, this plugin won't work for server side JS, necessitating modification of the plugin to determine where in your framework the WASM will get called. I never quite got this method to work in Remix, but if you do, feel free to drop me a DM on Twitter or Mastodon! I suspect that a more traditional SPA app, that doesn't do SSR, would fare better.

Packaging Webassembly in NPM

Given the aforementioned WASM loading tricks, I'd like to propose a fairly simple way to package Webassembly so that it can be used by the most people. The previously mentioned rust-remix repo's rust_functions package provides an example of this, but I'll cover the salient points below.

Essentially we'll be providing two package sets, generated with both the web target for client side use, and the node target for server side use. This can bloat your npm package size, so if your users aren't likely to use a bundler, it may be worth splitting these into two packages.

{
    "name": "rust_functions",
    "version": "1.0.0",
    "license": "MIT",
    "scripts": {
      "build": "npm run build:browser && npm run build:node",
      "build:browser": "wasm-pack build --target web --out-dir ./build/browser && rimraf ./build/browser/package.json",
      "build:node": "wasm-pack build --target nodejs --out-dir ./build/node && rimraf ./build/node/package.json"
    },
    "sideEffects": false,
    "files": [
      "build"
    ],
    "types": "./build/browser/rust_functions.d.ts",
    "exports": {
      ".": {
        "browser": "./build/browser/rust_functions.js",
        "node": "./build/node/rust_functions.js"
      },
      "./binary.wasm": {
        "browser": "./build/browser/rust_functions_bg.wasm",
        "node": "./noop.js"
      }
    },
    "devDependencies": {
      "rimraf": "^3.0.2"
    }
  }

Making use of a dual build script and a node's conditional imports, we can choose which generated js file is called depending on the environment. Remix's server environment for node based targets uses node, and the browser will load the browser. For the server, the JS will handle loading the .wasm file, and thus does not need to be called outside the package. For the browser environment, we need to make the .wasm file importable.

I've also spent some time experimenting with WASM packages on edge function environments like Deno Deploy. For that, you'll want to use the browser one, and be careful to load it

We also need to remove the package.json files that are autogenerated by wasm-pack for both builds, so that we can control it. This allows us to use the parent one at our leisure.

Performance Optimizations

Javascript is not a high performance language, and it likely never will be. Because of those limitations, it often tries to chear. At every opportunity, a smart JS developer will use a Browser API or native OS library to do high performance tasks. Computationally intensive tasks include, but are not limited to image manipulation, encryption, hashing and compression, large amounts of JSON processing, and 2D/3D graphics.

This has often limited architecture choices. For example, if image resizing and transcoding is expensive to do in the browser, we're forced to do it on a server somewhere instead. The push for edge computing has limited the ability to use this escape hatch, although workarounds for that have been found.

In my Remix talk, I posited the idea of Rust/Webassembly as the final level of "progressive enhancement". Because sometimes JS just isn't fast enough, we should consider whether Webassembly might enable us to give a better user experience.

As an example of some of the possible performance benefits, I'm going to do a little microbenchmark. Let's say you want to generate a SHA256 hash in the browser or on the edge when a user uploads a file.

We'll compare the performance of Rust's sha2 crate compiled to Webassembly and Amazon's @aws-crypto/sha256-js package, which is a pure JS implementation. In this test, we'll hash a 5MB file with both in a Remix loader, which runs in Node. It was run ten times in a local environment, and the results were averaged.

Using Webassembly for this application net an approximately 234% reduction in hash time. That can have real implications for page load time and server usage under load. For those curious, the Remix code looks like this:

export const loader: LoaderFunction = async ({ request }) => {
  const buf = randomBytes(50000000);
//JS Reference Implemntation
  const jsStart = Date.now()
  const hasher = new Sha256();
  hasher.update(buf);
  const rawJsHash = await hasher.digest();
  const jsHash = btoa(String.fromCharCode.apply(null, rawJsHash));
  const jsEnd = Date.now()
  const jsTime = jsEnd - jsStart;

  // Rust -> WASM Implementation using wasm-bindgen and node
  const wasmStart = Date.now();
  const hash2 = hash_wasm_server(buf);
  const wasmEnd = Date.now();
  const wasmTime = wasmEnd - wasmStart;
}

The rust function hash_wasm_server() is pretty close to the example, and is defined thusly:

#[wasm_bindgen]
pub fn hash_wasm_server(buf: Uint8Array) -> String {
    let mut hasher = Sha256::new();
    let owned_buf = buf.to_vec();
    hasher.update(owned_buf);
    let hash = hasher.finalize();
    Base64::encode_string(&hash)
}

I'm hardly the only person to consider this approach. Discourse published a nice blog post about how they optimized image processing for their forum software using Rust and WebAssembly. There are definitely some caveats to this idea, it's not a panacea, but for certain applications the performance benefits are very, very real.

Games

I remember the heady days of Flash gaming back in the early aughts, when I was in middle school, back when people had iPods and were just starting to experience AutoTune. Yes, back in the aughts (2000-2010ish) I played a ton of flash games in school. Games like QWOP and game sites like Popcap and Newgrounds ruled my childhood. The rise of Rust and it's native compilation into WebAssembly have allowed an entirely new generation of game engines to target the browser. Notable examples among these are Bevy, which is capable of both 2d and 3d gaming, and Macroquad, which is simpler to use, have enabled a new generation of these games to spring up.

Since you've come this far, maybe it's time to take a break and check out some of the lovely Rust Webassembly games on itch.io. I particularily enjoyed playing some of the Bevy Jam 2 submissions like Elemental Sorcerer and Shanty Quest Treble at Sea With Bevy's recent release of 0.9, it's never been a better time to start making games!

Conclusion

Hopefully we all know a little bit more about how Javascript and Webassembly interact, and how use it from within Remix. I'm pretty excited about how Webassembly is progressing, and how it can enable new projects and better user experiences.