4KB Rust & WASM: Finding functions

February 8, 2020

In the last post, we put together a minimal WASM file using Rust that runs directly in the browser. Next we need to make it do something by giving it a way to find and call javascript functions.

float w(vec3 p, vec3 l, vec3 l2, vec3 c) { p-=c; mat2 R; R[1]=l.xy;R[0]=vec2(l.y,-l.x); p.xz*=R; R[1]=l2.xy;R[0]=vec2(l2.y,-l2.x); p.yz*=R; return length(vec2(length(p.xy)-l.z,p.z))-.1; } void main() { vec2 a = (gl_FragCoord.xy * vec2(1.0/256.0)) - 0.5; vec3 O = vec3(a.xy,1.0); vec3 r = normalize(O); vec3 p = O; float d = 0., e = 0.; vec3 l = vec3(cos(time),sin(time), 0.5); vec3 l2 = vec3(cos(time*.37),sin(time*.37), 0.5); vec3 m = vec3(1.); vec3 c = vec3(0.,0.,1.6); float alpha = 1.0; vec3 o = vec3(.5+.49*sin(time*.3),.5+.49*sin(time*.7),.5+.49*sin(time*.42)); for (int i=0;i<20;i++) { e = w(p,l,l2,c); if(e>0.)d += e; p = O+d*r; } if(e>0.)alpha=max(1.0-e*3.,0.); vec3 E = vec3(.001,0.,0.); vec3 n = normalize(vec3( w(p+E,l,l2,c)-w(p-E,l,l2,c), w(p+E.yxy,l,l2,c)-w(p-E.yxy,l,l2,c), w(p+E.yyx,l,l2,c)-w(p-E.yyx,l,l2,c))); vec3 P=m-p,N=normalize(P); gl_FragColor=vec4(alpha*(pow(vec3(.01)+pow(dot(n,normalize(N-r)),2E2)+3.*fract(o)*max(0.,dot(n,N))/length(P),vec3(.45))),alpha); }

As we left it (see Part 1), our Rust program being compiled to WebAssembly was only able to call one javascript function we gave it to print to the debug console.

Since the target is to fit the final program in just 4096 bytes, we need to keep the amount of javascript as small as practical. The challenge will be to find the minimum needed to enable the Rust code to do everything it needs to.

Keeping javascript objects

Rust will need to be able to create, keep and use many javascript objects. So let’s do that with a stack (an array which we will only add/remove things from the end), called “s”:

s=[];

If we plan to keep all our javascript functions in s, then the single function we give Rust to call back to javascript could be:

(a,b)=>s[a](b)

That takes two parameters. The first is the absolute index of the function in the stack array, and the second is a value to pass to the function. It returns the result back to Rust.

To be useful, we are going to need some functions to call. Since we are going to need to work with the (javascript-side) stack, the best ones to start with are a Push (add one item to end of the stack) and Pop (remove one item from the end of the stack).

s=[U,O]=
  a=>s.push(a),
  _=>s.pop()
];

Note that we alse gave short names (U & O) to the functions so they can be cheaply called from other javascript functions later.

Let’s test this out from Rust (src/lib.rs):

#![no_std]

#[link(wasm_import_module = "i")]
extern {
    fn r(which:u32, a:u32) -> u32;
}

fn js_fn(which:u32, a:u32) -> u32 {
    unsafe {
        r(which, a)
    }
}

const FN_PUSH:u32 = 0;
const FN_POP:u32 = 1;

fn js_push(a:u32) -> u32 {
    js_fn(FN_PUSH, a)
}

fn js_pop() -> u32 {
    js_fn(FN_POP, 0)
}

const JS : &str = "-->\
<script>\
s=[U,O]=[\
  a=>s.push(a),\
  _=>s.pop()\
];\
fetch('').then(\
    r=>r.arrayBuffer().then(\
        b=>WebAssembly.instantiate(\
            b.slice(4),{\
                i:{r:(a,b)=>s[a](b)}\
            }).then(\
                o=>{o.instance.exports.d();}\
            )\
        )\
    )\
</script>";

