Writing a Tiny Stack VM in Rust
By: Mikail Khan
Tags: Interpreters, Programming, Rust, Programming Languages, Compilers

Well we can print stuff now, but the upper limit for program complexity is still quite low. In the next section, we’ll take this from a simple expression evaluator to aturing complete interpreter in just two steps.

First, we need jumps and conditional jumps. A jump takes us from one point in the code to another, and a conditional jump jumps only if a condition is met, otherwise it just moves on.

To add this, we need to be able to access the instruction list at any point, instead of just in sequence. In Rust, just add the current instruction pointer as a mutable variable, and change the for each loop to a while let loop:

fn interpret<'a>(program: Program<'a>) {

use Inst::*;



// instantiate stack as an empty Vector

let mut stack = Stack(Vec::new());

let mut pointer = 0;



while let Some(instruction) = program.get(pointer) {

pointer += 1;



match instruction {

...

}

}

}

Next, we add the Jump , JE , and JNE instructions. Jump unconditionally jumps, JE jumps if the top of the stack is equal to zero, and JNE jumps if the top of the stack is not equal to zero. Right now, we'll just specify an index in the instruction list to jump to.

The new Inst enum looks like:

type Pointer = usize;



enum Inst {

Push(Pointer),

...

Jump(Pointer),

JE(Pointer),

JNE(Pointer),

}

And to interpret the new instructions, we add:

fn interpret<'a>(program: Program<'a>) {

use Inst::*;



// instantiate stack as an empty Vector

let mut stack = Stack(Vec::new());

let mut pointer = 0;



while let Some(instruction) = program.get(pointer) {

pointer += 1;



match instruction {

...

Jump(p) => pointer = *p;

JE(p) => {

if stack.peek() == 0 {

pointer = *p;

}

}

JNE(p) => {

if stack.peek() != 0 {

pointer = *p;

}

}

}

}

}

This opens a world of possibilities. We can do loops pretty easily with JNE:

interpret(vec![

Push(10),

Print,

Push(1),

Sub,

JNE(0),

].as_slice());

This prints all the numbers from 10 to 0:

10

9

8

7

6

5

4

3

2

1

The next step is to make it possible to access or change arbitrary stack values. After this the VM will be Turing complete.

Here are the new instructions:

enum Inst {

...

Get(Pointer),

Set(Pointer),

}

Get(p) just indexes the stack at n and copies it to the top of the stack. Set(p) takes the value at the top of the stack and copies it to the stack at index p.

while let Some(instruction) = program.get(pointer) {

pointer += 1;



match instruction {

...

Get(p) => stack.push(*stack.0.get(*p).unwrap()),

Set(p) => {

let v = stack.pop();

*stack.0.get_mut(*p).unwrap() = v;

}

}

}

At this point, since we can Get values from the stack, we want to make our JNE and JE instructions pop the comparison value:

while let Some(instruction) = program.get(pointer) {

pointer += 1;



match instruction {

...

JE(p) => {

if stack.peek() == 0 {

stack.pop();

pointer = *p;

}

}

JNE(p) => {

if stack.peek() != 0 {

stack.pop();

pointer = *p;

}

}

}

}

Using these new instructions, we can write some slightly more complex programs. Here’s a program to sum the first 100 integers:

interpret(vec![

// setup

Push(0), // the accumulator

Push(0), // the index



// loop

// First, add the index to the accumulator

// stack: [accumulator, index]

Get(0),

Get(1),

// stack: [accumulator, index, accumulator, index]

Add,

// stack: [accumulator, index, accumulator + index]

Set(0),

Pop,

// stack: [accumulator + index, index]



// next, increment the index

Push(1), // the increment

// stack: [accumulator, index, 1]

Add,

// stack: [accumulator, index + 1]



// finally, compare the index with 100 and jump back to the start

// if they're not equal.

Get(1),

// stack: [accumulator, index, index]

Push(100),

Sub,

// stack: [accumulator, index, index - 100]

JNE(2),



// if index - 100 == 0, print the accumulator

Get(0),

// stack: [accumulator, index, 0, accumulator]

Print

].as_slice())

Just like that, the VM is Turing complete! It’s super awkward to use though. We can jump to different portions of the code, but since we jump by line number we have to update each jump whenever we change the code. Later, we’ll fix this by writing a super simple compiler that resolves text labels to line numbers.

// find_label takes a line split by spaces and the label it represents,

// or None if it does not represent a label.

fn find_label<'a>(i: Pointer, s: &'a [&'a str]) -> Option<Label> {

if let ["label", l] = s {

Some((l, i))

} else {

None

}

}



let input = ...; // a String



let line_splits = input

.split('

')

.map(|s| s.split_whitespace().collect::<Vec<_>>())

.filter(|s| !matches!(s.as_slice(), [] | ["--", ..]))

.collect::<Vec<_>>();



let labels: HashMap<&str, usize> = line_splits

.iter()

.enumerate()

.filter_map(|(i, s)| find_label(i, s.as_slice()))

.collect();

The following code is super Rusty so I won’t explain it in detail, but essentially it splits each line on space, iterates through them, and constructs a HashMap of label names to line numbers.

Now we need to use this info to actually compile the instructions. This is also pretty Rusty but it’s very concise.

fn parse_instruction(s: &[&str], labels: &Labels) -> Instruction {

use Instruction::*;



match s {

["Push", x] => Push(x.parse::<isize>().unwrap()),

["Pop"] => Pop,

["Add"] => Add,

["Sub"] => Sub,

["Mul"] => Mul,

["Div"] => Div,

["Jump", l] => Jump(*labels.get(l).unwrap()),

["JE", l] => JE(*labels.get(l).unwrap()),

["JNE", l] => JNE(*labels.get(l).unwrap()),

["Get", p] => Get(p.parse::<Pointer>().unwrap()),

["Set", p] => Set(p.parse::<Pointer>().unwrap()),

["Print"] => Print,

["label", ..] => Noop,

l => panic!("Invalid instruction: {:?}", l),

}

}

You might notice that lines starting with label get compiled to a Noop. This is just to make the line numbers easy to keep track of; you could do without it if you sort out all the off by one errors.

we map this function over the list of lines to actually “compile” the text to instructions:

let instructions: Vec<Instruction> = line_splits

.iter()

.map(|s| parse_instruction(s.as_slice(), &labels, &procedures))

.collect();

Now we can easily interpret the compiled instructions with interpret(instructions.as_slice()) . Using this, we can easily rewrite the sum example from before: