How to Cook with Rhai

I’ve been working on embedded projects for both work and play where I want to make said embedded devices (usually an ESP32) do something unorthodox: run a sandboxed, dynamically editable script at runtime, fast. One of them involves this amazing piece:

A 3.5’ tall, 50 lb traffic light I got off of Facebook Marketplace. It will get its own blog post at some point, but its relevance here is my desire to allow friends and apps to dynamically control the lamps whenever they wanted, at low latency. One of the more exciting applications is for Second Fridays, a multi-genre performance night I run every two weeks, often frequented by DJs who want a fun lighting setup. I quickly whittled down my options: driving each lamp change via a WebSocket was out of the question due to the glacial roundtrip time; running Python would blow up the flash immediately.

After some cursory research, I found Rhai, an embeddable scripting language written in Rust. Anyone who’s seen any of my work over the past few years will know that the prospect of writing a project in Rust is irresistible to me. The great thing about Rhai is it makes full use of Rust’s semantics to make registering functions a breeze:

fn set_lights(r: bool, y: bool, g: bool) {
    // ...
}

fn register_functions(engine: &mut Engine) {
    // Automatically handles arguments
    engine.register_fn("set_lights", set_lights);
}

In the script:

let count = 0;
loop {
    set_lights(count == 0, count == 1, count == 2);
    count = (count + 1) % 3;
    sleep(0.5);
}

Easy as that. I quickly got to work implementing Rhai into my traffic light’s firmware, all of which is open source.

The first few usecases worked great. I created a hook for my agents to indicate status, gave my friends keys to page me via a pattern of their making, and even did some manual pattern displays for Second Fridays. This last item inspired me to try to make something automatic that analyzes music in real time and creates patterns derived from the music’s energy (you will see results at the end). I had AI whip up a proof of concept, and immediately ran into an out of memory error on the ESP32.

My first instinct was the AI had slopped all over the place, blowing up the size of the script well past what it had to be, and I was correct. Comments, em-dashes, ASCII diagrams galore.

In hindsight this was a problem, but not the problem.

Working from this premise, I decided to make something that would get rid of all the fluff: a minimizer.

Rhaiper

Enter Rhaiper, the Rhai Reaper. It’s essentially a very basic Rolldown for Rhai. Take the source, throw out the comments, compact the code, and send it off. For example, here is a Fibonacci script, before compaction:

// This Rhai script calculates the n-th Fibonacci number using a
// really dumb algorithm to test the speed of the scripting engine.

const TARGET = 28;
const REPEAT = 5;
const ANSWER = 317_811;

fn fib(n) {
    if n < 2 {
        n
    } else {
        fib(n-1) + fib(n-2)
    }
}

print(`Running Fibonacci(${TARGET}) x ${REPEAT} times...`);
print("Ready... Go!");

let result;
let now = timestamp();

for n in 0..REPEAT {
    result = fib(TARGET);
}

print(`Finished. Run time = ${now.elapsed} seconds.`);

print(`Fibonacci number #${TARGET} = ${result}`);

if result != ANSWER {
    print(`The answer is WRONG! Should be ${ANSWER}!`);
}

And after:

const TARGET=28;const REPEAT=5;const ANSWER=317_811;fn fib(n){if n<2{n}else{fib(n-1)+fib(n-2)}}print(`Running Fibonacci(${TARGET}) x ${REPEAT} times...`);print("Ready... Go!");let result;let now=
timestamp();for n in 0..REPEAT{result=fib(TARGET)}print(`Finished. Run time = ${now.elapsed} seconds.`);print(`Fibonacci number #${TARGET} = ${result}`);if result!=ANSWER{print(`The answer is WRONG! Should be ${ANSWER}!`)
}

The difference between Rhaiper and Rhai’s built-in compactor is Rhaiper outputs a SourceMap v3 file to reassign position data post-hoc, and it doesn’t concatenate the entire program to one line (something that breaks Rhai’s 16-bit error position reporting), instead choosing to break at regular intervals. It’s pretty conservative compared to Rolldown and others, avoiding renaming variables and other more destructive optimizations due to the complexity of doing so. However, even with just those light optimizations, it was now small enough for the traffic light! I flashed the new firmware, sent over a lighting control script, and got a different out of memory error.

