Wax: A User-Friendly Syntax for WebAssembly

The runtime of wasm_of_ocaml, the OCaml-to-WebAssembly compiler developed alongside js_of_ocaml, is nearly 27,000 lines of WebAssembly written by hand. It is the library of primitives that compiled OCaml code calls into for strings, polymorphic comparison, hashing, marshalling and effect handlers.

I wrote Wax to offer another way to write that code: a Rust-like syntax for WebAssembly that compiles to exactly the same instructions. It comes with a toolchain, written in OCaml, that converts between Wax and WebAssembly's text and binary formats in every direction.

One function, three spellings

WebAssembly (Wasm) is a compact binary instruction format for a stack machine: instructions take their operands from a stack and push their results back onto it. Every major browser ships a Wasm engine, and standalone engines such as Wasmtime run it outside the browser. A compiled program is called a module. The binary is not meant for human eyes, so the standard also defines a text format, WAT.

Here is a function that adds two numbers. As it ships in the binary, its body is seven bytes:

Parameters count as local variables (locals): local.get 0 and local.get 1 read the two arguments, and the leading 00 says the function declares no locals beyond them.

In the text format, as a disassembler prints it, the same function reads:

(func $add (export "add") (param $x i32) (param $y i32) (result i32)
  (i32.add (local.get $x) (local.get $y)))

This is the folded form of the text format: each instruction wraps its operands, so the stack mostly stays out of sight. But each parenthesised form is still one instruction. Reading a local takes a local.get, and every operator carries its type: i32.add adds two 32-bit integers. The (export "add") makes the function callable from outside the module, from JavaScript for instance, under the name add.

And here it is in Wax:

#[export = "add"]
fn add(x: i32, y: i32) -> i32 {
    x + y;
}

The #[export = "add"] attribute is the Wax spelling of that export. Compiling the Wax version gives the seven bytes above, byte for byte.

The trouble with 27,000 lines of WAT

The text format was designed so that a tool could print a binary and a human could check what the tool did. It was never meant for someone to sit down and write twenty-seven thousand lines of it.

Yet that is how the wasm_of_ocaml runtime is written: about fifty .wat files, all by hand, for reasons we explain at the end. Its maintainers, me included, work in S-expressions where every variable read is a local.get and every loop is a loop with a branch (br) back to its start.

The WebAssembly GC proposal, now part of the standard, added structs, arrays, nullable references, subtyping and downcasts, all managed by the engine's garbage collector. They let wasm_of_ocaml represent OCaml values directly, rather than as bytes in a linear memory, the single flat array of bytes that Wasm programs otherwise manage themselves. They also make the text format wordier. Reading a field looks like this:

(struct.get $list $value (local.get $l))

And in Wax:

l.value

Both are the same two instructions, a local.get and a struct.get, and Wax can convert between the two spellings in either direction.

What you write is what runs

Here is a complete function that walks a linked list and sums its values:

type list = { value: i32, next: &?list };

#[export = "sum"]
fn sum(l: &?list) -> i32 {
    let total: i32 = 0;
    while l is &list {
        total += l.value;
        l = l.next;
    }
    total;
}

Three bits of notation cover most of what you will meet in Wax code:

• &t is a reference to t, and &?t a reference that may be null.
• e is &t is a type test, and e as &t a cast.
• Types are those of Wasm: i32, i64, f32, f64, and references.

wax sum.wax -f wat compiles it to the text format:

(type $list (struct (field $value i32) (field $next (ref null $list))))

(func $sum (export "sum") (param $l (ref null $list)) (result i32)
  (local $total i32)
  (local.set $total (i32.const 0))
  (loop $loop
    (if (ref.test (ref $list) (local.get $l))
      (then
        (local.set $total
          (i32.add (local.get $total)
            (struct.get $list $value (local.get $l))))
        (local.set $l (struct.get $list $next (local.get $l)))
        (br $loop))))
  (local.get $total)
)

Each Wax construct maps directly onto Wasm instructions and types. The while is a loop with a br back to its start, the += is an i32.add followed by a local.set, and the type declaration is the struct definition.

Three formats, every direction

Wax converts between any two formats, in both directions, and its formatter rewrites .wax and .wat files in place. It is all one binary:

wax input.wax -o output.wasm    # compile to binary
wax input.wax -f wat            # compile to text
wax program.wasm                # decompile a binary to Wax
wax input.wat                   # convert WAT to Wax
wax format -i input.wax         # reformat in place
wax check input.wax             # type-check only