#[no_mangle]
pub extern fn d() -> u32 { // entry point
    js_push(42);
    js_push(4096);
    let popped = js_pop();
    js_push(7);
    js_push(popped);

    JS.as_ptr() as u32 // Return this to prevent str js being removed from binary
}

use core::panic::PanicInfo;
#[panic_handler]
fn panic(_info: &PanicInfo) -> ! {
    loop {}
}

This compiles to 338 bytes.

Loading it to the browser, and opening the debug console, we see it worked as expected, adding 42, 7 and 4096 to the stack:

> s
< (5) [f, f, 42, 7, 4096]

A stack of arguments

So far, our Rust side is limited to calling javascript functions with only one argument. But we know that many functions are going to need more than one. So how can we call those?

Since we can already push values onto the stack from Rust, what we need is a way to collect a few of those on to the stack, and then call a function using those as arguments.

Javascript’s function objects have a method called “apply” which is what we need here.

Since “apply” takes an array of arguments, we first need a way to remove a specified number of values from the stack as an array:

a=>s.splice(-a,a)

We’ll give that a short name “P” so we can use it from within other functions. Now we can define our apply function:

a=>U(s[a].apply(O(),P(O())))

When called, this uses values from the stack as follows:

  • The function to call apply on is at index a in the stack
  • From the top of the stack, the “this” pointer if it is a class instance (otherwise can be zero/undefined)
  • How many arguments the function takes
  • That many values (last one top-most)

After calling, the result is pushed back to the stack.

Let’s try it out:

#![no_std]

#[link(wasm_import_module = "i")]
extern {
    fn r(which:u32, a:u32) -> u32;
}

fn js_fn(which:u32, a:u32) -> u32 {
    unsafe {
        r(which, a)
    }
}

const FN_PUSH:u32 = 0;
const FN_POP:u32 = 1;
const FN_APPLY:u32 = 3;
const FN_LOG:u32 = 4;

fn js_push(a:u32) -> u32 {
    js_fn(FN_PUSH, a)
}

fn js_pop() -> u32 {
    js_fn(FN_POP, 0)
}

fn js_apply(which:u32) -> u32 {
    js_fn(FN_APPLY, which)
}

const JS : &str = "-->\
<script>\
s=[U,O,P]=[\
  a=>s.push(a),\
  _=>s.pop(),\
  a=>s.splice(-a,a),\
  a=>U(t.apply(O(),P(O()))),\
  console.log,\
];\
fetch('').then(\
    r=>r.arrayBuffer().then(\
        b=>WebAssembly.instantiate(\
            b.slice(4),{\
                i:{r:(a,b)=>s[a](b)}\
            }).then(\
                o=>{o.instance.exports.d();}\
            )\
        )\
    )\
</script>";

#[no_mangle]
pub extern fn d() -> u32 { // entry point
    js_push(42);
    js_push(4096);
    let popped = js_pop();
    js_push(7);
    js_push(popped);
    js_push(3); // Number of arguments
    js_push(0); // "this" -> unusused for console.log
    js_apply(FN_LOG); // Should run "console.log.apply(0,[42,7,4096])"

    JS.as_ptr() as u32 // Return this to prevent str js being removed from binary
}

use core::panic::PanicInfo;
#[panic_handler]
fn panic(_info: &PanicInfo) -> ! {
    loop {}
}

It compiles to 415 bytes, and gives this output to the console (as expected):

42 7 4096

Notice how we added “console.log” to the stack so we could call it.

We could carry on like this and add the name of every javascript class and method we want to use. But some of those WebGL method names are quite long, so it seems like that might quickly use up a big part of our 4096 byte budget before we even get round to using the functions.

So, what to do?

A common approach here is to rename all functions (for example in the WebGL API), for example by keeping two characters only (say first and sixth) and hoping the result is unique.

We’re going to do something a little different. We will define a simple hash function that generates a fixed size number from a variable size string. Then we will apply that function to all the possible names in a namespace of interest, and pick the name that matches a hash value we pre-computed from the target name. To get the list of possible names, we will use the standard javascript function Reflect.ownKeys