I was somewhat confused as to what was happening at this point, so I dug in and looked at the allocation patterns of Rhai. This led me to an important discovery: Rhai makes thousands of tiny Box<> and Vec<> allocations on the heap during parsing. Each byte of source led to 24 bytes of allocated space on the ESP32. On most systems this is fine, but the ESP32 allocator has an 8-byte tag on each allocation. The allocation distribution below shows how a tree that’s already way bigger than the source is taxed even more (measuring follow.rhai)¹:

This amounts to around 10 KB of pure bookkeeping. I was pretty stuck after figuring this out, trying to think of any way to further squeeze bits out. I resorted to disabling parts of the runtime, more aggressive pruning, and a host of other measures to try to get my script running, to no avail. But then, inspiration struck to build a bytecode VM for Rhai.

Rhaigrain

To get an understanding of the architecture of Rhaigrain, let’s first look at where we’re coming from.

The AST Walker

The concept behind Rhai’s walker is to use the tokenized source code tree as a map for what to execute. Step through the script below (the same one we’ll compile to bytecode later on) to watch the walker descend into each node and return back up with a value:

It can get expensive to build this tree, resulting in around 10% more bytes allocated at peak than the final tree. This AST does come with some advantages, namely its ability to parse new code at runtime. Say I have this function:

fn foo(x, y, op) {
    if x < 0 || y < 0 {
        throw "error";
    }

    let result = eval(`${x} ${op} ${y}`);
    print(`Your result is: ${result}`);
}

The (simplified) AST before execution starts looks like this:

This may look unassuming, but eval is special: it evaluates Rhai code within its enclosing scope (just like the JS version). To handle this, another AST is allocated:

This flexibility makes it really easy to implement eval, import, and more. However, that comes at the aforementioned allocation cost that kills embedded platforms.

A Primer on Stack-based VMs

Stack-based VMs rely on their namesake, the stack data structure. To explain, consider the following demo:

fn do_op(a, b) {
    let result = 0;
    for i in 1..=a {
        result += i * b;
    }
    return result;
}

fn process_array(arr) {
    let prev = ();
    for itm in arr {
        if prev != () {
            let result = do_op(prev, itm);
            print(result);
        }
        prev = itm;
    }
}

process_array([2, 3, 4]);

A script is compiled down to a list of instructions, such as LoadLocal (push a variable that got a compile-time slot onto the stack), LoadNamed (push an external, non-compiled variable onto the stack by name, way slower than LoadLocal), Jump (set the program counter to a different location), or Call (call a function, which includes operators). The program counter determines the next instruction to be executed.

I created a prototype based on the concepts above that only supported a subset of the language, but it was enough for my use case.

Rhaigrain’s Architecture

Rhaigrain breaks down a program into discrete Chunks. A Chunk is metadata that describes the start and end addresses of a block of instructions (stored as indices into the whole bytecode) and the maximum size stack the Chunk will need. Each function is a Chunk, and the outermost scope is its own special Chunk, listed first. All the VM has to do is start executing at the first instruction.

To create a Chunk, the compiler “lowers” Rhai source to bytecode. Within a Chunk, the compiler allocates space for each local variable within the scope using Slots. Each slot within a Slots is indexed by a u16 and backed by a Scope at runtime: once a scope exits, all slots allocated by the scope are dropped. All statements and expressions are converted into opcodes and operands, with operands after opcodes.

This works for a grand majority of programs you could compile, but there are a few exceptions, such as eval(). These constructs use the lowering fallback Op::EvalAst. As a last resort, the compiler will store an AST in the program that can be executed at runtime. This is of course less than ideal for performance (and disables serialization), but it was critical in allowing for gradual adoption of the more obscure and complicated language features over time without breaking.

To give an idea of what real code looks like after being compiled, the fib(n) function from before compiles down to this:

