Wasm compiler

This article documents the next step in the process of allowing users to represent sound waves as arbitrary mathematical functions. I previously demonstrated a method of compiling s-expressions to a theoretical stack-based format that can be interpreted by a Scheme loop, but I now have my eyes set on an actual architecture: WebAssembly.

WebAssembly (wasm) is a new technology supported by the major browsers that compiles higher level languages like C or Rust down to static bytecode that runs sandboxed in the browser. This code can be considerably faster than Javascript, so it's worth considering for a performance-critical sector like real time audio. Projects like Emscripten make it easy to build wasm binaries from your language of choice, largely doing away with the need to manually write wasm instructions. Compiling user-specified functions is a different story however; I have to write my own mini s-expression→wasm compiler that lives in the Javascript runtime. What follows is a Scheme prototype of that compiler, since the JS code won't look nearly as nice.

Expected input

All functions are relatively normal Scheme arithmetic procedures.

(define (square x) (* x x))

(define (identity x) x)

(define (stupid x y)
  (if (> x 1)
    (if (< y 3)
      (+ y y y)
      (+ x x x))
    (if (<= y 8)
      (* x y y)
      (- y x))))

(define (mustbesame x y) (if (!= x y) 0 x))

There are some exceptions.

  1. All variables and constants are interpreted as 64 bit floats (f64/f32 is common in audio synthesis).
  2. All functions are referentially-transparent and return exactly one f64 value.
  3. There are no (let ...) bindings, lambdas, or recursions.
  4. All (if ...) expressions are ternary; they require exactly two branches.
  5. wasm has a single "not equal" instruction, which is represented by (!= ...) rather than the expected (not (= ...)).
  6. Function names must be readable by Javascript, which means the beloved (kebab-case-function ...) is out.

All of this makes for a poor general purpose programming experience, but it's good enough to express waves. These restrictions greatly simplify the compilation process as well.

Expected output

The above functions should compile into a single wasm module m.

m = new WebAssembly.Instance(new WebAssembly.Module(new Uint8Array(
  [0,97,115,109,1,0,0,0,1,12,2,96,1,124,1,124,96,2,124,124,1,124,3,5,4,0,0,
   1,1,7,43,4,6,115,113,117,97,114,101,0,0,8,105,100,101,110,116,105,116,
   121,0,1,6,115,116,117,112,105,100,0,2,10,109,117,115,116,98,101,115,97,
   109,101,0,3,10,117,4,7,0,32,0,32,0,162,11,4,0,32,0,11,79,0,32,0,68,0,0,
   0,0,0,0,240,63,100,4,124,32,1,68,0,0,0,0,0,0,8,64,99,4,124,32,1,32,1,32,
   1,160,160,5,32,0,32,0,32,0,160,160,11,5,32,1,68,0,0,0,0,0,0,32,64,101,4,
   124,32,1,32,1,32,0,162,162,5,32,1,32,0,161,11,11,11,22,0,32,1,32,0,98,4,
   124,68,0,0,0,0,0,0,0,0,5,32,0,11,11])))

All of the integers in the array are byte-sized wasm instructions. They are compiled by the browser into executable functions stored in m.exports. Paste the block above in the browser console and try it out.

>> m.exports.square(9)
←  81

A note on bytes

All the numbers in the big example array will eventually make sense. While it's common convention to represent bytes such as these in hexadecimal format, I've specified them as unsigned 8 bit integers, since this is how they are printed by Scheme's srfi-4 library. For a master list of wasm instructions in hex format, please consult this article.

Enough talk, get compiling

The actual s-expression→wasm compilation process is documented below. Dissecting a Minimal (Useful) Webassembly Module was an enormous help for me while writing this code, and I encourage you to read it as well. While most wasm tutorials are focused on compiling from C or writing raw WebAssembly in a human-readable text format, this piece breaks down the structure of the program byte by byte. My hope is to do something similar with my writing, but with less trivial examples. My code is also presented in a slightly different order, so please make sure to read closely.

Scheme

The compiler is written in Chicken Scheme, since it has a number of useful low-level memory constructs. It makes use of the following libraries:

(import
  (chicken bitwise)
  (chicken locative)
  (chicken memory)
  srfi-1
  srfi-4)

Fire up your text editor to a blank .scm file and follow along. Make sure that you understand these topics.

  1. folding
  2. alists
  3. quote syntax
  4. uint8 vectors
  5. creating pointers to vectors
  6. writing to pointers
  7. bit-shifts, AND, and OR

Numbers

WebAssembly has a small variety of numeric types, but wave functions are only concerned with two of them: 64 bit floats and variable-length unsigned integers. All numbers are expressed as a series of bytes. A f64 type is unsurprisingly represented by 8 unsigned integers.

> (f64 2.7182818284590452353602874713527)
  (105 87 20 139 10 191 5 64)

The easiest way to convert a float to a byte array is to allocate 8 bytes of memory, create a pointer to this array, and then write the number to it. Kooda was of great assistance in writing the following function.

(define (f64 n)
  (let* ((xs (make-u8vector 8 0))
         (p (make-locative xs)))
    (pointer-f64-set! p (exact->inexact n))
    (u8vector->list xs)))

All literal floats are preceded by the f64.const instruction, which is a byte valued at 68.

(define-constant CONST-F64 68)

Meaning a full representation of Euler's Number in the body of a function would be

68 105 97 20 139 10 191 5 64

Variable-length integers appear often in wasm as well. They are LEB128 encoded and represent the size of code blocks. Varints have no fixed length; the program will (theoretically) continue to read a series of bytes as a single number until it comes across a value with its upper bit set to 0, similar to a null terminator in a C string. WebAssembly in practice is only concerned with values 5 bytes in length or less. Because the first bit of a byte is used to signal the possibility of the next byte in the stream, only 7 bits of actual numerical data are stored in a given chunk. An unsigned varint is created by a loop that

  1. Isolates 7 bits from a number through a right shift and an AND operation against the number 127.
  2. Determines if more bits need to be read.
  3. If so, a bitwise OR operation between the 7 bit chunk and the number 128 is performed, setting the first bit of the chunk to 1. The loop recurses.
  4. If not, the first bit of the chunk is left unset to 0, signalling the end of the number.

The code

(define-constant SEVEN-BIT-CHUNK 127)
(define-constant NEXT-BYTE-EXISTS 128)

(define (varuint32 n)
  (let ((byte (bitwise-and n SEVEN-BIT-CHUNK))
        (new-n (arithmetic-shift n -7)))
    (if (<= new-n 0)
      `(,byte)
      (cons (bitwise-ior byte NEXT-BYTE-EXISTS) (varuint32 new-n)))))

This function is named (varuint32) to represent the 5 byte limit imposed by wasm, but it will create an arbitrarily long byte series.

> (varuint32 2384692347862139468721646123846)
  (198 246 218 164 165 247 216 208 149 218 155 218 149 195 7)

In the above example, 7 is the only byte with an unset leading bit.

((1 1 0 0 0 1 1 0)
 (1 1 1 1 0 1 1 0)
 (1 1 0 1 1 0 1 0)
 (1 0 1 0 0 1 0 0)
 (1 0 1 0 0 1 0 1)
 (1 1 1 1 0 1 1 1)
 (1 1 0 1 1 0 0 0)
 (1 1 0 1 0 0 0 0)
 (1 0 0 1 0 1 0 1)
 (1 1 0 1 1 0 1 0)
 (1 0 0 1 1 0 1 1)
 (1 1 0 1 1 0 1 0)
 (1 0 0 1 0 1 0 1)
 (1 1 0 0 0 0 1 1)
 (0 0 0 0 0 1 1 1))

Numbers are the most bit-mangled aspect of the compiler, so you're good to go if you understood all this. Everything else is just a matter of arranging fixnums representing wasm instructions in proper order.

Stacks

Function evaluation is expressed by pushing and popping values on/off a stack. If you are unfamiliar with this model, please read The Little Stacker, where I talk about evaluating s-expressions as stacks. Order-agnostic expressions such as addition and multiplication can simply be read onto the stack backwards and then have their operator applied (- (length exp) 2) times. (+ 1 2 3) would look like

f64.const 3
f64.const 2
f64.const 1
f64.add
f64.add

Of course this backwards model screws up subtraction, division, and comparisons. An (order-check) function is required to check for these specific operators and reverse their parameters.

(define (order-check xs)
  (let ((f (car xs)))
    (if (any (lambda (x) (eq? x f)) '(- / < > <= >=))
      `(,f ,@(reverse (cdr xs)))
      xs)))