Here we go:

h=>{ 
    t=Reflect.ownKeys(s[8]);
    U(t[
        t.map(
            a=>a.toString().split('').map(
                a=>a.charCodeAt()
            ).reduce(
                (a,b)=>((a<<1)+b*s[9]+(a>>11))&4095
            )
        ).findIndex(e=>e==h)
       ]
     )
    },

The function refers to a couple of fixed indices on the stack. The current namespace object will be at index 8, and a constant multiplier for the hash function will be at index 9. Sometimes we may need to change this constant to handle cases where two names in the namespace produce the same hash value (a collision).

Putting it all together, let’s try and print to the console after first dynamically looking up the console.log function:

#![no_std]

#[link(wasm_import_module = "i")]
extern {
    fn r(which:u32, a:u32) -> u32;
}

fn js_fn(which:u32, a:u32) -> u32 {
    unsafe {
        r(which, a)
    }
}

const FN_PUSH:u32 = 0;
const FN_POP:u32 = 1;
const FN_APPLY:u32 = 3;
const FN_DUP:u32 = 4;
const FN_SET:u32 = 5;
const FN_LOOKUP:u32 = 6;
const FN_GETNAME:u32 = 7;
const NAMESPACE:u32 = 8;
const MOD:u32 = 9;
const LOG:u32 = 10;

fn js_push(a:u32) -> u32 {
    js_fn(FN_PUSH, a)
}

fn js_pop() -> u32 {
    js_fn(FN_POP, 0)
}

fn js_dup(which:u32) -> u32 {
    js_fn(FN_DUP, which)
}

fn js_set(which:u32) -> u32 {
    js_fn(FN_SET, which)
}

fn js_apply(which:u32) -> u32 {
    js_fn(FN_APPLY, which)
}

fn js_lookup(which:u32) -> u32 {
    js_fn(FN_LOOKUP, which)
}

fn js_getname(hash:u32) -> u32 {
    js_fn(FN_GETNAME, hash)
}

const JS : &str = "-->\
<script>\
s=[U,O,P]=[\
  a=>s.push(a),\
  _=>s.pop(),\
  a=>s.splice(-a,a),\
  a=>U(s[a].apply(O(),P(O()))),\
  a=>U(s[a]),\
  a=>s[a]=O(),\
  a=>U(s[a][O()]),\
  h=>{\
    t=Reflect.ownKeys(s[8]);\
    U(t[\
        t.map(\
            a=>a.toString().split('').map(\
                a=>a.charCodeAt()\
            ).reduce(\
                (a,b)=>((a<<1)+b*s[9]+(a>>11))&4095\
            )\
        ).findIndex(e=>e==h)\
       ]\
     )\
    },\
  window,\
  411,\
];\
fetch('').then(\
    r=>r.arrayBuffer().then(\
        b=>WebAssembly.instantiate(\
            b.slice(4),{\
                i:{r:(a,b)=>s[a](b)}\
            }).then(\
                o=>{o.instance.exports.d();}\
            )\
        )\
    )\
</script>";

#[no_mangle]
pub extern fn d() -> u32 { // entry point
    js_getname(974); // 974="console"
    js_lookup(NAMESPACE); // Lookup console object from window
    js_set(NAMESPACE); // Set console as the namespace
    js_getname(2935); // 2935="log"
    js_lookup(NAMESPACE); // Lookup log object from console -> Now index 10

    js_push(42); // Push argument 42
    js_push(1); // 1 argument
    js_push(0); // "this" not needed
    js_apply(LOG); // Call console.log
    js_pop(); // Remove return value from calling console.log

    JS.as_ptr() as u32 // Return this to prevent str js being removed from binary
}

use core::panic::PanicInfo;
#[panic_handler]
fn panic(_info: &PanicInfo) -> ! {
    loop {}
}

Compiling and running this should print “42” to the browser debug console.

The compiled page size is now up to 637 bytes, and we still haven’t started finding or calling any WebGL functions. But at least we have the basic infrastructure in place to do that, and still thousands of bytes of our budget left for it. The journey to 4KB will continue in the next post.