150. The VM loop in Rust

The same loop, a different switch

Chapter 47’s dispatch loop in TypeScript, trimmed:

for (;;) {
  const instruction = readByte();
  switch (instruction) {
    case OP.CONSTANT:
      this.push(readConstant()); break;
    case OP.NIL: this.push(null); break;
    // ... 35 more cases
    default:
      throw new Error(`unknown opcode`);
  }
}

Here is drip-rs’s, adapted from run in src/vm.rs:

// the body of fn run(&mut self) -> Run
loop {
    let byte = self.read_byte();
    let op = OpCode::from_byte(byte)
        .expect("valid opcode");
    match op {
        OpCode::Constant => {
            let v = self.read_constant();
            self.push(v);
        }
        OpCode::Pop => {
            self.pop();
        }
        OpCode::Subtract => {
            let (a, b) =
                self.pop_numbers()?;
            self.push(Value::Number(a - b));
        }
        // ... one arm per opcode
    }
}

The shape is the same: read a byte, decode it, jump to the code for that instruction, repeat. Two lines differ. from_byte (chapter 148) turns the byte into an OpCode, and returns None for a byte no opcode has. And there is no default.

loop {
    let byte = self.read_byte();
    let op = OpCode::from_byte(byte)
        .expect("valid opcode");
    match op {
        OpCode::Constant => { .. }
        OpCode::Multiply => { .. }
        OpCode::Add => self.add()?,
        OpCode::Print => { .. }
        OpCode::Nil => self.push(Nil),
        OpCode::Return => { .. }
        // ... 31 more arms
    }
}

No default case

The match is over OpCode, an enum with 37 variants, and a Rust match must be exhaustive^[1]. Add a 38th opcode to the enum and forget its arm, and the VM no longer compiles. Here is rustc 1.97 on a three-variant version of the same mistake (trimmed):

error[E0004]: non-exhaustive patterns:
  `Op::False` not covered
 --> e2.rs:9:11
  |
9 |     match op {
  |           ^^ pattern `Op::False`
  |              not covered

In TypeScript, the same slip would surface only when a program ran the new instruction and hit default. In drip-rs, a bad byte is still possible, but only in one place: from_byte returning None. The compiler only ever writes valid opcodes, so the expect there is a statement of trust, and it would panic loudly if that trust were broken.

Errors are returned, not thrown

The TypeScript VM’s runtimeError throws, and interpret catches it. drip-rs returns errors as values instead. run returns Run, short for Result<(), RuntimeError>. An arm that can fail uses ?, as pop_numbers()? does above, so the error leaves run at once. Here is the start of pop_numbers:

let (Value::Number(a), Value::Number(b))
    = (self.peek(1), self.peek(0))
else {
    let m = "Operands must be numbers.";
    return Err(self.runtime_error(m));
};

That is a let ... else: if the pattern doesn’t match, the else block runs, and it must leave the function. RuntimeError carries no data, because runtime_error has already written the message and the stack trace to stderr.

The loop can’t hold the chunk

The TypeScript loop starts by caching the running frame’s code in local variables, and calls reload() after every call and return. clox does the same with a local frame pointer^[2]. Try that in Rust, on a toy VM:

fn run(&mut self) {
    let code = &self.code;
    for &b in code {
        self.push(b as f64);
    }
}

rustc 1.97 refuses (trimmed):

error[E0502]: cannot borrow `*self` as
  mutable because it is also borrowed
  as immutable
12 |  let code = &self.code;
   |             ---------- immutable
   |               borrow occurs here
14 |      self.push(b as f64);
   |      ^^^^^^^^^^^^^^^^^^^ mutable
   |               borrow occurs here

self.push needs &mut self: all of the VM, code included. A shared borrow of a field and a mutable borrow of the whole can’t overlap (chapter 125). The borrow checker is being right, too. In drip-rs, the chunk belongs to an ObjFunction in the heap’s arena, and allocating can grow that arena’s Vec and move every object in it. A reference into it could dangle.

So drip-rs keeps no long-lived reference at all. It copies the running frame’s state into fields of the Vm, which the README calls registers:

pub struct Vm {
    // ...
    ip: usize,          // next byte to run
    chunk: Rc<Chunk>,   // running bytecode
    slots: usize,       // frame's slot 0
}

fn read_byte(&mut self) -> u8 {
    let byte = self.chunk.code[self.ip];
    self.ip += 1;
    byte
}

Each read borrows self for one expression, and the borrow is over before push runs. The chunk is behind an Rc (chapter 137) because two owners need it: the function in the heap, and the VM while that function runs. Cloning an Rc just bumps a count. The compiler freezes each chunk with Rc::new when the function is done, so it is never changed afterwards and needs no RefCell. call saves the caller’s ip into its frame and loads the callee’s ip, chunk and slots. pop_frame loads the caller’s back.

What the registers bought

The README says the registers made fib “about 1.4× faster”. The committed benchmark run has drip-rs at 1460 ms on fib, against 5685 ms for the TypeScript VM.

To see the effect myself, I built two variants of drip-rs in a scratch directory. Neither is committed. They differ only in read_byte, read_constant and the jumps. One keeps ip and the Rc<Chunk> in the frame and goes through self.frames.last_mut() for every byte. The other has no cached chunk at all: every byte walks from the frame to its closure, to the function, to the chunk, through the heap. All three passed the 246-program corpus. I timed fib.drip with release builds (chapter 142), best of 5, alternating, on a shared 4-core cloud VM, so expect noise:

registers (committed)    1396 ms
ip + Rc<Chunk> in frame  1716 ms
walk through the heap    2983 ms

Every Drip instruction reads at least one byte, so the cost of reaching that byte is paid on every step of the loop.

Your turn

#[derive(Clone, Copy, Debug)]
enum Op {
    Num(f64),
    Add,
    Sub,
    Mul,
    Neg,
    Ret,
}

struct Vm {
    code: Vec<Op>,
    ip: usize,
    stack: Vec<f64>,
}

impl Vm {
    fn new(code: Vec<Op>) -> Vm {
        let stack = Vec::new();
        Vm { code, ip: 0, stack }
    }

    fn push(&mut self, v: f64) {
        self.stack.push(v);
    }

    fn pop(&mut self) -> f64 {
        self.stack.pop().expect("empty")
    }

    fn step(&mut self) -> bool {
        todo!()
    }

    fn run(&mut self) -> f64 {
        todo!()
    }
}

#[derive(Clone, Copy, Debug, PartialEq)]
enum Value {
    Num(f64),
    Bool(bool),
}

#[derive(Clone, Copy)]
enum Op {
    Push(Value),
    Add,
    Neg,
    Ret,
}

use Value::*;

fn pop(s: &mut Vec<Value>) -> Value {
    s.pop().expect("stack underflow")
}

fn add(a: Value, b: Value)
    -> Result<Value, String> {
    todo!()
}

fn neg(v: Value)
    -> Result<Value, String> {
    todo!()
}

fn run(code: &[Op])
    -> Result<Value, String> {
    todo!()
}

What you learned

• Match-based dispatch: decode the byte with from_byte, then one exhaustive match with an arm per opcode. A forgotten opcode is a compile error (E0004), not a default case.
• Runtime errors are Result values. Arms use ?, and let ... else handles operands of the wrong type.
• The loop can’t keep &self.chunk across self.push: a shared and a mutable borrow would overlap, and the arena could move.
• Registers: ip, chunk: Rc<Chunk> and slots are fields of the Vm, synced with the frame on calls and returns. In my measurement, fib took 1396 ms with them and 2983 ms walking through the heap.