fn #8 (1 params)  [185..249]  max_stack 3 // fn fib
    185            LoadLocal(0) // n
    188            Const(11) // 2
    191  9:10      Call { name: 13, argc: 2, op: Some(1), capture_parent_scope: false } // less_than(n, 2)
    197  9:10      JumpIfFalse { target: 210 } // n < 2 ?
    202            LoadLocal(0) // n
    205            Jump(248) // go to return
    210            LoadLocal(0) // n
    213            Const(12) // 1
    216  12:14     Call { name: 14, argc: 2, op: Some(2), capture_parent_scope: false } // subtract(n, 1)
    222  12:9      Call { name: 8, argc: 1, op: None, capture_parent_scope: false } // fib($RESULT) => fib(n - 1)
    226            LoadLocal(0) // n
    229            Const(11) // 2
    232  12:25     Call { name: 14, argc: 2, op: Some(2), capture_parent_scope: false } // subtract(n, 2)
    238  12:20     Call { name: 8, argc: 1, op: None, capture_parent_scope: false } // fib($RESULT) => fib(n - 2)
    242  12:18     Call { name: 15, argc: 2, op: Some(3), capture_parent_scope: false } // add($RESULT~1, $RESULT) => fib(n - 1) + fib(n - 2)
    248            Return // return whatever is on top of the stack

Note the strict bound on the stack size max_stack 3 (storing either n and 2, under the limit, or the results of the two recursive calls and their sum, at the limit) and the deduplicated constants (addresses 188 and 229).

Program { source: None, bytes: 249, max_stack: 2, consts: 13, names: 16, residuals: 0, compiled_fns: 1, walked_fns: 0, positions: true }

residuals (AST fragments left over): 0

main  [0..185]  max_stack 2
      0            Checkpoint
      1            Const(0)
      4            DeclareLocal { name: 0, is_const: true }
      7            Unit
      8            Pop
      9            Checkpoint
     10            Const(1)
     13            DeclareLocal { name: 1, is_const: true }
     16            Unit
     17            Pop
     18            Checkpoint
     19            Const(2)
     22            DeclareLocal { name: 2, is_const: true }
     25            Unit
     26            Pop
     27            Checkpoint
     28            Const(3)
     31  16:1      Call { name: 3, argc: 1, op: None, capture_parent_scope: false }
     35            Pop
     36            Checkpoint
     37            Const(4)
     40  17:1      Call { name: 3, argc: 1, op: None, capture_parent_scope: false }
     44            Pop
     45            Checkpoint
     46            Unit
     47            DeclareLocal { name: 4, is_const: false }
     50            Unit
     51            Pop
     52            Checkpoint
     53  20:11     Call { name: 5, argc: 0, op: None, capture_parent_scope: false }
     57            DeclareLocal { name: 6, is_const: false }
     60            Unit
     61            Pop
     62            Checkpoint
     63            Const(5)
     66  22:11     IterInit
     67            Unit
     68            DeclareLocal { name: 7, is_const: false }
     71  22:11     IterNext { exit: 99, indexed: false }
     76            StoreShared(5)
     79  22:20     Tick
     80            Const(0)
     83  23:14     Call { name: 8, argc: 1, op: None, capture_parent_scope: false }
     87  23:12     AssignLocal { slot: 3, var_name: 4, op: None }
     92            Unit
     93            Pop
     94            Jump(71)
     99            UnwindTo(5)
    102            Unit
    103            Jump(108)
    108            Pop
    109            Checkpoint
    110            InterpolateStart
    111            Const(6)
    114  26:7      InterpolateAppend
    115  26:34     Chain(0)
    118  26:34     InterpolateAppend
    119            Const(7)
    122  26:43     InterpolateAppend
    123            InterpolateEnd
    124  26:1      Call { name: 3, argc: 1, op: None, capture_parent_scope: false }
    128            Pop
    129            Checkpoint
    130            InterpolateStart
    131            Const(8)
    134  28:7      InterpolateAppend
    135            Const(0)
    138  28:28     InterpolateAppend
    139            Const(9)
    142  28:35     InterpolateAppend
    143            LoadLocal(3)
    146  28:40     InterpolateAppend
    147            InterpolateEnd
    148  28:1      Call { name: 3, argc: 1, op: None, capture_parent_scope: false }
    152            Pop
    153            Checkpoint
    154            LoadLocal(3)
    157            Const(2)
    160  30:11     Call { name: 12, argc: 2, op: Some(0), capture_parent_scope: false }
    166  30:11     JumpIfFalse { target: 183 }
    171            Const(10)
    174  31:5      Call { name: 3, argc: 1, op: None, capture_parent_scope: false }
    178            Jump(184)
    183            Unit
    184            Return