(order-check (- 2 1)) ensures that the s-exp is actually evaluated as

f64.const 2
f64.const 1
f64.sub

Operations

All relevant wasm arithmetic codes are held in a constant associative-list FUNCTIONS. These are the byte instructions for f64.add, f64.sub, etc.

(define-constant FUNCTIONS
 '((+ 160)
   (- 161)
   (* 162)
   (/ 163)
   (= 97)
   (!= 98)
   (< 99)
   (> 100)
   (<= 101)
   (>= 102)))

This alist will be referenced later in the compilation process.

Understanding check

Consider this expression.

(- (+ 3 3 3) (+ 3 3))

It is represented by the following bytes.

68 0 0 0 0 0 0 8 64
68 0 0 0 0 0 0 8 64
68 0 0 0 0 0 0 8 64
160
160
68 0 0 0 0 0 0 8 64
68 0 0 0 0 0 0 8 64
160
161

Take a moment to consider what each byte is doing, and why it appears in the order that it does.

Since we're performing subtraction, (order-check) reverses the parameters to read.

(- (+ 3 3) (+ 3 3 3))

Then the expressions are read backwards onto the stack, starting with the first 3 in (+ 3 3 3).

68 is the byte code for f64.const, and

0 0 0 0 0 0 8 64

is the f64 representation of 3. We can see that three 3s are pushed on to the stack.

160 is the byte of f64.add. It is called twice, since (- (length (+ 3 3 3)) 2) is 2.

(+ 3 3) is pushed onto the stack next with two more calls to 68 and one more addition with 160.

These two sums are finally subtracted from one-another with 161, which represents f64.sub.

There's nothing particularly scary about all these numbers.

Branching

Implementing conditional evaluation in the previous article was something of an ordeal, requiring perfectly-measured jump instructions and two compiler passes. WebAssembly's branching model is merciful in comparison. This simple compiler only needs to make use of if, else, and end instructions.

(define-constant BRANCH-IF 4)
(define-constant BRANCH-ELSE 5)
(define-constant END 11)

The skeletal stack of a ternary if statement looks like

<boolean expression ...>
IF
<return type of IF statement>
<true branch instructions ...>
ELSE
<false branch instructions ...>
END

This isn't too far off from how branching is expressed in high level code. The only difference is the placement of a boolean before the if instruction. The return type of a conditional is always a f64 of course, which is represented by

(define-constant TYPE-F64 124)

With that in mind, consider the expression

(if (> 2 1) 333 444)

which compiles to

68 0 0 0 0 0 0 0 64
68 0 0 0 0 0 0 240 63
100
4
124
68 0 0 0 0 0 208 116 64
5
68 0 0 0 0 0 192 123 64
11

The first two lines are f64.const instructions for 2 and 1 respectively.

The next line is 100 for f64.gt, which compares the two values on the stack and returns a boolean value, in this case: true.

The if statement then commences with a 4 on the next line, and a f64 return type as 124 after that.

The true branch is a f64.const for 333.

else is represented by a 5.

The false branch is a f64.const for 444.

The conditional block ends with 11 for end.

Very simple stuff.

Variables

Parameters to functions and local bindings (which my compiler does not concern itself with) are retrieved by the get_local instruction.