Running the compiler backwards

Compilers usually run one way. The most interesting engineering in Wax is in the reverse arrow. To show it at work, this is caml_string_equal from the wasm_of_ocaml runtime, as I wrote it. It implements OCaml's String.equal:

(func $caml_string_equal (export "caml_string_equal")
   (param $p1 (ref eq)) (param $p2 (ref eq)) (result (ref eq))
   (local $s1 (ref $bytes)) (local $s2 (ref $bytes))
   (local $len i32) (local $i i32)
   (if (ref.eq (local.get $p1) (local.get $p2))
      (then (return (ref.i31 (i32.const 1)))))
   (local.set $s1 (ref.cast (ref $bytes) (local.get $p1)))
   (local.set $s2 (ref.cast (ref $bytes) (local.get $p2)))
   (local.set $len (array.len (local.get $s1)))
   (if (i32.ne (local.get $len) (array.len (local.get $s2)))
      (then (return (ref.i31 (i32.const 0)))))
   (local.set $i (i32.const 0))
   (loop $loop
      (if (i32.lt_s (local.get $i) (local.get $len))
         (then
            (if (i32.ne (array.get_u $bytes (local.get $s1) (local.get $i))
                        (array.get_u $bytes (local.get $s2) (local.get $i)))
               (then (return (ref.i31 (i32.const 0)))))
            (local.set $i (i32.add (local.get $i) (i32.const 1)))
            (br $loop))))
   (ref.i31 (i32.const 1)))

Running it through wax gives:

#[export]
fn caml_string_equal(p1: &eq, p2: &eq) -> &eq {
    if p1 == p2 {
        return 1 as &i31;
    }
    let s1 = p1 as &bytes;
    let s2 = p2 as &bytes;
    let len = s1.length();
    if len != s2.length() {
        return 0 as &i31;
    }
    let i = 0;
    while i <s len : (i += 1) {
        if s1[i] as i32_u != s2[i] as i32_u {
            return 0 as &i31;
        }
    }
    1 as &i31;
}

In wasm_of_ocaml, &eq is the type of any OCaml value. Values that are not allocated, such as integers and booleans, are &i31: 31-bit integers packed into a reference, so 1 as &i31 is OCaml's true. A string is a &bytes, an array of 8-bit integers.

Some of the notation comes from Wasm itself. Wasm integers have no sign; each operation picks one. So <s is signed less-than (i32.lt_s), and s1[i] as i32_u reads a byte as an unsigned integer (array.get_u). Finally, a bare #[export] exports the function under its own name.

Now look at what the decompiler had to do to get there:

• Locals moved to where they are used. The text format declares every local at the top of the function. In Wax, each one became a let at its first assignment.
• The loop became a while. The decompiler recognised the loop ending in br $loop as a while. The increment, which must sit just before the jump back, moved into the loop header as a continue-expression, the : (i += 1) part, so the body holds only the body.
• Redundant types disappeared. s1 and s2 carry no annotation, because the casts already say what they are.

Across the whole file, about 480 lines of WAT shrink to about 320 lines of Wax, and they are less cluttered as well as shorter: no local.get on every variable read, no type spelled out on every operator, no locals declared far from where they are used. The result type-checks and compiles back. When the input is WAT, comments come through too, licence header included.

Decompile, recompile, same instructions

The reverse arrow is how a project with an existing .wat codebase gets to try Wax at all, so it has to be trustworthy, and "the module still does the same thing" is not good enough.

Recompiling the decompiled file gives back the same instruction stream: for string.wat, all 805 instructions, in the same order. Only names and numbering may change, and a running module cannot tell the difference. A local gets a new index when its declaration moves, and a function that was anonymous in the binary gets a name.

By default, the decompiler also tidies up a few idioms, reshaping a sequence into an equivalent one: an i32.eq followed by i32.eqz (a test for zero) becomes a single !=, and a cast that cannot fail is dropped. --faithful turns this off, so the result recompiles instruction for instruction, and a nightly fuzzing run holds it to that on machine-generated modules, unreachable code included. The default mode is checked on behaviour instead: modules from the official WebAssembly test suite go to Wax and back, and the suite's assertions are re-run against the result.

Errors, tools and editors

Wax has its own type checker, so a mistake is reported in your source, at the exact spot, rather than as a validation failure when an engine loads the module:

Error: This operator cannot be applied to operands of types i32 and f64.
 ──➤  hello.wax:3:7
1 │ #[export = "add"]
2 │ fn add(x: i32, y: f64) -> i32 {
3 │     x + y;
  ·       ^
4 │ }