fn #8 (1 params)  [185..249]  max_stack 3
    185            LoadLocal(0)
    188            Const(11)
    191  9:10      Call { name: 13, argc: 2, op: Some(1), capture_parent_scope: false }
    197  9:10      JumpIfFalse { target: 210 }
    202            LoadLocal(0)
    205            Jump(248)
    210            LoadLocal(0)
    213            Const(12)
    216  12:14     Call { name: 14, argc: 2, op: Some(2), capture_parent_scope: false }
    222  12:9      Call { name: 8, argc: 1, op: None, capture_parent_scope: false }
    226            LoadLocal(0)
    229            Const(11)
    232  12:25     Call { name: 14, argc: 2, op: Some(2), capture_parent_scope: false }
    238  12:20     Call { name: 8, argc: 1, op: None, capture_parent_scope: false }
    242  12:18     Call { name: 15, argc: 2, op: Some(3), capture_parent_scope: false }
    248            Return

artifact 583 bytes, sidecar 85 positions

You can check if any ASTs were emitted using Program::residual_count(), which returns the number of residuals. Zero residuals means your program can be exported to bytecode. There are a ton of really interesting architectural decisions I didn’t go over, go take a look at the source code if you’re interested!

Result

A Program is loaded from raw bytes after compilation (potentially on another device). The important thing here is allocations drop dramatically: no allocation is needed for the program instructions themselves (they are read zero-copy from the container they arrived in), and stack allocations can be bound very tightly within a single backing store through static analysis. The only dynamic allocation is for the pools of constants, functions, and other identifiable items. This property leads to constant retained bytes and allocations with respect to the size of the script¹:

After eliminating the “death by a thousand cuts malloc()s” failure mode, I tried follow.rhai again:

That’s a 93% reduction in retained bytes and a 95% reduction in allocations with no parser peak at all! The bytecode is smaller than the script source, too, coming in at 7,419 bytes, or 90% of the original 8,251 byte script. It’s also up to 60% faster than the walker³:

We are still behind in some places so there’s more work to be done, but the results are already very promising. Incidentally the power of Rhaigrain means Rhaiper is obsolete on arrival; I am happy I built Rhaiper as an exercise, but Rhaigrain is here to stay.

All of these improvements finally led to a working DJ-connected traffic light⁴:

Shipping

In order to ship Rhaigrain, I had to gain access to some of Rhai’s internals (more than was exposed with the internals feature). Even though this was a purely personal project and could’ve existed as a fork, I decided to open a PR to potentially help others get more access to poke around within Rhai.

In my PR body I explained that I was making a bytecode VM and needed access to some more internals. Instead of just merging the small 150 line PR, the maintainer Stephen encouraged me to submit the entirety of Rhaigrain as a PR. He explained that he had been looking for someone to contribute something like it for a while, so this was the perfect opportunity to bring it in. I was surprised, but also very flattered! I proceeded to open an absolute monster of a PR, clocking in at a 19,000 line diff. I thought it had no chance of getting merged, but after a lot of back and forth with Stephen, into main it went!

After the merge, Stephen added me as a collaborator on the main Rhai repo, and now I review VM PRs! I’m really happy I get to contribute to a project I rely on so much, and Stephen has been a joy to work with. Soon after, we worked together on adding a bunch of missing features to what was now the grain module in Rhai: optimizations, evaluation order fixes, strict CI, and error reporting. I’m much happier with the state of grain now, but there’s always more work to do. I look forward to eventually adding import support and a better embedded experience that completely removes AST codepaths.

I strongly encourage you to contribute to Rhai or your own favorite open source project, there is always a massive difference to be made. You could be the person making a contribution the maintainers have been waiting for! You can try grain using the grain feature flag in the 1.26.0 or later release of Rhai.