(define-constant GET-LOCAL 32)

The sole argument to get_local is an unsigned varint representing the index of a parameter in the function args list. For a function (λ (x y) ...)

32 0
32 1

would push x and y onto the stack respectively.

Environments

An individual function is compiled with an environment. This is simply an alist formed by appending the FUNCTIONS constants and key-value pairs for a function's variables. If we were to compile (λ (x y) ...) once again, its environment would be

((+ 160)
 (- 161)
 ...
 (x 32 0)
 (y 32 1)) 

An environment is made from the arguments list (x y) with the (make-env) procedure.

(define (make-env vars)
  (append FUNCTIONS
          (map (lambda (x n) `(,x ,GET-LOCAL ,@(varuint32 n)))
               vars (iota (length vars)))))

A symbol x in an s-expression is evaluated against this environment, along with the number n, which is (- (length exp) 2).

(define-constant NOP 1)

(define (variadic? x)
  (or (eq? x '+) (eq? x '-) (eq? x '*) (eq? x '/)))

(define (env-ref env x n)
  (letrec ((ref (assoc x env)))
    (cond ((not ref) `(,NOP))
          ((variadic? x) (make-list n (cadr ref)))
          (else (cdr ref)))))

NOP means that no operation is performed. It is returned if an unidentified symbol is detected. If the symbol is a variadic operation (+ - * /), then a list representing n applications of it is returned. If not, the regular key value of the symbol is returned in a list.

The function body

s-expressions are compiled recursively, either through a call to (make-branch) in the event of an if statement, or (make-stack) for any other procedure. Detecting a branch is as simple as checking for the presence of an if at the head of a list.

(define (branch? xs)
  (eq? (car xs) 'if))

All arguments to the compiler are guaranteed to be lists with (list-wrap).

(define (list-wrap x)
  (if (not (list? x)) `(,x) x))

It's simply a matter of left-folding over the expression xs, checking for branches and flipping the order of parameters where necessary.

(define (make-stack env xs)
  (if (branch? xs)
    (make-branch env xs)
    (let ((n (- (length xs) 2)))
      (foldl (lambda (acc x)
               (cond ((list? x) `(,@(make-stack env x) ,@acc))
                     ((number? x) `(,CONST-F64 ,@(f64 x) ,@acc))
                     (else `(,@(env-ref env x n) ,@acc))))
             '()
             (order-check xs)))))