I put a lot of effort into making messages like this one clear. Each of the hundreds of places where the parser can fail has its own message, saying what was expected there and, when an unclosed bracket is the likely cause, pointing back at it. These messages are generated from the grammar, with just a few written by hand, so they keep up as the language evolves. Automated checks make sure each one talks about what you wrote, statements and brackets rather than the parser's internals, and that what it says was expected is correct. Type errors get the same care.

Beyond type checking, Wax has:

• a formatter, and configurable lints that warn about suspicious code;
• conditional compilation, to include or leave out code depending on build flags;
• source maps, to step through the generated Wasm in browser DevTools.

You also get editor support. The VS Code extension runs the whole toolchain inside the editor, compiled to WebAssembly, so there is nothing else to install, and it works the same in desktop VS Code and on vscode.dev. Neovim and Emacs get syntax highlighting from a tree-sitter grammar, and the rest from wax lsp, a server for the Language Server Protocol (LSP) that most editors support: diagnostics as you type, types on hover, go to definition, find references, rename and completion. Any other editor with LSP support can use the same server.

None of this is reserved for .wax files. .wat files get the same type checking, with the same care in the messages. A type error underlines both the value of the wrong type and the instruction that expected another one:

Error:
  Type mismatch: this produces a value of type f64, but type i32 is expected.
 ──➤  bad.wat:3:30
1 │ (module
2 │   (func (export "f") (param $x i32) (result i32)
3 │     (i32.add (local.get $x) (f64.const 1))))
  ·                              ^^^^^^^^^^^
  ·      ^^^^^^^ expected here

They also get the formatter, the lints, conditional compilation, source maps, and the same editor support, so a project can use all of this on its existing .wat files.

Why not OCaml? Why not Rust?

OCaml programmers often ask why the runtime is not written in OCaml. Much of it exists to do what OCaml code cannot. Polymorphic comparison, hashing and marshalling walk any value whatever its type, so they must tell apart the Wasm structs, arrays, &i31 references and JavaScript values that represent it. Other primitives talk to the JavaScript environment the program runs in. OCaml compiled by wasm_of_ocaml only ever sees OCaml values, never the Wasm types behind them.

C and Rust both compile to Wasm, and they are the right tool for many applications, but not for a runtime like wasm_of_ocaml's, for two reasons.

First, they target the wrong half of Wasm. C and Rust compile to linear memory, which they manage with their own allocator. They have no way to express the GC proposal's structs, arrays and references, which the engine manages, and those are the values the OCaml runtime works on. A string is a GC array, a record is a GC struct, and a runtime primitive receives them as references and returns references.

The second reason is control. The runtime must decide exactly which types, casts and instructions the engine sees, because they are the contract with the code the OCaml compiler generates. caml_string_equal, for instance, must take two &eq and return an &eq, and must read a string as the &bytes array the compiled code allocated. Put a compiler in between, and it makes those choices for you.

Wax makes none. Every construct is a spelling of a Wasm instruction or type, and you can check by reading the WAT it produces.

Next: the wasm_of_ocaml runtime

Wax is still young, at release 0.2, but wasm_of_ocaml already uses it: its CI runs every runtime .wat file through Wax, so the runtime is type-checked and linted on each change. The next step is to switch to Wax for parsing those files when building wasm_of_ocaml. Moving the runtime itself to Wax syntax comes later, and it will come down to running the decompiler: the whole runtime already converts to Wax that type-checks and compiles back, conditional compilation included, and shrinks from nearly 27,000 lines to about 21,000.

Try it

The quickest way in is the playground. It runs the compiler in your browser, with nothing to install and nothing running on a server. Type Wax on the left and watch the WAT and the diagnostics update live on the right, or paste in a .wat file and read it the other way.

When you are ready to install it:

opam install wax    # or: npm install -g @wax-wasm/wax

or grab a prebuilt binary from the releases. In VS Code, install the Wax extension from the Marketplace; for Neovim, Emacs and other editors, the editor guide has ready-made configurations. The documentation has the language guide, the full Wax-to-WebAssembly correspondence and the CLI reference, and the code is on GitHub under Apache-2.0.

Until Next Time

If you write WebAssembly by hand, or read what your compiler emits, give Wax a try, and tell us on the issue tracker where it falls short.

You can connect with us on Bluesky, Mastodon, Threads, and LinkedIn or sign up for our mailing list to stay updated on our latest projects. We look forward to hearing from you!