(define (make-branch env xs)
  (let ((pred (make-stack env (list-wrap (cadr xs))))
        (branch1 (make-stack env (list-wrap (caddr xs))))
        (branch2 (make-stack env (list-wrap (cadddr xs)))))
  `(,@pred ,BRANCH-IF ,TYPE-F64 ,@branch1 ,BRANCH-ELSE ,@branch2 ,END)))

These are bigger procedures, but they make use of all the concepts outlined above. The heart of the process involves the (foldl), where a symbol x is evaluated either as an f64, a symbol in the environment env, or a recursive call to (make-stack), and then appended to the accumulator list acc. This results in a properly ordered list of wasm byte instructions.

Consider the function

(λ (x y) (if (> x y) (+ x x) (+ y y)))

Its body would be compiled as

(make-stack
  (make-env '(x y))
  '(if (> x y) (+ x y) (- x y)))

and result in

(32 0
 32 1
 100
 4
 124
 32 1
 32 0
 160
 5
 32 0
 32 1
 161
 11)

You should be able to read this.

The full representation of a function body in a wasm module includes

  1. The size of all information that follows, as an unsigned varint.
  2. The number of locally-declared variables (always 0 in our case).
  3. Declarations of local variables (blank in our case).
  4. The code of the body.
  5. An end instruction.

A full function definition of the form

(define (name var1 var2 ...) (body ...))

is compiled with

(define (make-code f)
  (let* ((stack (make-stack (make-env (cdadr f)) (list-wrap (caddr f))))
         (body `(0 ,@stack ,END)))
    `(,@(varuint32 (length body)) ,@body)))

Which invokes (make-env) and (make-stack) on the relevant sections, then appends the length of the resulting bytes

> (make-code '(define (λ x y) (if (> x y) (+ x y) (- x y))))
  (21 0 32 0 32 1 100 4 124 32 1 32 0 160 5 32 0 32 1 161 11 11)

The name and signature of the function are not specified anywhere in the body. This is just an anonymous array of instructions. Other parts of the wasm module are needed to point to it.

Stepping back

With the representation of numbers and functions accounted for, it's now possible to examine the full structure of a WebAssembly module. Valid wasm code contains much more than a series of constructions. Our use case requires a module which contains

  1. A magic number
  2. A version number
  3. A function signature section
  4. A function declaration section
  5. A function export section
  6. A function body section

Magic and version numbers

All modules begin with the byte string "0asm", then the version number as a 32 bit integer. This is simply "1 0 0 0" for the time being.

(define-constant MAGIC-WORD '(0 97 115 109))
(define-constant VERSION '(1 0 0 0))

Sections

Everything after the preamble is segregated into a section. All sections follow the same format.

  1. Section ID (1 byte)
  2. Section length (unsigned varint)
  3. Number of specific entries (unsigned varint)
  4. Section data (arbitrary bytes)

There are special sections with a more complex specification, but they aren't needed here.

The section IDs are as follows

(define-constant SIGNATURES 1)
(define-constant FUNCTION-TYPES 3)
(define-constant EXPORTS 7)
(define-constant CODE 10)

A generic section creating procedure is defined later in this article, but for now it makes more sense to examine the sections themselves.

Function signatures

A function's signature is

  1. A byte identifying it as a function type
  2. The number of arguments it takes (varint)
  3. A list of which type each argument is (all f64 in this case)
  4. The number of values returned (always 1 in this case)
  5. The return type (always f64 in this case)

With these restrictions in mind, it's quite simple to generate a signature of the form

(define (λ vars ...) (body ...))

with

(define-constant TYPE-FUNC 96)

(define (function-signature f)
  (let* ((parameters (cdr (cadr f)))
         (n (length parameters)))
    `(,TYPE-FUNC ,@(varuint32 n) ,@(make-list n TYPE-F64) 1 ,TYPE-F64)))

The signature for

(define (λ x y) (+ x y))

is

(96 2 124 124 1 124)

This section only lists unique signatures, meaning that a module containing eleven monadic functions and twenty dyadic functions would only contain two signatures. If we consider the earlier example functions

(define (square x) ...)
(define (identity x) ...)
(define (stupid x y) ...)
(define (mustbesame x y) ...)

we realize that only two signatures are necessary. This is reflected in the wasm module

.... 1 12 2 96 1 124 1 124 96 2 124 124 1 124 ...

section tag:       1
length of section: 12
signature count:   2
monadic signature: 96 1 124 1 124
dyadic signature:  96 2 124 124 1 124

Function declarations

The next section is where each individual function is associated with a signature. An association is nothing more than an index in the signatures section. Since the four example functions are (in order) monadic, monadic, dyadic, and dyadic, it follows that the declaration section appears as

... 3 5 4 0 0 1 1 ...

section tag:       3
length of section: 5
declaration count: 4
monadic functions: 0 0
dyadic functions:  1 1

If there is an alist si—a signature index where all redundant signatures are removed, and a 2d list ss—a simple list of all functions' signatures in order, then these signatures can be matched to their index with

(define (match-to-types si ss)
  (map (lambda (x) (varuint32 (list-index (lambda (y) (equal? y x)) si))) ss))

Function exports

The actual names of the functions are specified in the exports section. An export contains

  1. Length of the following data (varint)
  2. The name of a function as a byte string
  3. The type of data being exported (always 0 for function)
  4. The function number (varint)

A name can be converted to bytes with

(define (symbol->bytes s)
  (map char->integer (string->list (symbol->string s))))

And a function f numbered n can be exported thus

(define-constant EXPORT-FUNC 0)

(define (export-signature f n)
  (let* ((name (symbol->bytes (car (cadr f)))))
    `(,@(varuint32 (length name)) ,@name ,EXPORT-FUNC ,@(varuint32 n))))

Observe the bytes of the exports section.

... 7 43 4 6 115 113 117 97 114 101 0 0 8 105 100 101 110
    116 105 116 121 0 1 6 115 116 117 112 105 100 0 2 10
    109 117 115 116 98 101 115 97 109 101 0 3 ...

section tag:            7
length of section:      43
export count:           4
function 1 name length: 6
function 1 name:        115 113 117 97 114 101
function 1 export type: 0
function 1 number:      0
function 2 name length: 8
function 2 name:        105 100 101 110 116 105 115 121
function 2 export type: 0
function 2 number:      1
function 3 name length: 6
function 3 name:        115 116 117 112 105 100
function 3 export type: 0
function 3 number:      2
function 4 name length: 10
function 4 name:        109 117 115 116 98 101 115 109 101
function 4 export type: 0
function 4 number:      3

Function bodies

The last section contains the code itself. The remainder of the byte array represents this section. Body entries are discussed in the earlier function compilation section.

... 10 117 4 bodies ...

section code:      10
length of section: 117
function count:    4
function bodies:   bodies ... 

Putting it together

A 2d list of functions of the form

((define (λ x) (body ...))
 (define (λ' x y) (body ...)) ...)

Is rendered into a 3d list of the form

(((function-signature λ) (export-signature λ 0) (make-code λ))
 ((function-signature λ') (export-signature λ 1) (make-code λ')) ...)

with the `(compile-functions) procedure

(define (compile-functions fs)
  (map (lambda (f n)
         `(,(function-signature f) ,(export-signature f n) ,(make-code f)))
       fs (iota (length fs))))

The columns of this list are isolated and operated upon in the (compile-wasm) procedure, which creates the individual sections and concatenates them into a single flat list.

(define (compile-wasm xs)
  (let* ((compiled (compile-functions xs))
         (signatures (map car compiled))
         (exports (map cadr compiled))
         (bodies (map caddr compiled))
         (signature-index (delete-duplicates signatures)))
    `(,@MAGIC-WORD ,@VERSION
      ,@(section SIGNATURES signature-index)
      ,@(section FUNCTION-TYPES (match-to-types signature-index signatures))
      ,@(section EXPORTS exports)
      ,@(section CODE bodies))))

This lisp list can be printed as a Javascript expression with

(define (list->js-array xs)
  (let* ((backwards (reverse xs))
         (last (car backwards))
         (all-but (reverse (cdr backwards))))
    (display "new WebAssembly.Instance(new WebAssembly.Module(new Uint8Array([")
    (for-each (lambda (x) (display x) (display ",")) all-but)
    (display last)
    (display "])))\n")))

Finally, a macro for a (wasm ...) syntax allows the user to write Scheme definition forms without having to quote them.

(define-syntax wasm
  (syntax-rules ()
    ((_ xs ...) (list->js-array (compile-wasm '(xs ...))))))

With this, you are free to define a block of arithmetic functions and have them compile to a standalone wasm module.

> (wasm
    (define (plus x y) (+ x y))
    (define (minus x y) (- x y)))
  new WebAssembly.Instance(new WebAssembly.Module(new Uint8Array([0,97,115,109,1,0,0,0,1,7,1,96,2,124,124,1,124,3,3,2,0,0,7,16,2,4,112,108,117,115,0,0,5,109,105,110,117,115,0,1,10,17,2,7,0,32,1,32,0,160,11,7,0,32,0,32,1,161,11])))

Conclusion

Even this radically simplified compiler took pages of text to explain. The true version of these functions will have to fill buffers with their values, which will involve memory allocation and looping. Looking forward to that. Then everything I've demonstrated here in (relatively) elegant Scheme will have to be re-written to Javascript. Fun fun. I think the end product will be worth all these efforts though.