Rodrigodd

GameRoy: JIT compilation in High-Accuracy Game Boy Emulation

2023-09-02T18:30:00+00:00

Over the past two years, I have spent a lot of time working on my Game Boy emulator, GameRoy. It has reached a good point, where it has a GUI interface (including a debugger and disassembler) and passes numerous tests (comparable to some of the most accurate emulators). I even made a port to Android!

However, there’s always been something that I wanted to do, even before developing this emulator: implementing a dynamic recompiler, also known as a JIT compiler. Since the first time I researched how an emulator works, one of the initial descriptions I found was about the distinction between interpreters and dynamic recompilers. This made me very interested in the subject.

Reinspired by JIT compilers (such as those in Java, JavaScript and LuaJIT), by Dolphin’s development posts, and primarily by this blog post by Brook Heisler regarding a JIT compiler for NES, I began implementing one in my emulator.

One intriguing point in Brook Heisler’s post was a limitation in how interrupts were handled in the compiled code. When developing his JIT compiler, he encountered a trade-off between performance and precision in how to handle interrupts in his emulator. He ultimately decided to sacrifice performance for precision, but I was not satisfied with that approach.

After pondering for some time, I came up with a solution to this problem - a way to achieve maximum performance without sacrificing precision.

In this blog post, I will describe the process and considerations of implementing a JIT compiler in my emulator, and how I solved the problem of handling interrupts.

It’s worth noting that I am presenting the topics in the same order they came up in my thought process, which just happens to be the reverse order in which they were implemented (excluding the interpreter).

Interpreter x Just-In-Time compilation

An emulator has the job of reproducing a computer system’s behavior in software. A computer system, including the Game Boy, consists of many components, each requiring emulation. Among these, the CPU stands out as the most crucial component, responsible for directly executing the instructions of the programs that you intend to run on your emulator.

There are two main ways of emulating the CPU (or any programmable component).

A very basic interpreter

The simplest way to emulate a system’s CPU is by implementing an interpreter. Essentially, the interpreter reads instructions from the emulated memory, decodes them, and emulates their execution by updating the emulator’s state.

For the Game Boy, the instruction to be executed is read byte by byte (since the Game Boy is an 8-bit system) from the memory address pointed by the program counter (PC), incrementing the PC after each read. To decode the instruction only the first byte is needed for most instructions. Since there are only 255 possible values for a byte, we can use a look-up table. The Game Boy CPU also has 0xCB prefixed instructions that are decoded from the second byte, using another look-up table.

I am using Rust for my emulator implementation, so I can use a match statement to decode the instruction. Each branch of the match lead to a function that implements the instruction’s execution. This function reads and writes to registers or memory, reading additional bytes from the instructions when necessary (e.g., for immediate values, like the address of jump instructions).

A very simplified interpreter would be something like this:

fn interpret_instruction(gb: &mut GameBoy) {
    let opcode = gb.read(gb.cpu.pc);
    gb.cpu.pc += 1;
    match opcode {
        0x00 => nop(gb),
        0x01 => load_bc_im16(gb),
        0x02 => load_bc_a(gb),
        0x03 => inc_bc(gb),
        0x04 => inc_b(gb),
        // 252 more branches like that...
    }
}

// ...

// INC B 1:4 Z 0 H -
fn inc_b(gb: &mut GameBoy) {
    gb.cpu.b += 1;    // increment B
    gb.cpu.f &= 0x1F; // clear Z, N, H flags
    if gb.cpu.b == 0 { gb.cpu.f |= 0x80; }       // set Z flag on zero
    if gb.cpu.b & 0xF == 0 { gb.cpu.f |= 0x80; } // set H flag on half carry
}

The interpret_instruction function is invoked in a loop, executing one instruction per iteration.

One issue with this approach is that the speed at which the interpreter can execute the code is far behind the theoretical speed that your computer’s CPU could execute the code.

The interpreter requires multiple instructions of the host CPU, not only to execute the instruction, but also to read from memory and decode it, task that the host CPU does in a single instruction. Not only that, the interpreter needs to decode the same instruction multiple times.

But that isn’t a problem, at least not for the Game Boy, whose CPU runs at a clock speed of just a little over 1 MHz, while you might be reading this on a device with much more than 1 GHz. In fact, GameRoy can run Game Boy games at over 60 times the original system’s speed on my notebook with a 2.6 GHz Intel i5-4200U.

Even so, this is still 41 times slower than the theoretical speed that my CPU can execute the code! (If you assume the same number of instructions per instruction not very reasonable, actually). Yet, there is a very interesting technique that can be used to close this gap: recompilation.

A very basic JIT compiler

The second way of emulating the CPU is by translating the instructions of the Game Boy CPU into the instructions of the host CPU and then executing them. This method is called recompilation. In the case of the Game Boy, as well as in many other systems, you can not recompile all the code contained in a program beforehand (known as static recompilation). Therefore, you need to perform that translation at runtime, which makes the technique be called dynamic recompilation, or more commonly called, Just-In-Time compilation (or JIT).

To achieve this, you need to create a compiler. The compilation process is very similar to the interpreter, where you read, decode and execute instructions. But instead of executing the instruction directly, it will generate the corresponding machine code that will.

The machine code can be generated just once and then executed multiple times, bypassing the overhead of reading and decoding instructions of the interpreter.

In a simplified form, the code for the compiler would look something like this:

fn compile_instruction(ctx: &mut JitCompiler, gb: &GameBoy, code: &mut Vec<u8>) {
    let opcode = gb.read(ctx.curr_pc);
    match opcode {
        0x00 => nop(ctx, gb),
        0x01 => load_bc_im16(ctx, gb),
        0x02 => load_bc_a(ctx, gb),
        0x03 => inc_bc(ctx, gb),
        0x04 => inc_b(ctx, gb),
        // 252 more arms like that...
    }
}

// ...

// INC B 1:4 Z 0 H -
fn inc_b(_: &mut JitCompiler, code: &mut Vec<u8>) {
    code.extend(&[
        0x0f, 0xb6, 0x47, 0x01, // movzx  eax,BYTE PTR [rdi+0x1] ; read F
        0x0f, 0xb6, 0x4f, 0x02, // movzx  ecx,BYTE PTR [rdi+0x2] ; read B
        0xfe, 0xc1,             // inc    cl                     ; increase B
        0x88, 0x4f, 0x02,       // mov    BYTE PTR [rdi+0x2],cl  ; write B
        0x24, 0x1f,             // and    al,0x1f                ; clear Z, N, H flags
        0x8d, 0x50, 0x80,       // lea    edx,[rax-0x80]         ; calculate the flags...
        0x0f, 0xb6, 0xc0,       // movzx  eax,al                 ; ...
        0x0f, 0xb6, 0xd2,       // movzx  edx,dl
        0x0f, 0x45, 0xd0,       // cmovne edx,eax
        0x8d, 0x42, 0x20,       // lea    eax,[rdx+0x20]
        0xf6, 0xc1, 0x0f,       // test   cl,0xf
        0x0f, 0xb6, 0xc0,       // movzx  eax,al
        0x0f, 0x45, 0xc2,       // cmovne eax,edx
        0x88, 0x47, 0x01,       // mov    BYTE PTR [rdi+0x1],al  ; write F
    ]);
}

This compile_instruction function reads and decodes an instruction, and then writes the machine code that executes that instruction to a buffer.

(For the curious, I got that assembly/machine code by simply implementing the instruction in Rust, and then disassembling it with objdump -d.)

So, to create a JIT compiler, you only need to run compile_instruction on the sequence of instructions starting at the address the PC is currently, write the code to a memory page with executable permission (you can use a crate like memmap2 for that), transmute it into a function pointer, and call it! Simple as that!

Well, actually, some instructions that involve function calls in their implementation or include branching are much trickier to implement than just appending a static array of bytes to a buffer. You also need to make sure that the code for each instruction can work together, using the same registers and so on. You also need a prologue and epilogue for the function that satisfies a calling convention.

However, this already shows the basics of a JIT compiler.

Also, in my actual implementation I used the crate dynasm for emitting the machine code (this crate has a macro that translates assembly code into machine code at compile time), so the code is much more maintainable than the array of raw bytes that I presented here.

Optimizations

The cool thing about compiling code is that you can do optimizations! For example, one of the (few) optimizations that I made in the JIT compiler for my emulator is the omission of unnecessary flag calculations (directly inspired by the one described in Brook Heisler’s post).

If you tried to read what the assembly code is doing, you may have noticed that only four instructions are used to increment the value of register B, and the remaining 10 are being used to update the conditional flags in the register F. But the value of F may be overwritten by the next instructions before it could be read from a conditional jump instruction for example. In that scenario, updating the value of F does not affect the behavior of the emulation, so it is a waste to compute.

If we compile all these instructions together, we can detect that kind of scenario and omit the flag calculation, greatly reducing the amount of code emitted!

// INC B 1:4 Z 0 H -
fn inc_b(ctx: &mut JitCompiler, code: &mut Vec<u8>) {
    if !ctx.flag_is_need {
        code.extend(&[0xfe, 0x47, 0x02]); // inc BYTE PTR [rdi+0x2]
        return;
    }

    code.extend(&[
        0x0f, 0xb6, 0x47, 0x01, // movzx  eax,BYTE PTR [rdi+0x1] ; read F
        0x0f, 0xb6, 0x4f, 0x02, // movzx  ecx,BYTE PTR [rdi+0x2] ; read B
        0xfe, 0xc1,             // inc    cl                     ; increase B
        0x88, 0x4f, 0x02,       // mov    BYTE PTR [rdi+0x2],cl  ; write B
        0x24, 0x1f,             // and    al,0x1f                ; clear Z, N, H flags
        0x8d, 0x50, 0x80,       // lea    edx,[rax-0x80]         ; calculate the flags...
        0x0f, 0xb6, 0xc0,       // movzx  eax,al                 ; ...
        0x0f, 0xb6, 0xd2,       // movzx  edx,dl
        0x0f, 0x45, 0xd0,       // cmovne edx,eax
        0x8d, 0x42, 0x20,       // lea    eax,[rdx+0x20]
        0xf6, 0xc1, 0x0f,       // test   cl,0xf
        0x0f, 0xb6, 0xc0,       // movzx  eax,al
        0x0f, 0x45, 0xc2,       // cmovne eax,edx
        0x88, 0x47, 0x01,       // mov    BYTE PTR [rdi+0x1],al  ; write F
    ]);
}

Fantastic! The optimized version is 14 times smaller than the unoptimized one!

There are, of course, many other optimizations that can be done. These could reach the point where the optimizations become by far the most complex part of a compiler. Instead of implementing them your-self, you could use a third party compiler infrastructure, like LLVM or Cranelift, which implement many optimizations for you.

But my emulator is, at least for now, only targeting x86-64 machine code. So I will refrain from using any of them.

But there is a big problem in trying to apply any non-trivial optimization, even the simple flag omission one: the fact that the CPU needs to handle interrupts.

Interrupts

Interrupts are a mechanism that allows the CPU to stop the execution of the current code and start executing a function that handles the interrupt. For example, in the Game Boy, you can use the interrupt triggered by the VBlank of the PPU (which occurs right after the screen is fully rendered) to update the game logic and prepare the next frame for rendering.

There are multiple sources of interrupts in the Game Boy, which are triggered by one of the various components. So to handle the interrupts we also need to update them.

And if you want a precise enough emulator, you need to update the components and check for interrupts before every single instruction. Not only that, if you pursue cycle-accuracy, you also need to update the components before each memory access. One way of doing that is by using a tick function.

Our interpreter becomes something like this:

fn interpret_instruction(gb: &mut GameBoy) {
    if gb.interrupts_enabled & gb.interrupts_flags != 0 {
        // this will write to the stack, and change PC to the interrupts handler
        handle_interrupt(gb);
    }

    let opcode = gb.read(gb.cpu.pc);
    gb.cpu.pc += 1;
    tick(gb, 4); // each memory access takes 4 cycles
    match opcode {
        // ...
        0x36 => ld_hl_im8(gb),
    }
}

fn tick(gb: &mut GameBoy, cycles: u64) {
    gb.clock_count += cycles;

    gb.update_timer(clock_count);
    gb.update_ppu(clock_count);
    gb.update_sound_controller(clock_count);
    gb.update_serial(clock_count);
}

// LD (HL),d8 2:12 - - - -
fn ld_hl_im8(gb: &mut GameBoy) {
    let value = gb.read(gb.cpu.pc);
    gb.cpu.pc += 1;
    tick(gb, 4);

    let address = gb.cpu.h as u16 * 0x100 + gb.cpu.l as u16;
    gb.write(address, value);
    tick(gb, 4);
}

Ignoring all the complexity of emulating each component and some quirks with interrupt handling timing, this is roughly how the interpreter handles the interrupts.

However, handling interrupts poses a problem for our JIT compiler. Not only might constantly checking for interrupts introduce a considerable amount of overhead, but we can no longer apply many of the optimizations that we would like to do, like omitting calculation, because the code needs to always be prepared to branch to the interrupt handler.

You could try to check for interrupts only every 10 instructions or something; it may be precise enough for most games (but not all). But my emulator pursues maximum accuracy, and the point of the JIT compiler is to determine if it’s possible to implement it, without sacrificing precision or performance.

And I have figured out a way to achieve that: by estimating when the next interrupt will happen.

Estimating the next interrupt

The JIT compiler in GameRoy works this way: given the current address in the PC register and the bank currently mapped to that address, the compiler produces a block of instructions starting at that address. This block is then compiled down into machine code, and stored in a cache.

Before executing the block, it checks if an interrupt will happen during its execution. If not, the block is executed, otherwise, we fall back to the interpreter.

This allows us to apply all the optimizations that we want, without worrying about interrupts (for the most part).

The disadvantage is that we still need to have an interpreter around, and the fallback may decrease the emulator’s performance. However, I already had an interpreter working, so there’s no problem with that. Only a small fraction of the instructions need to be interpreted, so the loss of performance is not so bad. An alternative is to use, as a fallback, a version of the compiler that handles interrupts.

The main challenge in this approach is estimating when the next interrupt will happen. Depending on how a component generates interrupts, it may just not be possible to estimate when the next interrupt will happen more efficiently than just emulating the component until it occurs.

Thankfully, in the case of the components of the Game Boy, it is possible to estimate the instant of the next interrupt in constant time. But implementing the estimation is still tricky, with a lot of edge cases. Thanks to a lot of fuzzing-based testing, I am somewhat confident that I got it right.

If you want to take a look at how the implementation looks like, here are links for the timer, for the PPU, and for the serial. The fuzz tests are at the bottom of the same file, if you also want to take a look at them.

Lazy Update

And with that we fix the problem with interrupts! But we still need to update the components. But thankfully we only need to update the components before accessing a memory address that is mapped to a component or right before an interrupt is triggered. So we can just update the components in the memory access function, right before reading or writing to memory, and before checking for interrupts in the interpreter.

To implement that, each component will contain a clock count of the last time it was updated, and a function that updates the component to the current clock count. The cool part is that this allows us to further optimize the emulation of each component. By decreasing the number of times we update each component we save some work, and emulating many cycles at once allows more efficient emulation.

The implementation of the PPU is a good example of that. The PPU is the component responsible for transforming the tile map and sprite data into pixels for the screen. The final rendering result is somewhat straightforward to emulate if the input data doesn’t change. But the necessity of taking into account that the CPU may be changing the input during rendering makes it necessary to emulate the complex inner workings of the PPU.

But when we lazily update the PPU we can be sure that there were no changes since the last update, allowing us to bypass most of the complexity of the PPU, just emulating the final rendering result, which is much faster.

In the figure below we can view the time spent emulating the PPU, when updating it every instruction and when updating only when necessary.

The draw_scan_line function renders the final result of a scan line, without emulating the entire PPU’s inner workings. Doing the math, it is possible to see that the emulation of the PPU gains a speed-up of almost 10 times, increasing the overall speed of the emulator by 4.

One thing to notice is that before this optimization, only 14% of the time was spent in the CPU. That means the CPU emulation was hardly the bottleneck of the emulator, and the JIT compiler would not provide a significant speed-up. But after this optimization, the CPU emulation is responsible for 63% of the time, giving some room for the JIT compiler to shine.

The Compiler a little more in depth

Above, I’ve just provided a simplified version of how a JIT compiler would work. However, the actual implementation concerns itself with more details.

First, it needs to decide which group of instructions to compile. For that, I have a step for “tracing a block”, which basically starts at a given address (very likely the current address in the PC register) and decodes instructions one after the other, until it reaches an unconditional branch (like a jump or a call). This produces a block for that address.

Currently, the blocks are just linear sequences of instructions, but I could very well also expand the block tracing to follow branches and potentially compile big sections of the program at once. However, my block caching is still too suboptimal for that (I explain it in the “Remaining Work” section).

An important point is that I only trace blocks when the PC is pointing to an address in the memory ROM (range 0x0000 to 0xFFFF). If it is pointing outside of it, like the RAM, I just fall back to the interpreter. This allows me to avoid handling self-modifying code, which is a general problem for dynamic recompilers. Also, games in general don’t run much code in RAM, so it’s not that big of a loss.

Another point is that the Game Boy uses bank switching, which means that different memory banks may be mapped into the same addresses at different times. So, a block needs to be identified not only by its starting address but also by a bank number.

After tracing a block, the compiler compiles it down to a function. It emits machine code for the function prologue (align the stack, push registers, etc.), and then emits the machine code for each instruction.

Instructions that update registers work like the one I gave as an example before. They just load the registers from memory, do the operation, and store them back. A much more efficient way would be to keep the Game Boy registers in the x86-64 registers for as long as possible, but I still haven’t implemented that.

Instructions that may branch check if the target address is inside the block, and if it is, they emit a jump to the target address inside the block. This uses dynasm helper for emitting jumps. If the target isn’t in the block, it just emits the prologue for the block (pop registers and call ret).

Instructions that read or write to memory emit function calls to Rust functions that handle the memory access. This allows me to handle cases where it may change the emulator’s state, like the lazy update of components. But memory access with an immediate address may be emitted inline if possible.

Another important point is that writing to memory may trigger an interrupt, or at least change its timing. This means that after each write, the compiler needs to emit another check for interrupts, exiting if necessary. The same is true for bank switching: a write to ROM could switch the bank that the block is currently executing, so it needs to check if it did change and exit the block.

Also, because writes can trigger interrupts and the interrupts can read the content of flags, each write also compromises inter-instruction optimizations, like the omission of flag calculations that I mentioned earlier.

And because there are now multiple checks for interrupts throughout the block, each check only needs to make sure that the next interrupt will not happen before the next interrupt check, including the initial check that happens before entering the block.

After finishing emitting the block, the compiler moves the buffer with the machine code to an executable memory region and stores it in a HashMap, which maps (bank, address) to the block. Now, whenever the emulator needs to execute this block again, it just gets it from the HashMap (or compiles it if it isn’t there yet).

Results

The entire point of implementing a JIT compiler is to improve the performance of the emulator. So, how much of a speed-up did I get?

First, let’s compare the speed of the interpreter and the JIT compiler. For that, I ran the emulator with two different games and measured the time it took to run each game for 100 seconds. The first game is The Legend of Zelda: Link’s Awakening, and the second is Tobu Tobu Girl. Both games have a nice start screen animation, which may be a good benchmark for the speed of the emulator.

The emulator was capable of emulating Tobu Tobu Girl more than 100 times faster than the real system! However, when comparing the two implementations, it’s possible to see that the JIT compiler decreases the execution time by about 40% in Tobu Tobu Girl and about 30% in Zelda. This is not the massive speed-up I was expecting, but it is still very significant.

We can also take a closer look at the time spent emulating each component. By using blondie or inferno to capture a stack trace and passing it to my plot script, I can generate the following pie charts:

From the pie chart above, we can see that a big chunk of the time is spent emulating the PPU. We can also see that Zelda has a lot less opportunities for optimizing the PPU, and therefore spends more time in it, which explains why it is much slower and has a smaller speed up.

Looking just at the interpreter (blue) and JIT (red) sections, we can roughly estimate that the JIT compiler is about 4 times faster than the interpreter. However, a considerable amount of time is still spent just querying the block in the HashMap. And even with the JIT compiler, some time is still spent in the interpreter.

Another aspect we can explore is comparing the emulator with others out there. I picked 5 emulators that are relatively popular and have a high degree of accuracy: SameBoy, BGB, Beaten Dying Moon, Emulicious and Gambatte-Speedrun.

You can get a feel for how accurate each emulator is by viewing the test results in daid’s GBEmulatorShootout. Notice that these tests include DMG, GBC, and SGB versions of the Game Boy, but GameRoy only emulates DMG, so it appears a little behind the others.

The emulators were run on Windows, with their default GUI frontend, configured to emulate DMG and with “turbo mode” or “fast forward” or whatever it is called enabled. Due to the use of a GUI frontend, the measurement being based on video recording, and the intro being run only two or three times, the results are not that precise. But hopefully they are good enough to compare them with each other.

To measure the performance of each emulator, I ran the same two games as before, and measured the time it took to run the intro animation of each game, using a video recorded with OBS Studio. For Tobu Tobu Girl, I let the intro run three times, and for Zelda, I let it run twice.

An important point is that both BGB and Gambatte (at least, I’m not sure if the others do something similar) have optimizations that skip the rendering of some frames of the LCD, which is something that I haven’t implemented yet in my emulator. That optimization can be disabled in BGB, so I included it in the benchmarks to have a more direct performance comparison. For Gambatte it is not possible to disable it¹.

Here are the results²:

Based on these results (and if you ignore the emulators that do frameskip optimization and the lack of rigor in the benchmark), I can somewhat claim that GameRoy is the fastest emulator!

Interestingly, even the interpreted version of my emulator is faster than most of the other emulators when running Tobu Tobu Girl. This is likely due to the PPU optimization achieved by implementing lazy updates and estimating the next interrupt, as The Legend of Zelda emulation speed is still comparable to the other emulators and it less affected by PPU optimization.

But both Gambatte and BGB with frameskip optimization are a lot faster than GameRoy. BGB is more than 5 times faster than its non-optimized version (which is closer to the 4 times speed-up that I got with PPU optimization in the Interpreter version of GameRoy) and is more than 2 times faster than GameRoy.

If I implement the frameskip optimization, somehow eliminating the entire time spent in the PPU, my emulator would become about 2 times faster than it is now, which still would not be enough to beat BGB. So the claim that GameRoy is the fastest emulator may be a little exaggerated.

Errata: The emulator benchmarks included in the first published version of this blog post were using an older version of GameRoy, didn’t include Gambatte, and were running BGB with its maximum speed capped at 10x. In that scenario, I had claimed that GameRoy was the fastest emulator by far. I hope I have fixed it now.

Remaining work

I may have achieved the state-of-the-art regarding dynamic recompilation applied to Game Boy emulation (the only related work that I know of is JitBoy), but my implementation is still far from perfect.

The first thing is that I am not implementing some obvious optimizations. For example, I am not keeping the Game Boy registers loaded in the x64 registers, but instead, I am loading and saving them to memory every instruction. The only register that I optimize is the PC register, whose value I can know at compile time, so I only update it on block exit.

The flags are the most suboptimal, as I am computing the value of the flag, encoding it in the F register bit flag, and then later decoding them in the branch instructions, where I could be using the value of each flag directly between instructions.

But I will probably not implement these optimizations, at least not directly. Instead, I plan to write another backend for the compiler using Cranelift, which will likely automatically handle these optimizations and more. I have already done some experiments with it before, and it seems to work pretty well. It will also help with porting the JIT compiler to run on Android.

Another inefficiency in my JIT compiler is regarding redundant recompilations. Currently, I am building an entire block for each (bank, address) entry point, even if that same address is already in the middle of another block. This means that many of the blocks are overlapping, introducing unnecessary compilation time and memory usage.

I don’t know if this is a big problem (I need to do some profiling first), but I could fix it by including in the HashMap a key for the address of each instruction in the block and moving the block prelude to a Rust naked function, that jumps to the middle of the block.

However, this would introduce some complications, like filling the HashMap with too many keys, and the need to ensure that each compiled instruction can be jumped to. This last point may be a problem for optimizations, but I haven’t given it too much thought yet.

Conclusion

And that’s it! I was not able to achieve the massive speed-up I was expecting, but I still managed to make my emulator the fastest one out there, even if the generated code is far from optimum.

I would like to thank Brook Heisler for his blog post, “Experiments In NES JIT Compilation”, which inspired me to try this out, and gave me some useful links on how to implement a JIT compiler, like Eli Bendersky’s Adventures In JIT Compilation series, which I re-implemented in Rust as a learning exercise.

If you want to take a look at GameRoy, I invite you to check out its GitHub repository.

Another point is that the frameskip in Gambatte, unlike BGB, is not dynamic. You need to provide a fixed number of frames to skip, and the unskipped frames are still emulated at a maximum of 60 frames per second. I configured the frameskip to 75³, which made the GUI lag a littke (hopefully, this means it is running at maximum speed), but be aware that this choice is arbitrary and may not be optimal. ↩
You can also see the raw data collected here. ↩
The emulator has a hard cap of 16 in the frameskip configuration, but I modified the source code to disable that limit. ↩

Compiling Brainfuck code - Part 4: A Static Compiler

2022-12-31T19:50:00+00:00

This is the fourth post of a blog post series where I reproduce Eli Bendersky’s Adventures In JIT Compilation series, but this time using the Rust programming language.

In this final part we will use the compiled code that we generated in the previous parts, and learn how to put them in an executable file, achieving a static compiler.

Executable files and Object Files

To build a static compiler, we need to be able to save our generated assembly in a file, in a way that the operating system can execute.

Each OS have its own executable file format. On Windows, executable files are in the PE file format, on macOS in the Mach-O. On Linux, executable files are in the ELF file format, which we will be exploring a little more in deep here.

Another important piece of the compilation process are the object files, and the process of linking. Object files contain compiled code, but only part of the code that forms a program. They are used by compilers to split the source code compilation (C compilers can compile each .c file to a separated object file), or for distributing precompiled libraries, for example. Again, each OS has its own object file format, which commonly uses the same format as executables, or executables use an extended version of them. On Linux, objects are also ELF files.

To link all these object files into a complete and executable program, you use a linker. Each object file has a list of symbols, that, among other things, are used to indicate what function the file define, and which ones they are using (that may be defined in another file). It also has a list of relocations, that are used to indicate which points in the memory need to be modified, and how, to connect the object files. The linker is the program that read these files and does all these modifications, and outputs an executable.

We will see more practical examples of this further down.

Exploring ELF files

As always, we will start by making a compiler for the simplest program that we can achieve. In the previous parts we were using the add1 function as example, but that don’t make into an interesting enough program. So let’s make a good old “hello world” program.

On Linux x86-64 assembly, it would be:

section .data
    hello:    db 'Hello world!',10
    helloLen: equ $-hello

section .text
    global _start 

_start:
    mov rax, 1           ; 'write' system call = 1
    mov rdi, 1           ; file descriptor 1 = STDOUT
    mov rsi, hello       ; string to write
    mov rdx, helloLen    ; length of string to write
    syscall              ; call the kernel

    ; Terminate program
    mov rax, 60          ; 'exit' system call = 60
    mov rdi, 0           ; exit with error code 0
    syscall              ; call the kernel

In the code above we have two sections. Each section holds a chunk of arbitrary data. The sections can have any name, but the linker will map the section with the name .data to a segment that will have read and write permissions when executed, and the .text to a segment with read and execute permissions.

In the section .data we declared the string 'Hello World!\n' (the 10 there is the ASCII code for line-break) with the label hello, and a constant for the string length (compute by taking the current address ($) minus the hello address).

And in the section .text we declare that the label _start is global (meaning that other objects can see it), and declaring it preceding the code for printing the string (using the ‘write’ syscall) and exiting (with the ‘exit’ syscall).

Every label above will map to a symbol in the object file, that later will be used for applying relocations, debug purposes, etcetera. The _start symbol in specific will be used by the linker to tell where is the entry point of our executable (and that is why it must be global).

To assemble and link the program, we can use the assembler NASM and the linker ld:

$ nasm hello.as -f elf64 -o hello.o
$ ld hello.o -o hello

And running it:

$ ./hello
Hello World!

It works!

Now, let’s take a close look in the object file and the executable file. For this, we can use the readelf utility. You can use the -a flag to show in one go all the information that I will be explaining, but I will go in parts here.

First, lake take a look at the hello.o object file:

$ readelf hello.o -h
ELF Header:
  Magic:   7f 45 4c 46 02 01 01 00 00 00 00 00 00 00 00 00
  Class:                             ELF64
  Data:                              2's complement, little endian
  Version:                           1 (current)
  OS/ABI:                            UNIX - System V
  ABI Version:                       0
  Type:                              REL (Relocatable file)
  Machine:                           Advanced Micro Devices X86-64
  Version:                           0x1
  Entry point address:               0x0
  Start of program headers:          0 (bytes into file)
  Start of section headers:          64 (bytes into file)
  Flags:                             0x0
  Size of this header:               64 (bytes)
  Size of program headers:           0 (bytes)
  Number of program headers:         0
  Size of section headers:           64 (bytes)
  Number of section headers:         7
  Section header string table index: 3

First we have the ELF header, with basic information like version, ABI, etc. It also contains the location in the file of the program headers (only relevant for executable files) and the section headers.

You may want to take a look on the man page for the ELF format to see in more detail what each field in the header means, and also for the parts below.

$ readelf hello.o -S
There are 7 section headers, starting at offset 0x40:

Section Headers:
  [Nr] Name              Type             Address           Offset
       Size              EntSize          Flags  Link  Info  Align
  [ 0]                   NULL             0000000000000000  00000000
       0000000000000000  0000000000000000           0     0     0
  [ 1] .data             PROGBITS         0000000000000000  00000200
       000000000000000d  0000000000000000  WA       0     0     4
  [ 2] .text             PROGBITS         0000000000000000  00000210
       0000000000000027  0000000000000000  AX       0     0     16
  [ 3] .shstrtab         STRTAB           0000000000000000  00000240
       0000000000000032  0000000000000000           0     0     1
  [ 4] .symtab           SYMTAB           0000000000000000  00000280
       00000000000000a8  0000000000000018           5     6     8
  [ 5] .strtab           STRTAB           0000000000000000  00000330
       0000000000000020  0000000000000000           0     0     1
  [ 6] .rela.text        RELA             0000000000000000  00000350
       0000000000000018  0000000000000018           4     2     8
Key to Flags:
  W (write), A (alloc), X (execute), M (merge), S (strings), I (info),
  L (link order), O (extra OS processing required), G (group), T (TLS),
  C (compressed), x (unknown), o (OS specific), E (exclude),
  l (large), p (processor specific)

Following it, we have the section headers. The first header is always a null one. Here we can see that we have the headers for the .data and .text sections that we have written in our assembly code. But we also have some other sections that hold, respectively, the strings for the name of the sections headers, the symbols, the string for the symbol names and such, and lastly the relocations for the .text section.

Each section header tells some information about the section, like where its content is in the file. It may also be linked to another section (like .symtab has a link to the table that contains the symbol names).

The readelf can pretty print the contents of some sections, as we will see in a moment, but you can also use -x to view the hex dump:

$ readelf hello.o -x .shstrtab

Hex dump of section '.shstrtab':
  0x00000000 002e6461 7461002e 74657874 002e7368 ..data..text..sh
  0x00000010 73747274 6162002e 73796d74 6162002e strtab..symtab..
  0x00000020 73747274 6162002e 72656c61 2e746578 strtab..rela.tex
  0x00000030 7400                                t.

Here for example we can see that the section “section header string table” is a list of null byte terminated strings, with the first string being empty.

$ readelf hello.o -s

Symbol table '.symtab' contains 7 entries:
   Num:    Value          Size Type    Bind   Vis      Ndx Name
0000000000000000     0 NOTYPE  LOCAL  DEFAULT  UND
0000000000000000     0 FILE    LOCAL  DEFAULT  ABS hello.as
0000000000000000     0 SECTION LOCAL  DEFAULT    1
0000000000000000     0 SECTION LOCAL  DEFAULT    2
0000000000000000     0 NOTYPE  LOCAL  DEFAULT    1 hello
000000000000000d     0 NOTYPE  LOCAL  DEFAULT  ABS helloLen
0000000000000000     0 NOTYPE  GLOBAL DEFAULT    2 _start

With -s we can see the symbol table. The first one is always null. There are the symbols that we defined in our assembly, and some others introduced by the assembler like the file name and symbols for each section (used for relocation).

Each symbol contains a value. In the case of hello and start it is the address of the symbol in its respective section (both are at the start of the section, so they are 0) and in helloLen is the length of the string.

$ readelf hello.o -r

Relocation section '.rela.text' at offset 0x350 contains 1 entry:
  Offset          Info           Type           Sym. Value    Sym. Name + Addend
00000000000c  000200000001 R_X86_64_64       0000000000000000 .data + 0

And lastly we have the relocations. Here we have only a single one for the .text section that applies at offset 0xc, use the symbol 2 (the first 2 bytes of info) on the linked section (.symtab), is of type R_X86_64_64 (each processor has its own set of relocation types) and has an addend of 0. This means that it replaces a 64bit word at 0xc by the value of the symbol .data.

If we look at the disassembly of the machine code, we can see that in offset 0xc we have the last 8 bytes of the movabs instruction, the one that loads the address of the hello string:

$ objdump -d -M intel hello.o

hello.o:     file format elf64-x86-64


Disassembly of section .text:

0000000000000000 <_start>:
   0:   b8 01 00 00 00          mov    eax,0x1
   5:   bf 01 00 00 00          mov    edi,0x1
   a:   48 be 00 00 00 00 00    movabs rsi,0x0
  11:   00 00 00
  14:   ba 0d 00 00 00          mov    edx,0xd
  19:   0f 05                   syscall
  1b:   b8 3c 00 00 00          mov    eax,0x3c
  20:   bf 00 00 00 00          mov    edi,0x0
  25:   0f 05                   syscall

This means that during linking, after the linker have set the location in memory that the data section should be loaded, it will path that instruction to the correct address of the hello string.

And we can see all this in action, by inspecting the ./hello executable file:

$ readelf hello -S
There are 6 section headers, starting at offset 0x2158:

Section Headers:
  [Nr] Name              Type             Address           Offset
       Size              EntSize          Flags  Link  Info  Align
  [ 0]                   NULL             0000000000000000  00000000
       0000000000000000  0000000000000000           0     0     0
  [ 1] .text             PROGBITS         0000000000401000  00001000
       0000000000000027  0000000000000000  AX       0     0     16
  [ 2] .data             PROGBITS         0000000000402000  00002000
       000000000000000d  0000000000000000  WA       0     0     4
  [ 3] .symtab           SYMTAB           0000000000000000  00002010
       00000000000000f0  0000000000000018           4     6     8
  [ 4] .strtab           STRTAB           0000000000000000  00002100
       0000000000000031  0000000000000000           0     0     1
  [ 5] .shstrtab         STRTAB           0000000000000000  00002131
       0000000000000027  0000000000000000           0     0     1
Key to Flags:
  W (write), A (alloc), X (execute), M (merge), S (strings), I (info),
  L (link order), O (extra OS processing required), G (group), T (TLS),
  C (compressed), x (unknown), o (OS specific), E (exclude),
  l (large), p (processor specific)

Here we can see that the address of the .text and .data sections have been set up to 0x401000 and 0x402000. Now if we take a look in the disassembly:

$ objdump -d -M intel hello

hello:     file format elf64-x86-64


Disassembly of section .text:

0000000000401000 <_start>:
  401000:       b8 01 00 00 00          mov    eax,0x1
  401005:       bf 01 00 00 00          mov    edi,0x1
  40100a:       48 be 00 20 40 00 00    movabs rsi,0x402000
  401011:       00 00 00
  401014:       ba 0d 00 00 00          mov    edx,0xd
  401019:       0f 05                   syscall
  40101b:       b8 3c 00 00 00          mov    eax,0x3c
  401020:       bf 00 00 00 00          mov    edi,0x0
  401025:       0f 05                   syscall

Look! The address in the movabs instruction, at offset 0xc was updated to 0x402000 (in little endian)!.

Another thing you may notice if you run readelf hello -a is that it now contains a program header table:

$ readelf hello -l

Elf file type is EXEC (Executable file)
Entry point 0x401000
There are 3 program headers, starting at offset 64

Program Headers:
  Type           Offset             VirtAddr           PhysAddr
                 FileSiz            MemSiz              Flags  Align
  LOAD           0x0000000000000000 0x0000000000400000 0x0000000000400000
                 0x00000000000000e8 0x00000000000000e8  R      0x1000
  LOAD           0x0000000000001000 0x0000000000401000 0x0000000000401000
                 0x0000000000000027 0x0000000000000027  R E    0x1000
  LOAD           0x0000000000002000 0x0000000000402000 0x0000000000402000
                 0x000000000000000d 0x000000000000000d  RW     0x1000

 Section to Segment mapping:
  Segment Sections...
   00
   01     .text
   02     .data

A program header contains information about things that the system need to prepared for the program execution. Here we have three headers of type LOAD that tells which segments of the file need to be loaded in memory.

The first one tells to load 232 (0xe8) bytes starting at offset 0 of the file (that is, the ELF header and program headers) with read permissions in virtual address 0x40_0000 (or physical address 0x40_0000 but most PCs today uses virtual memory). (I am not sure why the ELF header is loaded in memory.)

The second one loads the .text section, with read and exec permissions, and the third loads .data with read and write permissions.

The program headers themselves don’t say which sections they are mapping (although readelf shows the mapping), but you can check this yourself by comparing the offset and address of the segments here, with the ones in the section header of each section.

A minimal hello world

Again, it’s always good to start experimenting with an idea by implementing the simplest example that we can make. In our case that would be compiling a hello world program. But the one that I showed is still not minimal enough.

Because our 'Hello World!\n' string is only read during the execution, it does not need to live in a section with read and write permission like the .data. We could put it in a .rodata, which is read only, but we can take advantage that the .text segment has read permission and put the string there, saving an entire section:

section .text
    global _start 

_start:
    mov rax, 1           ; 'write' system call = 1
    mov rdi, 1           ; file descriptor 1 = STDOUT
    lea rsi, [rel hello] ; string to write
    mov rdx, helloLen    ; length of string to write
    syscall              ; call the kernel

    ; Terminate program
    mov rax, 60          ; 'exit' system call = 60
    mov rdi, 0           ; exit with error code 0
    syscall              ; call the kernel
    hello:    db 'Hello world!',10
    helloLen: equ $-hello

And by putting the string alongside the code, we can make sure the string is always at a constant offset from the code. This means that I could replace the mov of the absolute address of hello in the code, by a lea rsi, [rel hello] rip relative addressing, saving a relocation!

If you take a look at the disassembly now, you will see that that instruction is compiled down to lea rsi, [rip+0x13]:

$ nasm hello.as -f elf64 -o hello.o
$ objdump -d -M intel hello.o

hello.o:     file format elf64-x86-64


Disassembly of section .text:

0000000000000000 <_start>:
   0:   b8 01 00 00 00          mov    eax,0x1
   5:   bf 01 00 00 00          mov    edi,0x1
   a:   48 8d 35 13 00 00 00    lea    rsi,[rip+0x13]        # 24 <hello>
  11:   ba 0d 00 00 00          mov    edx,0xd
  16:   0f 05                   syscall
  18:   b8 3c 00 00 00          mov    eax,0x3c
  1d:   bf 00 00 00 00          mov    edi,0x0
  22:   0f 05                   syscall

0000000000000024 <hello>:
  24:   48                      rex.W
  25:   65 6c                   gs ins BYTE PTR es:[rdi],dx
  27:   6c                      ins    BYTE PTR es:[rdi],dx
  28:   6f                      outs   dx,DWORD PTR ds:[rsi]
  29:   20 77 6f                and    BYTE PTR [rdi+0x6f],dh
  2c:   72 6c                   jb     9a <hello+0x76>
  2e:   64 21 0a                and    DWORD PTR fs:[rdx],ecx

(Also, these strange instruction are there only because objdump is trying to interpret the string data as code now.)

The object crate

Now we can finally start making a small compiler. First, lets use dynasm, like we did in part two, to get our hello world program in Rust:

let mut code: VecAssembler<X64Relocation> = VecAssembler::new(0);

let hello_str = b"Hello world!\n";
let len = hello_str.len() as i32;
dynasm!(code
    ; mov eax,1            // 'write' system call = 1
    ; mov edi,1            // file descriptor 1 = STDOUT
    ; lea rsi, [>hello]    // string to write
    ; mov edx, DWORD len   // length of string to write
    ; syscall              // call the kernel

    // Terminate program
    ; mov eax,60           // 'exit' system call
    ; mov edi,0            // exit with error code 0
    ; syscall              // call the kernel
    ; hello:
    ; .bytes hello_str
);

let code = code.finalize().unwrap();

We can directly run that code using memmap2, to make sure everything is working:

let mut buffer = memmap2::MmapOptions::new()
    .len(code.len())
    .map_anon()
    .unwrap();

buffer.copy_from_slice(code.as_slice());

let buffer = buffer.make_exec().unwrap();

let hello: unsafe extern "C" fn() -> ! = std::mem::transmute(buffer.as_ptr());
hello()

If you run this, it will print Hello World!.

But we should be compiling it to an object file. For this we will use the object crate. As its docs says, “the object crate provides a unified interface to working with objects files across platforms”.

This crate will allow us to write our object files, without needing to worry about the details of each object format, making it easier to make our compiler cross-platform. It also allows writing ELF and PE executables directly, but we would need to use its lower level API for each format, so we will not cover that here.

So first we crate a Object struct with the specification of our host:

let mut obj = object::write::Object::new(
    object::BinaryFormat::Elf,
    object::Architecture::X86_64,
    object::Endianness::Little,
);

Next we can add the _start symbol to the object:

let start = obj.add_symbol(Symbol {
    name: b"_start".to_vec(),
    kind: object::SymbolKind::Text,
    scope: object::SymbolScope::Linkage,
    weak: false,
    flags: SymbolFlags::None,

    value: 0,
    size: 0,
    section: object::write::SymbolSection::Undefined,
});

We set its name; its kind, text, because it points to executable code; its scope, linkage, which make the symbol global, like we did in the assembly; it doesn’t need to be weak (weak symbols can be overridden by non-weak symbols); and it won’t need any special flags.

Value, size and section will be updated when we add data to the symbol, so I put some default values there for now.

Now we can crate the .text section, and by calling add_symbol_data we append our code to it and at same time we update the _start symbol. x86 instructions don’t need to be aligned, so the alignment is 1.

let text = obj.section_id(object::write::StandardSection::Text);
obj.add_symbol_data(start, text, &code, 1);

And that is it! Now we can emit it, and save it to a file!

let mut out = Vec::new();
obj.emit(&mut out).unwrap();

std::fs::write("hello.o", out).unwrap();

If we take a look at the object generated, we see that everything is as we expected:

$ readelf -Srs hello.o
There are 5 section headers, starting at offset 0xd8:

Section Headers:
  [Nr] Name              Type             Address           Offset
       Size              EntSize          Flags  Link  Info  Align
  [ 0]                   NULL             0000000000000000  00000000
       0000000000000000  0000000000000000           0     0     0
  [ 1] .text             PROGBITS         0000000000000000  00000040
       0000000000000031  0000000000000000  AX       0     0     1
  [ 2] .symtab           SYMTAB           0000000000000000  00000078
       0000000000000030  0000000000000018           3     1     8
  [ 3] .strtab           STRTAB           0000000000000000  000000a8
       0000000000000008  0000000000000000           0     0     1
  [ 4] .shstrtab         STRTAB           0000000000000000  000000b0
       0000000000000021  0000000000000000           0     0     1
Key to Flags:
  W (write), A (alloc), X (execute), M (merge), S (strings), I (info),
  L (link order), O (extra OS processing required), G (group), T (TLS),
  C (compressed), x (unknown), o (OS specific), E (exclude),
  l (large), p (processor specific)

There are no relocations in this file.

Symbol table '.symtab' contains 2 entries:
   Num:    Value          Size Type    Bind   Vis      Ndx Name
     0: 0000000000000000     0 NOTYPE  LOCAL  DEFAULT  UND
     1: 0000000000000000    49 FUNC    GLOBAL HIDDEN     1 _start

Now let’s link and run it:

$ ld hello.o -o hello
$ ./hello
Hello world!

It works!

A zero-dependency Static Compiler

Now we know everything necessary to make the first version of our brainfuck compiler!

First, I will take an old version of our JIT compiler, the one right after we start using dynasm, while we were still using syscalls. This is because now our compiled code will not have access to Rust functions (unless we compile them to a library and link them, as we will see later).

We also will need to make some changes. In our JIT compiler we were passing the program’s memory as an argument to the compiled code, but we will not be able to do this anymore. So instead we will allocate that memory on the stack.

To allocate memory from the stack is very easy, we only need to decrease the value rsp by how much memory we need to allocate (the stack grows downwards). We also need to make sure to clear the memory, because brainfuck programs assume that they are initially zeroed¹.

Note that we are allocating a relatively large value, 30000 bytes, but modern desktop systems have at least 1 MiB of reserved space for the stack, so that should not be a problem.

Another change is that we can no longer finish our program with a ret, because there is no one to return to. So we replace it with a exit syscall.

So here is what the code becomes:

struct Program {
    code: Vec<u8>,
}
impl Program {
    fn new(source: &[u8]) -> Result<Program, UnbalancedBrackets> {
        let mut code: VecAssembler<X64Relocation> = VecAssembler::new(0);

        // r12 will be the adress of `memory`
        // r13 will be the value of `pointer`
        // r13 is set to 0
        dynasm! { code
            ; .arch x64
            ; push rbp
            ; mov rbp, rsp
            ; xor r13, r13

            // allocate 30_0000 bytes on stack for the memory
            ; sub rsp, 30_000
            ; mov r12, rsp

            // zero the memory
            ; xor eax, eax
            ; mov r11, r12
            ; loop_:
            ; mov QWORD [r11], rax
            ; add r11, 8
            ; cmp r11, rbp
            ; jne <loop_
        };

        let mut bracket_stack = Vec::new();

        for b in source {
            match b {
                b'+' => dynasm! { code
                    ; .arch x64
                    ; add BYTE [r12 + r13], 1
                },
                b'-' => dynasm! { code
                    ; .arch x64
                    ; add BYTE [r12 + r13], -1
                },
                b'.' => dynasm! { code
                    ; .arch x64
                    ; mov rax, 1           // write syscall
                    ; mov rdi, 1           // stdout's file descriptor
                    ; lea rsi, [r12 + r13] // buf address
                    ; mov rdx, 1           // length
                    ; syscall
                },
                b',' => dynasm! { code
                    ; .arch x64
                    ; mov rax, 0           // read syscall
                    ; mov rdi, 0           // stdin's file descriptor
                    ; lea rsi, [r12 + r13] // buf address
                    ; mov rdx, 1           // length
                    ; syscall
                },
                b'<' => dynasm! { code
                    ; .arch x64
                    ; sub r13, 1
                    ; mov eax, 29999
                    ; cmovb r13, rax
                },
                b'>' => dynasm! { code
                    ; .arch x64
                    ; add r13, 1
                    ; xor eax, eax
                    ; cmp r13, 30000
                    ; cmove r13, rax
                },
                b'[' => {
                    let start_label = code.new_dynamic_label();
                    let end_label = code.new_dynamic_label();
                    dynasm! { code
                        ; .arch x64
                        ; cmp BYTE [r12+r13], 0
                        ; je =>end_label
                        ; =>start_label
                    };

                    bracket_stack.push((start_label, end_label));
                }
                b']' => {
                    let (start_label, end_label) = match bracket_stack.pop() {
                        Some(x) => x,
                        None => return Err(UnbalancedBrackets(']', code.offset().0)),
                    };

                    dynasm! { code
                        ; .arch x64
                        ; cmp BYTE [r12 + r13], 0
                        ; jne =>start_label
                        ; => end_label
                    };
                }
                _ => continue,
            }
        }

        if !bracket_stack.is_empty() {
            return Err(UnbalancedBrackets(']', code.offset().0));
        }

        dynasm! { code
            ; .arch x64
            ; xor rax, rax
            ; ->exit:
            ; mov rdi, rax // exit error code
            ; mov rax, 60  // exit syscall
            ; syscall
        }

        Ok(Program {
            code: code.finalize().unwrap(),
        })
    }
}

Now the only thing we need to do is to write it to an object file:

fn to_object(&self) -> Vec<u8> {
    let mut obj = object::write::Object::new(
        object::BinaryFormat::Elf,
        object::Architecture::X86_64,
        object::Endianness::Little,
    );

    let start = obj.add_symbol(Symbol {
        name: b"_start".to_vec(),
        value: 0,
        size: 0,
        kind: object::SymbolKind::Text,
        scope: object::SymbolScope::Linkage,
        weak: false,
        section: object::write::SymbolSection::Undefined,
        flags: SymbolFlags::None,
    });

    let text = obj.section_id(object::write::StandardSection::Text);
    obj.add_symbol_data(start, text, &self.code, 16);
    let mut out = Vec::new();
    obj.emit(&mut out).unwrap();

    out
}

And in the main:

use std::{
    path::{Path, PathBuf},
    process::ExitCode,
};
fn main() -> ExitCode {
    let mut args = std::env::args();

    let input_file = args.nth(1).unwrap();
    let source = match std::fs::read(&input_file) {
        Ok(x) => x,
        Err(err) => {
            eprintln!("Error reading '{}': {}", input_file, err);
            return ExitCode::from(2);
        }
    };

    let program = match Program::new(&source) {
        Ok(x) => x,
        Err(UnbalancedBrackets(c, address)) => {
            eprintln!(
                "Error parsing file: didn't found pair for `{}` at instruction index {}",
                c, address
            );
            return ExitCode::from(3);
        }
    };

    let file_name = Path::new(&input_file).file_name().unwrap();
    let output_file = PathBuf::from(&file_name).with_extension("o");

    let obj = program.to_object();
    std::fs::write(output_file, obj).unwrap();

    ExitCode::from(0)
}

And if we run:

$ cargo run -p singlepass-compiler -- programs/factor.bf
    Finished dev [unoptimized + debuginfo] target(s) in 0.03s
     Running `target/debug/singlepass-compiler programs/factor.bf`
$ ld factor.o -o factor
$ ./factor
123456
123456: 2 2 2 2 2 2 3 643

It works!

Linking object files

Now let’s take a look at linking multiple object files.

First let’s modify our hello world example to instead of calling the write syscall directly, it will call a function defined in another file. So we declared that my_write is an external symbol, and call it, passing the string address and length as argument:

section .text
    global _start 
    extern my_write

_start:
    lea rdi, [rel hello] ; string to write
    mov rsi, helloLen    ; length of string to write
    call my_write

    ; Terminate program
    mov rax, 60          ; 'exit' system call = 60
    mov rdi, 0           ; exit with error code 0
    syscall              ; call the kernel
    hello:    db 'Hello world!',10
    helloLen: equ $-hello

And in a file called write.as, we declared my_write as global:

section .text
    global my_write

my_write:
    push rbp        ; push rbp to keep the stack 16 bytes aligned, and
    mov rbp, rsp    ; also setup the stack frame.

    mov rdx, rsi    ; length of string to write
    mov rsi, rdi    ; string to write
    mov rax, 1      ; 'write' system call = 1
    mov rdi, 1      ; file descriptor 1 = STDOUT
    syscall         ; call the kernel

    pop rbp
    ret

If we assemble and link everything, we see that it continues to work:

$ nasm hello.as -f elf64 -o hello.o
$ nasm write.as -f elf64 -o write.o
$ ld hello.o write.o -o hello
$ ./hello
Hello world!

Now if we take a look at hello.o we will see that it now contains relocations:

$ readelf -r hello.o

Relocation section '.rela.text' at offset 0x300 contains 1 entry:
  Offset          Info           Type           Sym. Value    Sym. Name + Addend
00000000000d  000500000002 R_X86_64_PC32     0000000000000000 my_write - 4

The relocation is of type R_X86_64_PC32, which means that it writes a 32-bit word whose value is the offset from the relocation to the given symbol. This is because the call instruction receives an offset as argument.

But the call instruction receives an offset relative to the address of the next instruction, so instead of writing S - P (symbol minus relocation address) it should be written S - (P + 4), i.e., S - P - 4. So the relocation has an addend of -4.

So let’s translate that new version of the code to our example compiler. The code now is:

let mut code: VecAssembler<X64Relocation> = VecAssembler::new(0);

let hello_str = b"Hello world!\n";
let len = hello_str.len() as i32;

let relocation_offset;
dynasm!(code
    ; lea rdi, [>hello]    // string to write
    ; mov rsi, DWORD len   // length of string to write
    ; call DWORD 0
    ;; relocation_offset = code.offset().0 as u64 - 4

    // Terminate program
    ; mov eax,60           // 'exit' system call
    ; mov edi,0            // exit with error code 0
    ; syscall              // call the kernel
    ; hello:
    ; .bytes hello_str
);

let code = code.finalize().unwrap();

As the address of my_write is unknown at this point, the call instruction receives a dummy value for now and the offset where the relocation will take place is saved in relocation_offset.

So, we when we are building the Object, we add the my_write symbol, and the relocation. The new symbol is mostly equal to the other one, the only difference is the name, and that we will not call add_symbol_data to it.

The relocation will have size of 32bits; its kind is Relative; in this particular case the encoding don’t change anything, so I set it to Generic; and I make sure to set the addend to -4, as we have seen previously:

let my_write = obj.add_symbol(Symbol {
    name: b"my_write".to_vec(),
    value: 0,
    size: 0,
    kind: object::SymbolKind::Text,
    scope: object::SymbolScope::Linkage,
    weak: false,
    section: object::write::SymbolSection::Undefined,
    flags: SymbolFlags::None,
});

let text = obj.section_id(object::write::StandardSection::Text);
obj.add_relocation(
    text,
    Relocation {
        offset: relocation_offset,
        size: 32,
        kind: object::RelocationKind::Relative,
        encoding: object::RelocationEncoding::Generic,
        symbol: my_write,
        addend: -4,
    },
)
.unwrap();

If we run and link to write.o we see that it works!

$ cargo run -p object-example -- obj
    Finished dev [unoptimized + debuginfo] target(s) in 0.03s
     Running `../target/debug/object-example obj`
$ ld hello.o write.o
$ ./hello
Hello world!

But now, let’s try, instead of linking to an object written in assembly, link to a static lib written in Rust.

For this, let’s rewrite write.as as write.rs:

#[no_mangle]
pub extern "sysv64" fn my_write(address: *const u8, len: usize) {
    let string = unsafe { std::slice::from_raw_parts(address, len) };
    let string = std::str::from_utf8(string).unwrap();
    print!("{}", string);
}

We need to add the attribute no_mangle to make sure that rustc don’t mangle the name of the symbol for this function.

So we can compile the code to a static library using rustc (I didn’t bother creating an entire cargo project for a single function):

$ rustc --crate-type=staticlib write.rs

This produces a libwrite.a, which is a static library. Static libraries are basically an archive containing multiple object files (you can list them by running ar -t libwrite.a).

Now we can try linking the library with our compiler object (make sure to put object files before the libraries that they depend on):

$ ld hello.o libwrite.a

If run the command above, you will notice that it will output almost 3000 lines of “undefined symbol” errors. That is because the Rust library that we build is static linked to the Rust standard library, that in turn depends on some libraries provided by the OS.

One way of fixing the error is by finding all libraries that define all the missing symbols, and add them as argument to the linker. This is kinda tricky, I am not even sure if this would work, and it will vary between Linux distros, so let’s instead use gcc to link everything. It will already take care of linking all system libraries.

The std also needs some extra static libraries. To get a list of them you can pass --print=native-static-libs to the rustc invocation:

$ rustc --crate-type=staticlib write.rs --print=native-static-libs
note: Link against the following native artifacts when linking against this static library. The order and any duplication can be significant on some platforms.

note: native-static-libs: -lgcc_s -lutil -lrt -lpthread -lm -ldl -lc

$ gcc hello.o libwrite.a -o hello -nostartfiles -lgcc_s -lutil -lrt -lpthread -lm -ldl -lc
$ ./hello
Hello world!

And it works! Note that I need to pass -nostartfiles, because I am declaring _start directly, instead of declaring a main function that would be called by the C runtime.

A Static Compiler

Now we can make the final version of our single-pass compiler. First, we create a library with all functions that we will use in our compiled brainfuck programs.

We will put them in a file called bf_lib.rs. There will be our read and write function, and also a function for exiting. Also, these function will not return errors, because there is no one to report them like in our JIT version. Instead, they only print the error message and exit.

I will also prefix each function with bf_ to avoid symbol conflicts:

use std::io::{Read, Write};

#[no_mangle]
pub extern "sysv64" fn bf_write(value: u8) {
    // Writing a non-UTF-8 byte sequence on Windows error out.
    if cfg!(target_os = "windows") && value >= 128 {
        return;
    }

    let mut stdout = std::io::stdout().lock();

    let result = stdout.write_all(&[value]).and_then(|_| stdout.flush());

    if let Err(err) = result {
        eprintln!("IO error: {}", err);
        std::process::exit(1);
    }
}

#[no_mangle]
pub unsafe extern "sysv64" fn bf_read(buf: *mut u8) {
    let mut stdin = std::io::stdin().lock();
    loop {
        let mut value = 0;
        let err = stdin.read_exact(std::slice::from_mut(&mut value));

        if let Err(err) = err {
            if err.kind() != std::io::ErrorKind::UnexpectedEof {
                eprintln!("IO error: {}", err);
                std::process::exit(1);
            }
            value = 0;
        }

        // ignore CR from Window's CRLF
        if cfg!(target_os = "windows") && value == b'\r' {
            continue;
        }

        *buf = value;
        break;
    }
}

#[no_mangle]
pub unsafe extern "sysv64" fn bf_exit() {
    std::process::exit(0);
}

Now in our compiler, we will call these functions instead of the syscall. We need to register the relocations for each call.

So Program becomes:

struct Program {
    code: Vec<u8>,
    write_relocations: Vec<usize>,
    read_relocations: Vec<usize>,
    exit_relocation: usize,
}

And the compilation:

let mut write_relocations = Vec::new();
let mut read_relocations = Vec::new();

b'.' => dynasm! { code
    ; .arch x64
    ; mov rdi, [r12 + r13] // cell value
    ; call DWORD 0
    ;; write_relocations.push(code.offset().0 - 4)
},
b',' => dynasm! { code
    ; .arch x64
    ; lea rdi, [r12 + r13] // cell address
    ; call DWORD 0
    ;; read_relocations.push(code.offset().0 - 4)
},

let exit_relocation;

dynasm! { code
    ; .arch x64
    ; call DWORD 0
    ;; exit_relocation = code.offset().0 - 4
}

Ok(Program {
    code: code.finalize().unwrap(),
    write_relocations,
    read_relocations,
    exit_relocation,
})

And we make sure to write these relocations on the generated object:

fn to_elf_object(&self) -> Vec<u8> {
    let mut obj = object::write::Object::new(
        object::BinaryFormat::Elf,
        object::Architecture::X86_64,
        object::Endianness::Little,
    );

    let mut add_symbol = |name: &[u8]| {
        obj.add_symbol(Symbol {
            name: name.to_vec(),
            value: 0,
            size: 0,
            kind: object::SymbolKind::Text,
            scope: object::SymbolScope::Linkage,
            weak: false,
            section: object::write::SymbolSection::Undefined,
            flags: SymbolFlags::None,
        })
    };

    let start = add_symbol(b"_start");
    let bf_write = add_symbol(b"bf_write");
    let bf_read = add_symbol(b"bf_read");
    let bf_exit = add_symbol(b"bf_exit");

    let text = obj.section_id(object::write::StandardSection::Text);
    obj.add_symbol_data(start, text, &self.code, 16);

    let mut add_call_reloc = |offset, symbol| {
        obj.add_relocation(
            text,
            Relocation {
                offset: offset as u64,
                symbol,
                size: 32,
                kind: object::RelocationKind::Relative,
                encoding: object::RelocationEncoding::Generic,
                addend: -4,
            },
        )
        .unwrap();
    };

    for offset in self.read_relocations.iter().copied() {
        add_call_reloc(offset, bf_read);
    }
    for offset in self.write_relocations.iter().copied() {
        add_call_reloc(offset, bf_write);
    }
    add_call_reloc(self.exit_relocation, bf_exit);

    let mut out = Vec::new();
    obj.emit(&mut out).unwrap();

    out
}

Now if we compile everything:

$ cargo run -p singlepass-compiler -- ../programs/factor.bf
    Finished dev [unoptimized + debuginfo] target(s) in 0.03s
     Running `../target/debug/singlepass-compiler programs/factor.bf`
$ rustc --crate-type=staticlib bf_lib.rs
$ gcc factor.o libbf_lib.a -nostartfiles -pthread -ldl -o factor
$ ./factor
123456
123456: 2 2 2 2 2 2 3 643

It works! Cool!

A Windows Port

Following the spirit of the precious parts, let’s now try porting our compiler to Windows!

At first, it is very straightforward. In the compiler, we need to tell object to crate a COFF object file on Windows, and we also need to replace the symbol _start by WinMain, which is the one used as entry point on Windows.

let (format, entry_name) = if cfg!(target_os = "windows") {
    (object::BinaryFormat::Coff, "WinMain")
} else if cfg!(target_os = "linux") {
    (object::BinaryFormat::Elf, "_start")
} else {
    unimplemented!("Only Linux and Windows are implemented")
};
let entry_name = entry_name.as_bytes();

let mut obj = object::write::Object::new(
    format,
    object::Architecture::X86_64,
    object::Endianness::Little,
);

let mut add_symbol = |name: &[u8]| { /* ... */ };

let start = add_symbol(entry_name);
// ...

Then, if we use the x86_64-pc-windows-gnu target, we can continue to link our objects with GCC:

> rustc --crate-type staticlib bf_lib.rs --target=x86_64-pc-windows-gnu
> cargo run -p singlepass-compiler -- ../programs/factor.bf
    Finished dev [unoptimized + debuginfo] target(s) in 0.17s
     Running `..\target\debug\singlepass-compiler.exe ../programs/factor.bf`
> gcc -o factor.exe factor.o libbf_lib.a -nostartfiles -ladvapi32 -luserenv -lkernel32 -lkernel32 -lws2_32 -lbcrypt
> factor.exe

If you test the executable now, you may notice that nothing happens. And if you open it on a debugger, you will notice that it is actually crashing on an access violation when clearing the first byte of the memory.

What is happening here is that, on Windows, only part of the reserved stack space is actually committed. When I tested it, for example, there were only 12 KiB of accessible space on the stack.

But at the end of the committed stack, there are guard pages, that whenever they are accessed, more pages are allocated for the stack. Pages have a size of 4096 bytes, but we are allocating 30000 bytes on the stack, which is causing us to miss the guard page, and access uncommitted memory.

The fix for this is very simple, we only need to make sure to not miss the guard page. An easy change is to start clearing our memory at the bottom end of the stack:

dynasm! { code
    ; .arch x64
    ; push rbp
    ; mov rbp, rsp
    ; xor r13, r13

    // allocate 30_0000 bytes on stack for the memory
    ; sub rsp, 30_000
    ; mov r12, rsp

    // zero the memory
    ; xor eax, eax
    ; mov r11, rbp
    ; loop_:
    ; add r11, -8
    ; mov QWORD [r11], rax
    ; cmp r11, r12
    ; jne <loop_
};

If we test our program now:

> rustc --crate-type staticlib bf_lib.rs --target=x86_64-pc-windows-gnu
> cargo run -p singlepass-compiler -- ../programs/factor.bf -o
    Finished dev [unoptimized + debuginfo] target(s) in 0.17s
     Running `..\target\debug\singlepass-compiler.exe ../programs/factor.bf`
> gcc -o factor.exe factor.o libbf_lib.a -nostartfiles -ladvapi32 -luserenv -lkernel32 -lkernel32 -lws2_32 -lbcrypt
> factor.exe
123456
123456: 2 2 2 2 2 2 3 643

It works!

We can also use the msvc tool chain. If we open a “Developer Command Prompt for VS 22” (I have the Visual Studio 2022 installed), we can run the following:

> rustc --crate-type staticlib bf_lib.rs --target=x86_64-pc-windows-msvc -Copt-level=2 -Clto -Cpanic=abort
> link /subsystem:console /entry:WinMain advapi32.lib advapi32.lib userenv.lib kernel32.lib kernel32.lib ws2_32.lib bcrypt.lib msvcrt.lib vcruntime.lib factor.o bf_lib.lib
Microsoft (R) Incremental Linker Version 14.33.31629.0
Copyright (C) Microsoft Corporation.  All rights reserved.

bf_lib.lib(bf_lib.bf_lib.58e9fd71-cgu.2.rcgu.o) : warning LNK4210: .CRT section exists; there may be unhandled static initializers or terminators
msvcrt.lib(tlssup.obj) : warning LNK4210: .CRT section exists; there may be unhandled static initializers or terminators
bf_lib.lib(bf_lib.bf_lib.58e9fd71-cgu.2.rcgu.o) : warning LNK4210: .CRT section exists; there may be unhandled static initializers or terminators
msvcrt.lib(tlssup.obj) : warning LNK4210: .CRT section exists; there may be
unhandled static initializers or terminators

> factor.exe
123456
123456: 2 2 2 2 2 2 3 643

I needed to include vcruntime.lib, even though --print=static-native-libs didn’t report that it was used. Even so, there were some unresolved symbols, but thankfully by setting panic=abort, enabling LTO and increasing the optimization level, these symbols have gone away.

It still outputs some warnings about unhandled static initializers, but thankfully our program don’t need them (I think), and everything works!

Future work

In this post we have seen the basic of object files, and how compile our programs to them, and them link them. But we only covered the very basic of this. There still things that we didn’t cover, like dynamic libraries, position independent executables, debug information, etc.

Also, all these file formats are system dependent, and the relocations in them vary greatly between architectures, so we again fall in the same problem that we discussed in the end of part 2. Thankfully code generators, like LLVM, also covers these problems.

A natural next step here would be to try exactly that using the Cranelift JIT compiler that we made in the last part as basis. You can use cranelift-module and cranelift-object to do that.

But this will be the last part of this blog post series. There is much more to explore on compilation, but brainfuck is too simple a language for them to be need.

But hopefully, these series has taught you at least the basics to let you be able to start writing your own compilers if the will ever came.

All the code developed during these series can be found here (but it may not contain the exactly snippets showed here).

Actually, on Linux, fresh allocated stack for a process is initially zeroed, and because we are the first function being executed (we are at the entry point) we can make sure that the memory is already clear. But I cleared it anyway, because it will be useful when we port this compiler to Windows. ↩

Compiling Brainfuck code - Part 3: A Cranelift JIT Compiler

2022-11-26T21:00:00+00:00

This is the third post of a blog post series where I will reproduce Eli Bendersky’s Adventures In JIT Compilation series, but this time using the Rust programming language.

In the previous part we made an x86 JIT compiler for brainfuck interpreter. In this part we will make a JIT compiler using the Cranelift machine code generator.

What is Cranelift

Cranelift is a machine code generator written in Rust, similar to tools like LLVM or GCC’s backend. As briefly discussed at the end of the last part, these tools are used to create an abstraction layer between language implementations and their compilation targets, by means of a target-independent intermediate representation (IR).

This means that instead of each language having a compiler for each target (a quadratic amount of work), each language compiler can instead target the code generator IR, and the code generator has a backend for each compilation target (a linear amount of work).

Furthermore, these tools also take to job of applying all kinds of optimizations to the generated machine code, since it allows it to be optimal for each target, and also because it improves the performance of dozens of language at once.

But Cranelift has focus on fast compile-time with reasonable runtime performance, and it is also not so mature, so it does not produce heavily optimized code. It is primarily design as a code generator for WebAssembly, so it makes sense to invest more in compile-time speed, where wasm programs are compiled constantly, differently from LLVM for example, that is has more focus on static compilation, where programs are compiled once and run many times.

Basics of Cranelift IR

Using Cranelift boils down to translation our code to Cranelift’s IR (CLIR), that later get optimized and compiled down to machine code. In the examples below I am showing the CLIR in its textual format, but it is more often manipulated in its in-memory data structure format, as we will see in the next section.

Cranelift’s IR, similar to many others code generators IR, uses static single-assignment form (SSA), which basically means that there is no mutable variables in this representation. Each instruction returns a new value, that is only uses as input of subsequent instructions.

This means that a mutable variable in a high level language would be represented by different values throughout the program. As example:

; x = 3 * (2 + 4)
v0 = iconst.i64 3
v1 = iconst.i64 2
v2 = iconst.i64 4
v3 = iadd v1, v2
v4 = imult v0, v3 // x is in v4

; x = 2*x + 1
v5 = imult v1, v4
v6 = iconst 1
v7 = iadd v5, v6 // now x is in v7

These instructions are contained in basic blocks. A basic block is sequence of instructions, with a single entry at the top, and no branches except at the very end. Branches can only target other blocks. In Cranelift the blocks can receive arguments that allow passing values between blocks.

To implement an if-else expression for example, you would need at least four blocks, the start block, an if block, an else block and the after block:

// x = if v0 != 0 { 1 } else { 2 }
block0(v0: i32):
    brnz v0, block1 // branch if not zero
    jump block2

block1:
    v1 = iconst.i32 1
    jump block3(v1)

block2:
    v2 = iconst.i32 2
    jump block3(v2)

block3(v3):
    // x is in v3

A set of basic blocks forms a function, which is the compilation unit of Cranelift. A function has a start block, which inherit the function parameters.

But let’s now see how these are generated in practice.

A Simple Example

Let’s try compiling the same example that we use in the previous part, the add 1 function. It receives one parameter, adds 1, and return it. In x86 assembly, it was only three instructions:

add rdi, 1    ; the argument is in rdi, I add 1 to it.
mov rax, rdi  ; the return value is rax, so I move rdi to rax
ret           ; return from the function

We will generate the following correspondent CLIR:

function u0:0(i64) -> i64 system_v {
block0(v0: i64):
    v1 = iadd_imm v0, 1
    return v1
}

For that we need to make use of the cranelift crate. That crate is an umbrella for two other crates, cranelift-codegen, the core of Cranelift that actually compiles CLIR to machine code, and cranelift-frontend, that is used to write the CLIR.

First we create a Function, that will contain all our CLIR. For that we need to give it a name (for debugging purposes only), and a signature. Its signature involves the calling convention (we will use System V, as always) and its parameters and returns (Cranelift support functions that return multiple values):

use cranelift::codegen::{
    ir::{types::I64, AbiParam, Function, Signature},
    isa::CallConv,
};

let mut sig = Signature::new(CallConv::SystemV);
sig.params.push(AbiParam::new(I64));
sig.returns.push(AbiParam::new(I64));

let mut func = Function::with_name_signature(UserFuncName::default(), sig);

To write CLIR to our function we will need a FunctionBuilder. First we create a FunctionBuilderContext, that is used to recycle resources between the compilation of multiple functions (something we won’t do). Then we pass both the Function and the FunctionBuilderContext to FunctionBuilder:

use cranelift::frontend::{FunctionBuilder, FunctionBuilderContext};

let mut func_ctx = FunctionBuilderContext::new();
let mut builder = FunctionBuilder::new(&mut func, &mut func_ctx);

Now we can start the translation! First we create a basic block, that will be the entry of our function. We append the parameters of our function to the block, and switch to it, allowing us to start writing instruction.

We must also remember to seal our block as soon as all branches to it are defined¹. Because our function have no branches we can seal it right after creating it.

let block = builder.create_block();
builder.seal_block(block);

builder.append_block_params_for_function_params(block);
builder.switch_to_block(block);

Now we can start writing instructions. We grab the value of the first (and only) parameter of the block, add 1 to it, and return it. And then we can finalize our function.

let arg = builder.block_params(block)[0];
let plus_one = builder.ins().iadd_imm(arg, 1);
builder.ins().return_(&[plus_one]);

builder.finalize();

With this our function is complete! If we want to debug what we have generated, we can print the CLIR:

println!("{}", func.display());

function u0:0(i64) -> i64 system_v {
block0(v0: i64):
    v1 = iadd_imm v0, 1
    return v1
}

It is exactly what I have showed at the start of the section (I have copy-pasted it there).

Now we only need to compile it. For that we will need to use Context, that contains everything to compile our function.

First we need to get a TargetISA, that represents the target that we went to compile to. We get one by passing a Triple from the target_lexicon crate to isa::lookup. We will use Tiple::host() to get the native ISA of our host (because we will run on it right after). In out case this is an x86-64 target.

The TargetISA also receives a settings::Flags, that contains some compilations settings, like the level of optimization.

use cranelift::codegen::{isa, settings};
use target_lexicon::Triple;

let builder = settings::builder();
let flags = settings::Flags::new(builder);

let isa = match isa::lookup(Triple::host()) {
    Err(err) => panic!("Error looking up target: {}", err),
    Ok(isa_builder) => isa_builder.finish(flags).unwrap(),
};

Now we can compile our function:

let mut ctx = Context::for_function(func);
let code = ctx.compile(&*isa).unwrap();

This returns a CompileCode struct, the machine code can be retrieved with the code_buffer method.

Now to run it, we can do the same process as we have done in the last part. But instead of using raw libc calls, lets use memmap2’s safe and cross-platform abstractions instead:

let mut buffer = memmap2::MmapOptions::new()
    .len(code.len())
    .map_anon()
    .unwrap();

buffer.copy_from_slice(code.code_buffer());

let buffer = buffer.make_exec().unwrap();

let x = unsafe {
    let code_fn: unsafe extern "sysv64" fn(usize) -> usize =
        std::mem::transmute(buffer.as_ptr());

    code_fn(1)
};

println!("out: {}", x);

And if we run the code it prints 2! It works!

But is the generated code any good? Let’s take a look at it. You can use ctx.set_disasm(true) and code.disasm to get a disassembly generated by Cranelift, which also contains some extra debug information. But it uses AT&T syntax, and in this series I am using Intel syntax, so I will save the code to a file and disassemble it with objdump:

std::fs::write("dump.bin", code.code_buffer()).unwrap();

$ objdump -b binary -D -M intel,x86-64 -m i386:x86-64 dump.bin

dump.bin:     file format binary


Disassembly of section .data:

0000000000000000 <.data>:
   0:   55                      push   rbp
   1:   48 89 e5                mov    rbp,rsp
   4:   48 89 f8                mov    rax,rdi
   7:   48 83 c0 01             add    rax,0x1
   b:   48 89 ec                mov    rsp,rbp
   e:   5d                      pop    rbp
   f:   c3                      ret

As we can see, if you ignore the pushes and moves use for setting up the stack frame, it generated machine code pretty close to what we have done before.

It is cool that Cranelift also takes care of generating the code for setting up the stack frame, but it is not so necessary for leaf functions like this, even more for such small function.

And in fact, Cranelift is supposed to do that, but this is currently not implemented for x86 (see issue #1148). But it appears to be implemented for AArch64…

So, let’s take a look at AArch64 output! We need to enable the “arm64” features of cranelift-codegen:

# make sure both are in the same version
cranelift = "0.89.2"
cranelift-codegen = { version = "0.89.2", features = ["arm64"] }

And change the target ISA to AArch64:

use target_lexicon::triple;
use std::str::FromStr,

let isa = match isa::lookup(triple!("aarch64")) {
    Err(err) => panic!("Error looking up target: {}", err),
    Ok(isa_builder) => isa_builder.finish(flags).unwrap(),
};

Of course, we can no longer execute the code, but we can disassemble its output:

$ aarch64-linux-gnu-objdump -b binary -maarch64 -D dump.bin

dump.bin:     file format binary


Disassembly of section .data:

0000000000000000 <.data>:
   0:   91000400        add     x0, x0, #0x1
   4:   d65f03c0        ret

And look! Only two instructions, no useless stack frame allocation!

Another thing we can look at is what happens when non-optimal code is emitted, for example, if I add 1 two times, one after the other, as we will do in the brainfuck compiler. Could Cranelift be able to optimized it to a single add?

First, let’s enable Cranelift optimizations. For this, we need to use the settings::builder(), that was used for creating the compiler flags:

let mut builder = settings::builder();
builder.set("opt_level", "speed").unwrap();
let flags = settings::Flags::new(builder);

Now we can emit some non-optimal code:

let arg = builder.block_params(block)[0];
let plus_one = builder.ins().iadd_imm(arg, 1);
let plus_two = builder.ins().iadd_imm(plus_one, 1);
builder.ins().return_(&[plus_two]);

And if we run again and disassemble:

$ objdump -b binary -D -M intel,x86-64 -m i386:x86-64 dump.bin

0000000000000000 <.data>:
   0:   55                      push   rbp
   1:   48 89 e5                mov    rbp,rsp
   4:   48 89 f8                mov    rax,rdi
   7:   83 c0 02                add    eax,0x2
   a:   48 89 ec                mov    rsp,rbp
   d:   5d                      pop    rbp
   e:   c3                      ret

Hey, look! A single add eax, 2. It automatically optimized the code for us!

A Cranelift Compiler

Now we can finally start compiling brainfuck. But first we need to take care of some details.

First, the SSA don’t have mutable values, but brainfuck has a mutable variable, the cell pointer (the mutability of the memory itself will be represented by load/store instructions). So we would need to somehow keep track of what value is the current value of pointer, and keeping passing it between blocks.

But thankfully, codegen-frontend already have a functionality that automatically manages variables. We only need to create a Variable, and declared inside our builder. Whenever we need its value, we call builder.use_var, and when we need to set it we call builder.def_var.

And second, we were using I64 as the type of our values in the last example, but for the pointer it should have the pointer type of our current target ISA (as we may change the compilation target later). We can get it by calling isa.pointer_type().

With all this together, we can start writing our compiler:

struct Program {
    code: Vec<u8>,
    memory: [u8; 30_000],
}
impl Program {
    fn new(source: &[u8]) -> Result<Program, UnbalancedBrackets> {
        let mut builder = settings::builder();
        builder.set("opt_level", "speed").unwrap();
        let flags = settings::Flags::new(builder);

        let isa = match isa::lookup(Triple::host()) {
            Err(_) => panic!("x86_64 ISA is not avaliable"),
            Ok(isa_builder) => isa_builder.finish(flags).unwrap(),
        };

        // we get the `pointer_type` from `isa`
        let pointer_type = isa.pointer_type();

        // receive memory address as parameter, and return pointer to io::Error
        let mut sig = Signature::new(CallConv::SystemV);
        sig.params.push(AbiParam::new(pointer_type));
        sig.returns.push(AbiParam::new(pointer_type));

        let mut func = Function::with_name_signature(UserFuncName::user(0, 0), sig);

        let mut func_ctx = FunctionBuilderContext::new();
        let mut builder = FunctionBuilder::new(&mut func, &mut func_ctx);

        // create the variable `pointer` (it is a offset from memory address)
        let pointer = Variable::new(0);
        builder.declare_var(pointer, pointer_type);

        let block = builder.create_block();
        builder.seal_block(block);

        builder.append_block_params_for_function_params(block);
        builder.switch_to_block(block);

        let memory_address = builder.block_params(block)[0];

        // initialize pointer to 0 
        let zero = builder.ins().iconst(pointer_type, 0);
        builder.def_var(pointer, zero);

        let mut stack = Vec::new();

        for (i, b) in source.iter().enumerate() {
            match b {
                b'+' => todo!(),
                b'-' => todo!(),
                b'.' => todo!(),
                b',' => todo!(),
                b'<' => todo!(),
                b'>' => todo!(),
                b'[' => todo!(),
                b']' => todo!(),
                _ => continue,
            }
        }

        if !stack.is_empty() {
            return Err(UnbalancedBrackets(']', source.len()));
        }

        builder.ins().return_(&[zero]);

        builder.finalize();

        let res = verify_function(&func, &*isa);

        if let Err(errors) = res {
            panic!("{}", errors);
        }

        let mut ctx = Context::for_function(func);
        let code = match ctx.compile(&*isa) {
            Ok(x) => x,
            Err(err) => {
                eprintln!("error compiling: {:?}", err);
                std::process::exit(8);
            }
        };

        let code = code.code_buffer().to_vec();

        Ok(Program {
            code,
            memory: [0; 30_000],
        })
    }

Now we only need to implement each instruction.

Whenever we need to read the value of the current cell, we will add memory_address with the value of pointer to get the cell address, and use load.i8 to read it. To write it back, we use store.i8.

These instructions need a MemFlags, which is used to tell if the memory may trap, if it is aligned, etc. Let’s use the default for now:

let mem_flags = MemFlags::new();

So for the + and -, we use iadd_imm:

b'+' => {
    let pointer_value = builder.use_var(pointer);
    let cell_address = builder.ins().iadd(memory_address, pointer_value);
    let cell_value = builder.ins().load(I8, mem_flags, cell_address, 0);
    let cell_value = builder.ins().iadd_imm(cell_value, 1);
    builder.ins().store(mem_flags, cell_value, cell_address, 0);
}
b'-' => {
    let pointer_value = builder.use_var(pointer);
    let cell_address = builder.ins().iadd(memory_address, pointer_value);
    let cell_value = builder.ins().load(I8, mem_flags, cell_address, 0);
    let cell_value = builder.ins().iadd_imm(cell_value, -1);
    builder.ins().store(mem_flags, cell_value, cell_address, 0);
}

For < and > it is similar, but we also need to wrap around. For that we use the select instruction, that will be compiled down to a conditional move:

b'<' => {
    let pointer_value = builder.use_var(pointer);
    let pointer_minus = builder.ins().iadd_imm(pointer_value, -1);

    let wrapped = builder.ins().iconst(pointer_type, 30_000 - 1);
    let pointer_value = builder.ins().select(pointer_value, pointer_minus, wrapped);

    builder.def_var(pointer, pointer_value);
}
b'>' => {
    let pointer_value = builder.use_var(pointer);
    let pointer_plus = builder.ins().iadd_imm(pointer_value, 1);

    let cmp = builder.ins().icmp_imm(IntCC::Equal, pointer_plus, 30_000);
    let pointer_value = builder.ins().select(cmp, zero, pointer_plus);

    builder.def_var(pointer, pointer_value);
}

Now for [ and ] we need to start creating blocks. So when we reach a [ we create two blocks, one for the inside of the brackets, and another for after it. We conditionally jump to one of the blocks based on the current cell value, using a brz/jump pair of instructions. We switch to the inner block (next instructions will be written to it), and we finish by pushing the blocks to the stack.

b'[' => {
    let inner_block = builder.create_block();
    let after_block = builder.create_block();

    let pointer_value = builder.use_var(pointer);
    let cell_address = builder.ins().iadd(memory_address, pointer_value);
    let cell_value = builder.ins().load(I8, mem_flags, cell_address, 0);

    builder.ins().brz(cell_value, after_block, &[]);
    builder.ins().jump(inner_block, &[]);

    builder.switch_to_block(inner_block);

    stack.push((inner_block, after_block));
}

On ], we get the blocks back from the stack, write our conditionally jump logic, and switch to the after_block. Also, at this point we already have written all branches to these blocks, so we can seal them.

b']' => {
    let (inner_block, after_block) = match stack.pop() {
        Some(x) => x,
        None => return Err(UnbalancedBrackets(']', i)),
    };

    let pointer_value = builder.use_var(pointer);
    let cell_address = builder.ins().iadd(memory_address, pointer_value);
    let cell_value = builder.ins().load(I8, mem_flags, cell_address, 0);

    builder.ins().brnz(cell_value, inner_block, &[]);
    builder.ins().jump(after_block, &[]);

    builder.seal_block(inner_block);
    builder.seal_block(after_block);

    builder.switch_to_block(after_block);
}

For the . and ,, we will do something similar to our final single-pass compiler of the last part, and implement them in a Rust function, and call them from the generated code.

For that we will use the call_indirect instruction. It needs the function signature and address:

let (write_sig, write_address) = {
    let mut write_sig = Signature::new(CallConv::SystemV);
    write_sig.params.push(AbiParam::new(I8));
    write_sig.returns.push(AbiParam::new(pointer_type));
    let write_sig = builder.import_signature(write_sig);

    let write_address = write as *const () as i64;
    let write_address = builder.ins().iconst(pointer_type, write_address);
    (write_sig, write_address)
};

let (read_sig, read_address) = {
    let mut read_sig = Signature::new(CallConv::SystemV);
    read_sig.params.push(AbiParam::new(pointer_type));
    read_sig.returns.push(AbiParam::new(pointer_type));
    let read_sig = builder.import_signature(read_sig);

    let read_address = read as *const () as i64;
    let read_address = builder.ins().iconst(pointer_type, read_address);
    (read_sig, read_address)
};

We also need a block to branch to in case of an error:

let exit_block = builder.create_block();
builder.append_block_param(exit_block, pointer_type);

And right after calling builder.finalize():

builder.switch_to_block(exit_block);
builder.seal_block(exit_block);

let result = builder.block_params(exit_block)[0];
builder.ins().return_(&[result]);

So the operations become:

b'.' => {
    let pointer_value = builder.use_var(pointer);
    let cell_address = builder.ins().iadd(memory_address, pointer_value);
    let cell_value = builder.ins().load(I8, mem_flags, cell_address, 0);

    let inst = builder
        .ins()
        .call_indirect(write_sig, write_address, &[cell_value]);
    let result = builder.inst_results(inst)[0];

    let after_block = builder.create_block();

    builder.ins().brnz(result, exit_block, &[result]);
    builder.ins().jump(after_block, &[]);

    builder.seal_block(after_block);
    builder.switch_to_block(after_block);
}
b',' => {
    let pointer_value = builder.use_var(pointer);
    let cell_address = builder.ins().iadd(memory_address, pointer_value);

    let inst = builder
        .ins()
        .call_indirect(read_sig, read_address, &[cell_address]);
    let result = builder.inst_results(inst)[0];

    let after_block = builder.create_block();

    builder.ins().brnz(result, exit_block, &[result]);
    builder.ins().jump(after_block, &[]);

    builder.seal_block(after_block);
    builder.switch_to_block(after_block);
}

And done! If we run it:

14s, about 4% faster than our single-pass implementation. It is definitely faster, but, well, I was hoping that it would have a bigger improvement by replacing all repeat instructions by a single one.

Let’s take a look at its generated code. Here is the code generated by the program +++++:

$ objdump -b binary -D -M intel,x86-64 -m i386:x86-64 dump.bin

dump.bin:     file format binary


Disassembly of section .data:

0000000000000000 <.data>:
   0:   55                      push   rbp
   1:   48 89 e5                mov    rbp,rsp
   4:   48 0f b6 17             movzx  rdx,BYTE PTR [rdi]
   8:   83 c2 01                add    edx,0x1
   b:   88 17                   mov    BYTE PTR [rdi],dl
   d:   83 c2 01                add    edx,0x1
  10:   88 17                   mov    BYTE PTR [rdi],dl
  12:   83 c2 01                add    edx,0x1
  15:   88 17                   mov    BYTE PTR [rdi],dl
  17:   83 c2 01                add    edx,0x1
  1a:   88 17                   mov    BYTE PTR [rdi],dl
  1c:   83 c2 01                add    edx,0x1
  1f:   88 17                   mov    BYTE PTR [rdi],dl
  21:   48 31 c0                xor    rax,rax
  24:   48 89 ec                mov    rsp,rbp
  27:   5d                      pop    rbp
  28:   c3                      ret

As we can see, contrary to what I expected, Cranelift was not capable of merging all sums together when they are between a load and a store.

But looking more closely, it was able to remove the redundant loads (movzx rdx,BYTE PTR [rdi]), but not the stores (mov BYTE PTR [rdi],dl).

Looking around on Cranelift’s GitHub repo, you will find this issue that explains what happened here: the remotion of redundant loads is implemented, but for store it is not.

So, I made a naive implementation of redundant load remotion here (this probably miscompile some code). Only removing the stores was not enough, because the optimization pass that merge multiple adds together happened in a previous pass, so my change also naively rerun all optimizations passes that looked important.

I got the following emitted code:

$ objdump -b binary -D -M intel,x86-64 -m i386:x86-64 dump.bin

dump.bin:     file format binary


Disassembly of section .data:

0000000000000000 <.data>:
   0:   55                      push   rbp
   1:   48 89 e5                mov    rbp,rsp
   4:   48 0f b6 0f             movzx  rcx,BYTE PTR [rdi]
   8:   83 c1 05                add    ecx,0x5
   b:   88 0f                   mov    BYTE PTR [rdi],cl
   d:   48 31 c0                xor    rax,rax
  10:   48 89 ec                mov    rsp,rbp
  13:   5d                      pop    rbp
  14:   c3                      ret

Much better! But running again shows that does not give much improvement (I forget to benchmark this change, sorry for the lack of precise timings).

Actually, this was to be expected, because an actual improvement only comes from optimizing the < and > instructions, was we have seen in the first part. And that one if much more trick to optimize automatically:

# this ones is for the program `>>>+`
$ objdump -b binary -D -M intel,x86-64 -m i386:x86-64 dump.bin

dump.bin:     file format binary


Disassembly of section .data:

0000000000000000 <.data>:
   0:   55                      push   rbp
   1:   48 89 e5                mov    rbp,rsp
   4:   b8 01 00 00 00          mov    eax,0x1
   9:   48 83 c0 00             add    rax,0x0
   d:   48 31 c9                xor    rcx,rcx
  10:   48 81 f9 30 75 00 00    cmp    rcx,0x7530
  17:   48 0f 43 05 49 00 00    cmovae rax,QWORD PTR [rip+0x49]        # 0x68
  1e:   00
  1f:   48 89 c1                mov    rcx,rax
  22:   48 83 c1 01             add    rcx,0x1
  26:   48 81 f8 30 75 00 00    cmp    rax,0x7530
  2d:   48 0f 43 0d 33 00 00    cmovae rcx,QWORD PTR [rip+0x33]        # 0x68
  34:   00
  35:   48 89 ca                mov    rdx,rcx
  38:   48 83 c2 01             add    rdx,0x1
  3c:   48 81 f9 30 75 00 00    cmp    rcx,0x7530
  43:   48 0f 43 15 1d 00 00    cmovae rdx,QWORD PTR [rip+0x1d]        # 0x68
  4a:   00
  4b:   4c 0f b6 44 17 00       movzx  r8,BYTE PTR [rdi+rdx*1+0x0]
  51:   41 83 c0 01             add    r8d,0x1
  55:   44 88 44 17 00          mov    BYTE PTR [rdi+rdx*1+0x0],r8b
  5a:   48 31 c0                xor    rax,rax
  5d:   48 89 ec                mov    rsp,rbp
  60:   5d                      pop    rbp
  61:   c3                      ret

Here it is less obvious how a program could detect that these three cmovae can be replaced by a single one.

But to be fair, Cranelift is design to have faster compilation times, and may not be prepared to optimized such a bad emitted IR.

So let’s optimize these repeated operations during translation. I will not implement the three remaining optimizations yet, because I think they will not be necessary.

First let’s make the source iterator be peekable, and filter out non-instruction symbols:

let mut ops = source
    .iter()
    .enumerate()
    .filter(|x| b"+-><.,[]".contains(x.1))
    .peekable();

while let Some((i, &b)) = ops.next() {
    // ...
}

And on +, -, < and >, we first count how many of these instructions we have, and then emit the optimized IR:

b'+' | b'-' => {
    let mut n = if b == b'+' { 1 } else { -1 };
    while ops.peek().map_or(false, |y| b"+-".contains(y.1)) {
        let b = *ops.next().unwrap().1;
        n += if b == b'+' { 1 } else { -1 };
    }

    let pointer_value = builder.use_var(pointer);
    let cell_address = builder.ins().iadd(memory_address, pointer_value);
    let cell_value = builder.ins().load(I8, mem_flags, cell_address, 0);
    let cell_value = builder.ins().iadd_imm(cell_value, n);
    builder.ins().store(mem_flags, cell_value, cell_address, 0);
}
b'<' | b'>' => {
    let mut n = if b == b'>' { 1 } else { -1 };
    while ops.peek().map_or(false, |y| b"<>".contains(y.1)) {
        let b = *ops.next().unwrap().1;
        n += if b == b'>' { 1 } else { -1 };
    }

    let pointer_value = builder.use_var(pointer);
    let pointer_plus = builder.ins().iadd_imm(pointer_value, n);

    let pointer_value = if n > 0 {
        let wrapped = builder.ins().iadd_imm(pointer_value, n - 30_000);
        let cmp =
            builder
                .ins()
                .icmp_imm(IntCC::SignedLessThan, pointer_plus, 30_000);
        builder.ins().select(cmp, pointer_plus, wrapped)
    } else {
        let wrapped = builder.ins().iadd_imm(pointer_value, n + 30_000);
        let cmp = builder
            .ins()
            .icmp_imm(IntCC::SignedLessThan, pointer_plus, 0);
        builder.ins().select(cmp, wrapped, pointer_plus)
    };

    builder.def_var(pointer, pointer_value);
}

Running it again:

Much faster! But it is still slower than our optimized compiler.

So I gave up on the idea that, currently, Cranelift could apply enough optimizations to replace our optimizations.

So, here is the implementation for the optimized instructions. The MoveUntil instructions was not included, because I noticed that it does not help in generating better optimized code:

Instruction::Clear => {
    let pointer_value = builder.use_var(pointer);
    let cell_address = builder.ins().iadd(memory_address, pointer_value);
    builder.ins().store(mem_flags, zero_byte, cell_address, 0);
}
Instruction::AddTo(n) => {
    let n = n as i64;
    let pointer_value = builder.use_var(pointer);
    let to_add = builder.ins().iadd_imm(pointer_value, n);

    let to_add = if n > 0 {
        let wrapped = builder.ins().iadd_imm(pointer_value, n - 30_000);
        let cmp = builder
            .ins()
            .icmp_imm(IntCC::SignedLessThan, to_add, 30_000);
        builder.ins().select(cmp, to_add, wrapped)
    } else {
        let wrapped = builder.ins().iadd_imm(pointer_value, n + 30_000);
        let cmp = builder.ins().icmp_imm(IntCC::SignedLessThan, to_add, 0);
        builder.ins().select(cmp, wrapped, to_add)
    };

    let from_address = builder.ins().iadd(memory_address, pointer_value);
    let to_address = builder.ins().iadd(memory_address, to_add);

    let from_value = builder.ins().load(I8, mem_flags, from_address, 0);
    let to_value = builder.ins().load(I8, mem_flags, to_address, 0);

    let sum = builder.ins().iadd(to_value, from_value);

    builder.ins().store(mem_flags, zero_byte, from_address, 0);
    builder.ins().store(mem_flags, sum, to_address, 0);
}

And the results:

Now it is faster that our previous implementation! At least, we got something from using Cranelift.

It is also important to notice that the times that I measure here takes in account the running time of the entire program, including the compilation time. Mandelbrot takes at least 100ms to compile, and factor 30 ms, which are already significant times.

But of course, optimizations is only one the strengths of using a machine code generator. Because now our compiler can target 3 different architectures! We have x86-64, AArch64, s390x and however backend they add in the future!

After changing ABI’s of the functions to the host default (we don’t need to stick to a single calling convention anymore), I was finally able to run our compiler on my smartphone (an AArch64), using Termux:

Cool!

Conclusion

At the start of this post I have no idea of how good Cranelift optimizations were. Ideally it would be able to automatically figure out all the optimizations that we came up in the first part. At the other end, it would emit suboptimal machine code, and be slower than our handcrafted implementation.

In the end, we verified that the second scenario didn’t occur, the generated code was actually faster than our implementation.

However, it was far from the first one: even the basic optimization of combining multiple + didn’t happen by default (at least I was able to implement it in Cranelift fairly easy, despite being hackish).

Yet, it accomplished our main goal of using a machine code generator: be able to target multiple architectures with only a fraction of the knowledge and work required.

Note that this post is not showing that Cranelift optimizations are bad, but how bad brainfuck is. Brainfuck only has a single persistent variable and almost every instruction need to read and write to memory.

This means that many of Cranelift optimizations passes are completely underexplored, because they are designed for languages that have multiples values being passed around (more than there are CPU registers), as is the norm for most programming languages.

Future work

Since the beginning of this series we are optimizing our brainfuck implementation more and more. In fact, we could continue creating more and more optimizations to our program, making the execution of out brainfuck programs be faster and faster.

But as we go applying these optimizations, the time necessary to parse and compile the code starts to become greater and greater. To such a point, that applying more optimizations will not actually decrease the time we need to wait to get a brainfuck program to run.

At that point, having such a degree of optimizations is not viable when creating a JIT compiler.

To solve this, we can instead make an AOT compiler². Instead of waiting for the compilation every time we need to run a program, we can instead compile it just one, ahead-of-time (AOT), and serialize it to disk. Now, the user can load and run the program as many times as they want, without needing to wait for the compilation.

And this is exactly what we will do in the next part. We will see how we can get our generated machine code and serialize it to an executable file, like a ELF file.

All the code developed during these series can be found here (may not be exactly equal to one showed here).

Sealing them as soon as possible allows the resolving of Variables to be faster. I present and use Variables when we start making our brainfuck compiler. ↩
I used the term “AOT” here to have a contrast with “JIT”, but it is a term more used in context where the language is expected to be JIT compiled. Normally, the term “compiled language” already implies that the language is compiled ahead of time. AOT is also used in scenarios where the compilation is not emitting machine code directly, so “static compiled” or “native compiled” may be better suitable to “compiled languages” than “AOT compiled”. ↩

Compiling Brainfuck code - Part 2: A Singlepass JIT Compiler

2022-11-07T11:00:00+00:00

This is the second post of a blog post series where I will reproduce Eli Bendersky’s Adventures In JIT Compilation series, but this time using the Rust programming language.

In the previous part we made an optimized brainfuck interpreter. In this part we will make a brainfuck JIT compiler.

The Interpreter Overhead

In the end of the previous post, I said that the implemented interpreter have yet some optimizations opportunities to explore, but that there always would be a significant performance overhead by its program loop.

To grasp exactly how big is this overhead, we can take a look at generated assembly for the main loop of our interpreter, were the majority of the runtime is spent.

In this case, I used cargo-show-asm to get the x86 assembly for Program::run, and (after a lot clean up) I got the following:

; 'program: loop (not including first iteration)
.PROGRAM:
    ; r13 is &self
    ; rcx is self.program_counter
    ; rdx is self.instructions.len()
    mov rcx, qword ptr [r13]
    mov rdx, qword ptr [r13 + 32]
    add rcx, 1
    mov qword ptr [r13], rcx
    cmp rdx, rcx                  ; break if instructions.len() == program_counter
    je .BREAK_PROGRAM             ; .

    ; rsi is pointer
    ; rax is self.instructios.ptr

    mov rsi, qword ptr [r13 + 8]
    mov rax, qword ptr [r13 + 16]

    ; `match instruction[pointer]` by using a jump table
    movzx edi, byte ptr [rax + rcx]     ; instr = instruction[program_counter]
    movsxd rdi, dword ptr [r12 + 4*rdi] ; rel_jump = jump_table[instr]
    add rdi, r12                        ; jump += rel_jump + &jump_table
    jmp rdi                             ; goto `jump`

; Increase (`+`)
.LBB8_1:
    add byte ptr [r13 + rsi + 40], 1
    jmp .PROGRAM

; Decrease (`-`)
.LBB8_0:
    add byte ptr [r13 + rsi + 40], -1
    jmp .PROGRAM

; MoveRight (`>`)
.LBB8_1:
    add esi, 1

    movzx eax, si
    mov ecx, eax
    shr ecx, 4
    imul ecx, ecx, 2237
    shr ecx, 22
    imul ecx, ecx, 30000
    sub eax, ecx
    movzx eax, ax
    mov qword ptr [r13 + 8], rax
    jmp .PROGRAM

; Etc..

I added some comments to the assembly to make it a little more clear what is going on, but the main point here is the amount of instructions that is used to implement the interpreter loop compared to the amount used for the brainfuck operations.

The + and - operations are being implemented by a single assembly instruction (the add byte ptr [r13 + rsi + 40], 1), but after executing each instruction it needs to do a jump and run everything in the .PROGRAM block, who has 12 instructions. Instructions count is not a very precise performance parameter, but this roughly means that the interpreter has an overhead of 13x when running a + or - operation.

This is a lesser problem for brainfuck operations with bigger implementations, like the > there, who has a very involved wrapping calculation (could be greatly improved by the memory length being a power of 2)¹, but the overhead is still very significant, almost 2x.

This roughly means that if we could eliminate the loop overhead, we can have something between 2x and 13x of performance increase!

And to achieve this we will implement a JIT compiler.

What is a compiler and a JIT compiler

A compiler is basically a program that transforms our high level code, easy to write and understand, to a lower level code, easy to execute. In this case, we will compile our brainfuck code to machine code, the code used by the computer’s processor to execute instructions.

Each processor architecture has a different set of instructions, that are encoded by different machine code formats.

In this post, we will be generating machine code for the x86-64 CPU architecture, which is the most common processor architecture used by personal computers today. But if you want to target mobile devices, or the newer Macs, you will probably also want to target AArch64.

A JIT compiler is a compiler that compiles just-in-time. That is, it compiles the program to machine code only just before executing it. This is the opposite of what normal compilers do, like Rust, who compiles the entire program into native executables before distributing them.

JIT compilers have a series of advantages, like not needing to distribute multiple versions of each program compiled for each architecture and operating system, or being able to apply optimizations specific to each particular computer and runtime workload.

But also, it needs less knowledge to implement a JIT compiler than to create a native executable, so it is easier to explore it here in this post (but I plan to generate native executable in a later post).

Basics of a JIT compiler

To write a JIT compiler, you need basically do things:

Generate machine code.
Execute that code, right after generating it.

Generating machine code

Machine code is basically the binary encoding for the instructions that the processor execute, and assembly is the human-readable encoding of these instructions. An assembler is the program that convert assembly to the machine code.

For example, take the assembly instruction that was implementing the + in the assembly at the start of the post. It was add byte ptr [r13 + rsi + 40], 1, which is an instruction that adds the immediate value 1 to the byte at the address given by r13 + rsi + 40 (on intel’s assembly syntax, the first operant of an instruction is the target of the operation, and the second is the source).

If you follow the encoding steps for this instruction (which I tried, but discovered that this would take much more time to learn that I expected), or use an assembler as a sane person would do, you will discover that this instruction is encoded by the byte sequence 0x41 0x80 0x44 0x35 0x28 0x01.

When we create our compiler, we will emit a series of instructions that implement our desired program. To encode multiple instructions, you simply need to concatenate the encoding of each instruction.

And we also need to make sure to terminate our program with some termination instruction, otherwise the processor would start executing any random bytes that happens to be after the last instruction. In our case it will be a ret instruction, which will return the execution from the generated code, back to our Rust code, was we will see in the next section. But this also could be a exit syscall, which would terminate the entire program.

As an example, there is the assembly for a function that adds 1 to its 64-bit argument and return it:

add rdi, 1    ; the argument is in rdi, I add 1 to it.
mov rax, rdi  ; the return value is rax, so I move rdi to rax
ret           ; return from the function

Passing it through the Netwide Assembler, followed by objdump, we get the machine code:

$ nasm add1.s -felf64
$ objdump -M intel -d add1.o

add1.o:     file format elf64-x86-64


Disassembly of section .text:

0000000000000000 <.text>:
   0:	48 83 c7 01          	add    rdi,0x1
   4:	48 89 f8             	mov    rax,rdi
   7:	c3                   	ret    

And we could write it in a Rust program:

let add1 = [
    0x48, 0x83, 0xc7, 0x01, // add rdi, 1
    0x48, 0x89, 0xf8,       // mov rax, rdi
    0xc3,                   // ret
];

Executing that code

And to execute that code we only need to cast it to a function pointer and call it!

let add1: unsafe extern "sysv64" fn(u64) -> u64 = unsafe {
    std::mem::transmute(add1.as_ptr())
};

let x = unsafe {
    add1(1)
};

println!("{x}");

Well, no. If you run that it will only output “Segmentation fault”.

But this is basically what we need to do. The problem is that operations systems have restrictions on what memory regions have permission to be read, write or execute. And by default, allocated memory has only read and write permissions, they don’t have execute permission to avoid arbitrary code execution vulnerabilities.

So to make this work, we only need to use some operating system API’s that can allocate memory with specified permissions, or can change the permission of a memory region.

On POSIX systems, like Linux, that would be mmap and mprotect. On Windows, that would be VirtualAlloc and VirtualProtect.

So in our Rust code, we could use the libc crate to call mmap on Linux, and get an executable memory region:

let add1: extern "sysv64" fn(u64) -> u64 = unsafe {
    let mem = libc::mmap(
        std::ptr::null_mut(),
        add1.len(),
        libc::PROT_READ | libc::PROT_WRITE | libc::PROT_EXEC,
        libc::MAP_PRIVATE | libc::MAP_ANONYMOUS,
        -1,
        0,
    );
    
    if mem == libc::MAP_FAILED {
        println!("mmap failed");
        return;
    }

    let mem = std::slice::from_raw_parts_mut(mem as *mut u8, add1.len());

    mem.copy_from_slice(&add1);

    std::mem::transmute(mem.as_ptr())
};

let x = add1(1);

println!("{x}");

And if you run this (here is a playground), it prints 2!

But what we are doing is completely circumventing the memory protection. Ideally we should never have exec and write permissions for the same memory.

We can achieve that by allocating memory with read and write permissions, write the code to it, and then change it to read and exec permissions. This can be done with the help of mprotect, as long as the given pointer is page aligned, which the mmap guaranties (playground):

let add1: extern "sysv64" fn(u64) -> u64 = unsafe {
    let mem = libc::mmap(
        std::ptr::null_mut(),
        add1.len(),
        libc::PROT_READ | libc::PROT_WRITE,
        libc::MAP_PRIVATE | libc::MAP_ANONYMOUS,
        -1,
        0,
    );

    if mem == libc::MAP_FAILED {
        println!("mmap failed");
        return;
    }

    std::slice::from_raw_parts_mut(mem as *mut u8, add1.len())
        .copy_from_slice(&add1);

    let result = libc::mprotect(
        mem,
        add1.len(),
        libc::PROT_READ | libc::PROT_EXEC,
    );

    if result == -1 {
        println!("mprotect failed");
        return;
    }

    std::mem::transmute(mem.as_ptr())
};

Also, notice that we are leaking the mmap allocated memory, you need to call munmap to free it. I didn’t do it here because the program will finish immediately after running the code, so the OS will already free it anyway².

The calling Convention

Before we start writing our JIT compiler we need to revise one last thing, the calling convention. In the previous snippets I was casting the program code to a extern "sysv64" fn(u64) -> u64, where the extern "sysv64" specifies that that function pointer is using the System V ABI.

The ABI dictates, among other things, what is the calling convention. That is what guarantees that the first argument will be in the rdi register, and the return value should be in the rax, for example.

Another important point of the calling convention are the preserved registers, and the scratch registers. The preserved registers are the ones that a function call does not change. This means that when we generate our function, we need to save these register on the stack before using them, and must restore them before returning.

The scratch registers are the opposite, they can be modified by a function call. This also means that if we will need a value of one of these register, we need to save them (in another register, or on the stack) before executing a call.

Here is the quick reference that I am using to check the register names³ and types.

So when we start writing our compiled code, we need to be sure to not use a scratch register to hold long persisted values (which will be only the address of the memory and the value of the pointer), and remember to save and restore any preserved register that we might be using.

A Singlepass JIT Compiler

So now we have everything necessary to build our Brainfuck JIT compiler. We will start with our basic interpreter, which will not have any of our applied optimizations, except for the precomputed paring brackets, that will be compiled down to conditional jumps.

This interpreter will have a single pass, which means that it will already emit the compiled code while still traversing the source code. In non-singlepass compilers, the source code may be entirely collected in a new intermediate representation, and then passed to the next pass.

First we need to decide which register will be used for what. I will use the r12 register for the address of the memory of the array of cells, which will initially be passed as an argument in the rdi, and r13 will hold the value of the pointer. Because they are preserved registers, we need to save them.

So we start our compiler by generation our code prelude and epilogue. In assembly, the prelude of a function is the code that push the necessary registers to the stack and set things like the stack frame, and the epilogue is where the registers are popped, etc.

Another thing that the prelude maintain is the stack alignment requirement, another aspect of the ABI. In the System V ABI, the stack need to be aligned by 16 bytes wherever we do a call⁴. Because the call of the current function push the return address to the stack, it starts with unaligned by 8 bytes. We will push two registers to the stack, increasing it by 16 bytes, keeping it with the same unalignment.

To fix that, we need to decrease the stack pointer by 8 (the stack grows downwards), or push another register to the stack. Because we already need to do that, I will take the opportunity to maintain the base pointer, rbp. rbp points to the start of the stack frame, and also the previous value of rbp. This is used for creating backtraces in a debugger for example.

So in the start of our prelude we need to push rbp to save the old value, and mov rbp, rsp to make it point to it.

Now we can put everything together, and start our compiler. I copied the code from the optimized interpreter, keeping the main function, but modifying the content of Program and its methods:

use std::io::Write;

struct Program {
    code: Vec<u8>,
    memory: [u8; 30_000],
}
impl Program {
    fn new(source: &[u8]) -> Result<Program, UnbalancedBrackets> {
        let mut code = Vec::new();

        // ; r12 will be the adress of `memory`
        // ; r13 will be the value of `pointer`
        // ; r12 is got from argument 1 in `rdi`
        // ; r13 is set to 0
        // push rbp
        // mov rbp, rsp
        // push r12
        // push r13
        // mov r12, rdi
        // xor r13, r13
        code.write_all(&[
            0x55, //
            0x48, 0x89, 0xe5, //
            0x41, 0x54, //
            0x41, 0x55, //
            0x49, 0x89, 0xfc, //
            0x4d, 0x31, 0xed,
        ])
        .unwrap();


        for b in source {
            match b {
                b'+' => todo!(),
                b'-' => todo!(),
                b'.' => todo!(),
                b',' => todo!(),
                b'<' => todo!(),
                b'>' => todo!(),
                b'[' => todo!(),
                b']' => todo!(),
                _ => continue,
            }
        }

        // ; when we push to the stack, we need to remeber
        // ; to pop them in the opossite order.
        // pop r13
        // pop r12
        // pop rbp
        // ret
        code.write_all(&[
            0x41, 0x5d, //
            0x41, 0x5c, //
            0x5d,
            0xc3,
        ])
        .unwrap();

        Ok(Program {
            code,
            memory: [0; 30_000],
        })
    }

    // ...
}

Now we need to implement each instruction. For the + and - we can use the already seen add byte ptr [r13 + rsi + 40], 1, but now using r12 and r13:

b'+' => {
    // add byte [r12 + r13], 1
    code.write_all(&[0x43, 0x80, 0x04, 0x2c, 0x01]).unwrap();
}
b'-' => {
    // add byte [r12 + r13], -1
    code.write_all(&[0x43, 0x80, 0x04, 0x2c, 0xff]).unwrap();
}

The > and < could also be a single instruction, add r13, 1 and add r13, -1, but we need to implement the wrapping behavior. I used module operator to implement that in the interpreter, but I discovered that operation is very slow to perform, and the compiler decided to use some compiler magic, that resulted in that big series of multiplications and shifts in the disassembly seen at the start of this post.

So instead, I will implement the wrapping using conditional logic. It is uncertain if the branching will be more performant than the first implementation, but it will be easier to write.

Using this playground’s ASM output, I got the following implementation:

playground::add1:
	lea	rcx, [rdi + 1]
	xor	eax, eax
	cmp	rcx, 30000
	cmovne	rax, rcx
	ret

playground::sub1:
	sub	rdi, 1
	mov	eax, 29999
	cmovae	rax, rdi
	ret

Which does not contain branching at all! Only a conditional move (didn’t know that instructions existed). So there is a high change that this will be better than the interpreter implementation.

So, after changing the register to the ones we are using, and inverting some conditions, we get the following:

b'<' => {
    // sub r13, 1
    // mov eax, 29999
    // cmovb r13, rax
    code.write_all(&[
        0x49, 0x83, 0xed, 0x01, //
        0xb8, 0x2f, 0x75, 0x00, 0x00, //
        0x4c, 0x0f, 0x42, 0xe8,
    ])
    .unwrap();
}
b'>' => {
    // add r13, 1
    // xor eax, eax
    // cmp r13, 30000
    // cmove r13, rax
    code.write_all(&[
        0x49, 0x83, 0xc5, 0x01, //
        0x31, 0xc0, //
        0x49, 0x81, 0xfd, 0x30, 0x75, 0x00, 0x00, //
        0x4c, 0x0f, 0x44, 0xe8,
    ])
    .unwrap();
}

Now lets address the . and ,. These operators interact with the stdout and stdin. To be able to access them, we will need to make use of system calls.

A system call, or syscall, is an instruction that call a function that lives in the operating system kernel. We will need to use the read and write syscalls.

These syscall are the ones that read and write from a file descriptor. To call them you need to set the syscall call number to rax, pass its arguments using the calling convention, and finally execute the syscall instruction.

The write syscall, on x86-64 Linux, has call number 1, and receives as parameter the file descriptor (which for stdout is 1) and a pointer and a length for the buffer whose content will be written. It returns the number of bytes written, or -1 in case of error.

You are expected to handle the error, or cases where the entire buffer was not written, but I will avoid writing too much complicate assembly for now, it will already be enough for our case:

b'.' => {
    // mov rax, 1 ; write syscall
    // mov rdi, 1 ; stdout's file descriptor
    // lea rsi, [r12 + r13] ; buf address
    // mov rdx, 1           ; length
    // syscall
    code.write_all(&[
        0xb8, 0x01, 0x00, 0x00, 0x00, //
        0xbf, 0x01, 0x00, 0x00, 0x00, //
        0x4b, 0x8d, 0x34, 0x2c, //
        0xba, 0x01, 0x00, 0x00, 0x00, //
        0x0f, 0x05, //
    ])
    .unwrap();
}

The read syscall is very similar, but its call number is 0, and it writes to the passed buffer. Also, now we pass the file descriptor 0, which is the stdin:

b',' => {
    // mov rax, 0 ; read syscall
    // mov rdi, 0 ; stdin's file descriptor
    // lea rsi, [r12 + r13] ; buf address
    // mov rdx, 1           ; length
    // syscall
    code.write_all(&[
        0xb8, 0x00, 0x00, 0x00, 0x00, //
        0xbf, 0x00, 0x00, 0x00, 0x00, //
        0x4b, 0x8d, 0x34, 0x2c, //
        0xba, 0x01, 0x00, 0x00, 0x00, //
        0x0f, 0x05, //
    ])
    .unwrap();
}

And at last we need to implement [ and ], which is the tricker. To implement it we will need to write a conditional relative jump. In assembly, it will be like:

cmp byte [r12+r13], 0
je near <jump_offset>

That is, we compare the value of the current cell to 0, and will jump if it is equal (je) or if not equal (jne) depending on whether it is [ or ]. The jump instruction receives as parameter the offset of the jump, and the near keyword there is to tell that the offset will be a i32 value. Could be a smaller value (by using short), depending on how big the jump is, but checking the jump length would complicate things a little, so let’s keep it simple.

First, to write the machine code of that instruction we need to know how the offset is encoded in it. Thankfully the instruction encoding is very simple, it is only two bytes of the opcode followed by the four little endian bytes of the offset.

But we also need to discover the offset. For that, we will do something very similar to the pair address precomputation in the previous post.

First we create a stack. And every time we encounter a [, we write its implementation (with the jump filled with a dummy offset) and push the byte index of the next instruction to the stack.

let mut bracket_stack = Vec::new();

for b in source {
    // ...

b'[' => {
    // ; note that the offset of 0 is a dummy value,
    // ; it will be fixed in the pair `]`
    // cmp byte [r12+r13], 0
    // je near .END
    code.write_all(&[
        0x43, 0x80, 0x3c, 0x2c, 0x00, //
        0x0f, 0x84, 0x00, 0x00, 0x00, 0x00,
    ])
    .unwrap();

    // push to the stack the byte index of the next instruction.
    bracket_stack.push(code.len() as u32);
}

And every time we encounter a ], we pop the index from the stack, compute the offset, and fix the offset of the last jump. Note that the offset that the encoded jump receives is from the address of the next instruction, to the address of the target instruction.

b']' => {
    let left = match bracket_stack.pop() {
        Some(x) => x as usize,
        None => return Err(UnbalancedBrackets(']', code.len())),
    };

    // cmp byte [r12 + r13], 0
    // jne near .START
    code.write_all(&[
        0x43, 0x80, 0x3c, 0x2c, 0x00, //
        0x0f, 0x85, 0xf1, 0xff, 0xff, 0xff,
    ])
    .unwrap();

    // the byte index of the next instruction
    let right = code.len();

    let offset = right as i32 - left as i32;

    // fix relative jumps offsets
    code[left - 4..left].copy_from_slice(&offset.to_le_bytes());
    code[right - 4..right].copy_from_slice(&(-offset).to_le_bytes());
}

And that is it! Now we only need to write the generated code to an executable buffer, cast to pointer and call it:

fn run(&mut self) -> std::io::Result<()> {
    unsafe {
        let len = self.code.len();
        let mem = libc::mmap(
            std::ptr::null_mut(),
            len,
            libc::PROT_READ | libc::PROT_WRITE,
            libc::MAP_PRIVATE | libc::MAP_ANONYMOUS,
            -1,
            0,
        );

        if mem == libc::MAP_FAILED {
            panic!("mmap failed");
        }

        // SAFETY: mem is zero initalized by the mmap.
        std::slice::from_raw_parts_mut(mem as *mut u8, len).copy_from_slice(&self.code);

        // mem.as_ptr() is page aligned, because it is get from mmap.
        let result = libc::mprotect(mem, len, libc::PROT_READ | libc::PROT_EXEC);

        if result == -1 {
            panic!("mprotect failed");
        }

        let code_fn: unsafe extern "sysv64" fn(*mut u8) = std::mem::transmute(mem);

        code_fn(self.memory.as_mut_ptr());

        let result = libc::munmap(mem, len);

        if result == -1 {
            panic!("munmap failed");
        }
    }
    Ok(())
}

And after using a debugger to step through the generated assembly, going back and fixing them, it’s finished! I finally get to measure its performance:

It is almost 7 times faster than the interpreted one! It is within the expected performance increase from the analysis of the interpreter overhead.

The comparison is not yet fair with the basic interpreter, because we still need to check and return the errors from the write and read syscall and implement the EOF behavior. But these changes will only make it a little slower (only a little, IO is the least common instruction).

Even so, the optimized interpreter is still at par with this JIT compiler. Nevertheless, nothing prevent us to apply the exact optimizations used in the interpreter, and on top of that compile them to machine code.

But before we continue further, we first need to address how we can improve our workflow.

Improving our workflow with Dynasm

My workflow for writing this JIT compiler was the following:

Inspect some assembly emitted from Rust, to learn how to implement the operation.
Write assembly to a temporary file.
Assemble and disassemble the assembly to get the machine code, using nasm test.s -felf64 && objdump -M intel -d test.o
Copy the hexadecimal output to my JIT source code.
Test it, go back to step 2 if didn’t work.

This workflow is not so bad, at least I don’t need to read Intel’s documentation to discover how to encode the instructions. Well… I still need to search a little when I need to dynamically set some parameter of the instruction.

Steps 1 and 5 are unavoidable, but steps 2, 3 and 4 creates a big overhead over the iterative process, and introduce new error points. If somehow a typo is introduced in one the bytes of the machine code, the only reasonable way of discovery which one is by disassembling them again.

What would really help would be if I could write the assembly directly in my Rust source code, and it assembles it to a byte sequence at compile time. Better yet if it could automatically figure out how to inject dynamic parameter in the machine code.

And every thing describe here can be implemented using Rust procedural macros. And has already been done!

Dynasm-rs is a dynamic assembler for Rust, a tool that ease the creation of programs that require run-time assembling. It is inspired by LuaJIT’s famous DynASM, but uses procedural macros instead of a C preprocessor.

To us, this means that we replace code like this:

// push rbp
// mov rbp, rsp
// push r12
// push r13
// mov r12, rdi
// xor r13, r13
code.write_all(&[
    0x55, //
    0x48, 0x89, 0xe5, //
    0x41, 0x54, //
    0x41, 0x55, //
    0x49, 0x89, 0xfc, //
    0x4d, 0x31, 0xed,
])
.unwrap();

By this:

dynasm! { code
    ; .arch x64
    ; push rbp
    ; mov rbp, rsp
    ; push r12
    ; push r13
    ; mov r12, rdi
    ; xor r13, r13
};

Much easier to maintain. And the dynasm! macro compiles that down to a single code.extend() call. But for the remainder of the assembly blocks, dynasm will rely on more methods of its DynasmApi, that can be found in the dynasmrt crate.

But we will need only a handful of them, so we don’t need to import dynasmrt (yet), and instead copy the trait to our code:

trait DynasmApi {
    fn push_i8(&mut self, value: i8);
    fn push_i32(&mut self, value: i32);
}
impl DynasmApi for Vec<u8> {
    fn push_i8(&mut self, value: i8) {
        self.push(value as u8)
    }
    fn push_i32(&mut self, value: i32) {
        self.extend_from_slice(&value.to_le_bytes())
    }
}

That will be enough to replace all raw machine code blocks:

fn new(source: &[u8]) -> Result<Program, UnbalancedBrackets> {
    let mut code = Vec::new();

    // r12 will be the adress of `memory`
    // r13 will be the value of `pointer`
    // r12 is got from argument 1 in `rdi`
    // r13 is set to 0
    dynasm! { code
        ; .arch x64
        ; push rbp
        ; mov rbp, rsp
        ; push r12
        ; push r13
        ; mov r12, rdi
        ; xor r13, r13
    };

    let mut bracket_stack = Vec::new();

    for b in source {
        match b {
            b'+' => dynasm! { code
                ; .arch x64
                ; add BYTE [r12 + r13], 1
            },
            b'-' => dynasm! { code
                ; .arch x64
                ; add BYTE [r12 + r13], -1
            },
            b'.' => dynasm! { code
                ; .arch x64
                ; mov rax, 1 // write syscall
                ; mov rdi, 1 // stdout's file descriptor
                ; lea rsi, [r12 + r13] // buf address
                ; mov rdx, 1           // length
                ; syscall
            },
            b',' => dynasm! { code
                ; .arch x64
                ; mov rax, 0 // read syscall
                ; mov rdi, 0 // stdin's file descriptor
                ; lea rsi, [r12 + r13] // buf address
                ; mov rdx, 1           // length
                ; syscall
            },
            b'<' => dynasm! { code
                ; .arch x64
                ; sub r13, 1
                ; mov eax, 29999
                ; cmovb r13, rax
            },
            b'>' => dynasm! { code
                ; .arch x64
                ; add r13, 1
                ; xor eax, eax
                ; cmp r13, 30000
                ; cmove r13, rax
            },
            b'[' => {
                // note that the offset of 0 is a dummy value,
                // it will be fixed in the pair `]`
                dynasm! { code
                    ; .arch x64
                    ; cmp BYTE [r12+r13], 0
                    ; je i32::max_value()
                };

                // push to the stack the byte index of the next instruction.
                bracket_stack.push(code.len() as u32);
            }
            b']' => {
                let left = match bracket_stack.pop() {
                    Some(x) => x as usize,
                    None => return Err(UnbalancedBrackets(']', code.len())),
                };

                dynasm! { code
                    ; .arch x64
                    ; cmp BYTE [r12 + r13], 0
                    ; jne i32::max_value()
                };

                // the byte index of the next instruction
                let right = code.len();

                let offset = right as i32 - left as i32;

                // fix relative jumps offsets
                code[left - 4..left].copy_from_slice(&offset.to_le_bytes());
                code[right - 4..right].copy_from_slice(&(-offset).to_le_bytes());
            }
            _ => continue,
        }
    }

    if !bracket_stack.is_empty() {
        return Err(UnbalancedBrackets(']', code.len()));
    }

    // when we push to the stack, we need to remeber
    // to pop them in the opossite order.
    dynasm! { code
        ; .arch x64
        ; xor rax, rax
        ; ->exit:
        ; pop r13
        ; pop r12
        ; pop rbp
        ; ret
    }

    Ok(Program {
        code,
        memory: [0; 30_000],
    })
}

Much better! But we still are relying on in how the offset of the jump is encoded int the machine code. But thankfully dynasm can also handle that, we only need to start using the dynasmrt crate.

So first can replace our simple Vec<u8> + ad-hoc extension trait, by dynasmrt::VecAssembler:

use dynasmrt::{
    dynasm, x64::X64Relocation, DynasmApi, 
    DynasmLabelApi, VecAssembler
};

let code = VecAssembler<X64Relocation>::new();

Now, on [, we create two dynamic labels, one for the start of the loop, and one for the end. We use => to use the dynamic label inside the macro. The end_label go as the parameter of the jump, and start_label goes to a standalone line, defining the label position. And we finish by putting both labels at the stack.

b'[' => {
    let start_label = code.new_dynamic_label();
    let end_label = code.new_dynamic_label();

    dynasm! { code
        ; .arch x64
        ; cmp BYTE [r12+r13], 0
        ; je =>end_label
        ; =>start_label
    };

    bracket_stack.push((start_label, end_label));
}

And on ], we pop the labels from the stack, and use them in a similar way:

b']' => {
    let (start_label, end_label) = match bracket_stack.pop() {
        Some(x) => x,
        None => return Err(UnbalancedBrackets(']', code.offset().0)),
    };

    dynasm! { code
        ; .arch x64
        ; cmp BYTE [r12 + r13], 0
        ; jne =>start_label
        ; =>end_label
    };
}

Now we don’t need to know how the jump offsets are encoded!

Dynasm also has some help utilities for running our generated code.

We can replace our unsafe calls to mmap, mprotect and munmap by dynasm’s MutableBuffer and ExecutableBuffer. Not only they encapsulate the unsafe code for us, but also provides a nice cross-platform abstraction⁵, that will be useful for my eventual Windows support.

So, our Program::run becomes:

fn run(&mut self) -> std::io::Result<()> {
    let mut buffer = MutableBuffer::new(self.code.len()).unwrap();
    buffer.set_len(self.code.len());

    buffer.copy_from_slice(&self.code);

    let buffer = buffer.make_exec().unwrap();

    unsafe {
        let code_fn: unsafe extern "sysv64" fn(*mut u8) = std::mem::transmute(buffer.as_ptr());
        code_fn(self.memory.as_mut_ptr());
    }

    Ok(())
}

Implement IO error handling and EOF behavior

Now with our enhanced workflow, we can continue to implement the final details necessary for our JIT compiler be 100% compatible with our basic interpreter. For that we need to retry the read/write syscall until all bytes are written, check if any of the calls return an error and repass it to the Rust code, and check if the stdin reached the end-of-file and fallback to the value 0.

That is a lot of code to be implemented in assembly. But thankfully we not need to! We can instead implement every thing in Rust! (minus the error repass)

We only need to declare a Rust function with the proper ABI, and call it directly from the assembly.

Not only it will make implementing all this much easier, we also get a cross-platform implementation.

We can directly copy the code from our interpreter and put in a extern "sysv64" fn(*const u8) function, and call it from our generated code.

But we need to adapt it to also return the IO error. Because Rust ABI is not stable, we cannot return the Error type directly. So instead I will put it in a Box, and return its pointer, using Box::into_raw().

Putting everything together, the read and write functions become:

extern "sysv64" fn write(value: u8) -> *mut std::io::Error {
    // Writing a non-UTF-8 byte sequence on Windows error out.
    if cfg!(target_os = "windows") && value >= 128 {
        return std::ptr::null_mut();
    }

    let mut stdout = std::io::stdout().lock();

    let result = stdout.write_all(&[value]).and_then(|_| stdout.flush());

    match result {
        Err(err) => Box::into_raw(Box::new(err)),
        _ => std::ptr::null_mut(),
    }
}

unsafe extern "sysv64" fn read(buf: *mut u8) -> *mut std::io::Error {
    let mut stdin = std::io::stdin().lock();
    loop {
        let mut value = 0;
        let err = stdin.read_exact(std::slice::from_mut(&mut value));

        if let Err(err) = err {
            if err.kind() != std::io::ErrorKind::UnexpectedEof {
                return Box::into_raw(Box::new(err));
            }
            value = 0;
        }

        // ignore CR from Window's CRLF
        if cfg!(target_os = "windows") && value == b'\r' {
            continue;
        }

        *buf = value;

        return std::ptr::null_mut();
    }
}

Now we only need to call these functions in the assembly. For that we get the functions addresses, move them to a register, and do an indirect call from the register. The call will return a non-null pointer in case of an error, so we can check that and return early if necessary.

To return early we will a jump to a dynasm’s global label, denoted by ->:

b'.' => {
    dynasm! { code
        ; .arch x64
        ; mov rax, QWORD write as *const () as i64
        ; mov rdi, [r12 + r13] // cell value
        ; call rax
        ; cmp rax, 0
        ; jne ->exit
    }
}
b',' => {
    dynasm! { code
        ; .arch x64
        ; mov rax, QWORD read as *const () as i64
        ; lea rdi, [r12 + r13] // cell address
        ; call rax
        ; cmp rax, 0
        ; jne ->exit
        ; ret
    }
}

And we define that label at the end, before we pop the stack and return:

// when we push to the stack, we need to remeber
// to pop them in the opossite order.
dynasm! { code
    ; .arch x64
    ; xor rax, rax // clear `rax` in case of no early return
    ; ->exit:
    ; pop r13
    ; pop r12
    ; ret
}

Finally, we get the error back to Rust, after calling the generated code:

unsafe {
    let code_fn: unsafe extern "sysv64" fn(*mut u8) -> *mut std::io::Error =
        std::mem::transmute(buffer.as_ptr());

    let error = code_fn(self.memory.as_mut_ptr());

    if !error.is_null() {
        return Err(*Box::from_raw(error));
    }
}

And done! We have an implementation that is 100% comparable with our first interpreter. Besides, these changes are already enough to make the program compatible with Windows.

If we look at its performance:

It continues to be within the margin of error of the previous one.

An Optimized JIT compiler

And now implementing A JIT compiled version of our optimized interpreter becomes very straightforward. For that I will just copy the enum Instruction and parsing code from the interpreter, and generate the machine from that.

This implementation will no longer be, despite the title of this post, a strictly speaking single pass implementation, because we are collection all instructions before doing the compilation. For this to remain as a single-pass implementation, we would need to emit and process Instructions one at a time. This can be done, but because some instructions are made up from smaller instruction (like Clear), this would be a little trick. So let’s not focus on that.

Here is the commit where I added the enum and parsing. I also clean up the unnecessary precomputation of paring bracket, that now happens in compilation.

Now we can start compiling the instructions. The implementation of JumpRight, JumpLeft, Input and Ouput continued to be exactly the same.

For the Add(u8), we only need to use value as the operant of the instructions, instead of the hard-coded 1:

Instruction::Add(n) => dynasm! { code
    ; .arch x64
    ; add BYTE [r12 + r13], BYTE n as i8
},

When implementing the Move(usize) instruction, I notice that x86 does not have add immediate for 64-bit values, so I downgrade all usizes used in instructions to i32.

I use this other playground to get the assembly for the Move(i32):

Instruction::Move(n) => {
    if n > 0 {
        dynasm! { code
            ; lea eax, [r13 + n]
            ; add r13, -(30000 - n)
            ; cmp eax, 30000
            ; cmovl r13d, eax
        }
    } else {
        dynasm! { code
            ; lea eax, [r13 + n]
            ; add r13d, 30000 + n
            ; test eax, eax
            ; cmovns r13d, eax
        }
    }
}

This show one of the benefits of JIT compilation. At this point I know if n is positive or negative, so because I am JIT compiling, I can generate code optimized for each case.

For the Clear we simply move 0 to the cell.

Instruction::Clear => dynasm! { code
    ; .arch x64
    ; mov BYTE [r12 + r13], 0
},

For the AddTo(i32), we compute that index of the target cell similarly to the Move(i32) instruction, but with the result in rax, and finish by reading the current cell, clearing it, and adding to the target cell.

Instruction::AddTo(n) => dynasm! { code
    ; .arch x64
    // rax = cell to add to
    ;;
    if n > 0 {
        dynasm! { code
            ; lea ecx, [r13 + n]
            ; lea eax, [r13 + n - 30000]
            ; cmp ecx, 30000
            ; cmovl eax, ecx
        }
    } else {
        dynasm! { code
            ; lea ecx, [r13 + n]
            ; lea eax, [r13 + 30000 + n]
            ; test ecx, ecx
            ; cmovns eax, ecx
        }
    }
    ; mov cl, [r12 + r13]
    ; add BYTE [r12 + rax], cl
    ; mov BYTE [r12 + r13], 0
},

Here we can see that we can use ;; to inline rust code inside the dynasm macro.

And for the MoveUntil(i32) instructions we, create a loop where we check for zero, exit the loop if true, otherwise do a move and repeat.

Instruction::MoveUntil(n) => dynasm! { code
    ; .arch x64

    ; repeat:

    // check if 0
    ; cmp BYTE [r12 + r13], 0
    ; je >exit

    // Move n
    ;;
    if n > 0 {
        dynasm! { code
            ; lea eax, [r13 + n]
            ; add r13, -(30000 - n)
            ; cmp eax, 30000
            ; cmovl r13d, eax
        }
    } else {
        dynasm! { code
            ; lea eax, [r13 + n]
            ; add r13d, 30000 + n
            ; test eax, eax
            ; cmovns r13d, eax
        }
    }

    ; jmp <repeat

    ; exit:
},

Here I again copied the code from Move(i32), and also used dynasm’s local labels, denoted by : and < or > depending on if the jump is backwards or forwards.

And our optimized JIT compiler is complete! Measuring its performance:

And was expected, the performance increase is big! More than 3x faster.

About tracing JIT

Our JIT compiler implementation is very straight forward. But there is much more to JIT compilation than we have implemented here. One of them, is if compiling and optimizing the code is even worth it.

For example, there are sections in the program that are only executed once, at startup. So first compiling it to machine and then executing it, may be slower than merely interpreting it directly. Compiling it, only start to be worth if the code is executed at least a couple of times.

But again, you can also spend a lot of time trying to optimize the generated code, but it only pays off if the code is executed more than a couple of times.

A technique that takes that in account is the tracing just-in-time compilation. This technique profile how many times each loop in the program is executed, and only starting applying some levels of optimization after passing certain thresholds.

Not only that, the compiler can trace which paths in the code are being executed, and apply specific optimizations to them. It can even incorporate runtime information in the optimizations, like type information in dynamically typed languages.

Future work

We manage to create more than 26x performance increase since our first Interpreter. But we can always do better. One possibility is to investigate what choice of instructions could result in a better performance. For example, using xor eax, eax instead of xor rax, rax would result in a smaller instruction, and using mov [r12 + r13], r14 where r14 is register with the value 0 would be faster than mov [r12 + r13], 0. We could also check if multiple instructions are reading from the same cell, and optimize that to one read. Etc.

But these types of optimizations need a deep understanding of the target instructions set and how each instruction is executed by the processor (knowledge that I don’t have). And this knowledge need to exist for each architecture that your program is targeting, and there are dozens of them!

To solve that problem is why exist projects such as GCC and LLVM. They make that languages implementations just need to worry about emitting they code to these projects corresponding intermediate representation (IR), and they take the job of choosing the optimal instructions to emit. And they also take care of applying optimizations to the code.

In fact, the optimizations applied by these projects are so good that they probably could automatically apply all the optimizations that we have come up.

And this is exactly what we will explore in the next part. The original series by Eli Bendersky have made use of LLVM, but I follow a more Rusty approach and use Cranelift instead. Cranelift is a code generator written in Rust, and used by projects such as Wasmtime and rustc_codegen_cranelift. It does not apply a lot of optimizations, so maybe we don’t get it be as fast as our optimized JIT implementation (we will see). But it does solve the problem of multiple architecture targeting, so I will at least be able to run our compiler on Android devices.

All the code developed during these series can be found here (may not be exactly equal to one showed here).

Actually, this can already be greatly improved by using conditional logic instead, as we will discover later when compiling this operation. ↩
But also because by leaking the memory, it will have a 'static lifetime, and it is a necessary condition to have a safe function pointer. Later in the post I am freeing the memory, so I start using a unsafe fn pointer. ↩
The names of the registers in x64 are not very intuitive, due to backward compatibility. The same register can be accessed in different sizes, and each size has a different name. Also, some registers are named after letters, and some after numbers. ↩
If you forget to do that, the called code can misbehave, and some instructions can crash trying to access unaligned memory. Actually, at first I forget about stack alignment, and became very confused when random functions start crashing. ↩
Actually these dynasm types uses memmap2 abstractions underneath, the de facto crate that offer a safe and cross-platform API for memory mapping. ↩

Compiling Brainfuck code - Part 1: An Optimized Interpreter

2022-10-21T18:00:00+00:00

This is the first post of a blog post series where I will reproduce Eli Bendersky’s Adventures In JIT Compilation series, but this time using the Rust programming language. I will also expand this series into compiling native executables.

The goal of this series is to give a small intro into writing compilers, and explore the tools in the Rust ecosystem that can be used for that.

Like in the original series, this will be achieved by writing a compiler for the brainfuck programming language, in increasing degrees of complexity.

The Brainfuck language

Brainfuck is one of the most famous esoteric programming languages. It was first designed by Urban Müller, with the goal of being a Turing complete language whose compiler could be the smallest possible.

By being esoteric, it isn’t practical to write programs in but its simplicity makes this language one of the easiest to implement, and will be used here to lower the barrier in exploring more advanced topics in language implementation, such as machine code generation.

The language only has 6 single character instructions. A program consists of a sequence of these instructions, that act upon an array of memory cells, each initialized with the value 0. There is also a pointer, initially pointing to the first cell, that is used to interact with the memory.

The instructions are as follows:

>: move the pointer to the right.
<: move the pointer to the left.
+: increase the value at the pointer by 1.
-: decrease the value at the pointer by 1.
.: output the value at the pointer to stdout.
,: set the value at the pointer to the value read from stdin.
[: jump forward to the matching ] if the value at the pointer is 0.
]: jump backward to the matching [ if the value at the pointer is not 0.
Any other symbol is treated as a comment and ignored.

For example, to print the numbers from 1 to 5, you can do:

set cell 0 to 48 ('0' in ASCII)
++++++++++ ++++++++++ ++++++++++ ++++++++++ ++++++++
set cell 1 to 5
>+++++
increment cell 0; print it; decrement cell 1; loop while cell 1 is not 0
[<+.>-]

The language does not have an official specification, so some details may vary from implementation to implementation, like the length of the array, or what the maximum value of a memory cell is, and what happens on overflow.

I will follow the original implementation here, where the array has a length of 30,000 cells, each cell is 1 byte in size, and the value does a two-complement wrapping on overflow. But my implementation will diverge on stdin EOF behavior and the pointer will wrap around when it goes out of bounds.

Part 1: a basic interpreter

The easiest form of implementing a programming language is to write an interpreter for it. For normal languages, this can be done, for example, by reading the source code from a file, parsing it in an abstract syntax tree, and walking through that tree, executing each of its instructions.

But for brainfuck, we do not need to do any actual parsing, and if we allow brackets to not have a match, any input file will be valid brainfuck code. This means that we can execute the brainfuck code directly from the source file.

So in Rust, we can do the following:

fn main() -> std::io::Result<()> {
    let file_name = std::env::args().nth(1).unwrap();
    let source = std::fs::read(&file_name)?;

    let mut pointer = 0;
    let mut program_counter = 0;
    let mut memory = [0u8; 30_000];
    'program: loop {
        match source[program_counter] {
            b'+' => todo!(),
            b'-' => todo!(),
            b'.' => todo!(),
            b',' => todo!(),
            b'>' => todo!(),
            b'<' => todo!(),
            b'[' => todo!(),
            b']' => todo!(),
            _ => {} // do nothing, its only a comment
        }
        program_counter += 1;

        if source.len() == program_counter {
            break 'program;
        }
    }
    Ok(())
}

We read the file given by the first argument into a Vec<u8>, we define the pointer, the program counter (the index of the current instruction) and the memory, and in a loop, we execute an instruction and advance the program. When the program counter reaches the end of the source, we terminate.

Now it only remains to implement the instructions. For + and - we increment and decrement the value at the pointer, making sure to use wrapping:

b'+' => memory[pointer] = memory[pointer].wrapping_add(1),
b'-' => memory[pointer] = memory[pointer].wrapping_sub(1),

For . and , we print to stdout and read from stdin, respectively:

b'.' => print!("{}", memory[pointer] as char),
b',' => {
    use std::io::Read;
    std::io::stdin().read_exact(&mut memory[pointer..pointer + 1])?
}

For > and < we increment and decrement the pointer, making sure to wrap around, and to not underflow the usize:

b'>' => pointer = (pointer + 1) % memory.len(),
b'<' => pointer = (pointer + memory.len() - 1) % memory.len(),

[ and ] are more complicated. If we pass the instruction condition, we start traversing the program, updating deep for each [ and ] we find. When the deep reaches 0, we have found the matching bracket. We also check if the program_counter goes out of bounds (when there is no pair) and terminate the program.

Note that at the end of the loop program_counter will be on the matching bracket, which would cause an infinite loop, but this will be fixed by the increment at the end of the program loop.

b'[' => {
    if memory[pointer] == 0 {
        let mut deep = 1;
        loop {
            if program_counter + 1 == source.len() {
                break 'program;
            }
            program_counter += 1;
            if source[program_counter] == b'[' {
                deep += 1;
            }
            if source[program_counter] == b']' {
                deep -= 1;
            }
            if deep == 0 {
                break;
            }
        }
    }
}
b']' => {
    if memory[pointer] != 0 {
        let mut deep = 1;
        loop {
            if program_counter == 0 {
                break 'program;
            }
            program_counter -= 1;
            if source[program_counter] == b']' {
                deep += 1;
            }
            if source[program_counter] == b'[' {
                deep -= 1;
            }
            if deep == 0 {
                break;
            }
        }
    }
}

And we are done! That code does not handle errors graciously, and unmatching bracket simply terminates the program, but it’s only 68 lines of code!

But of course we can do much better, and not only in error handling. Brainfuck code could have a lot of comment characters, and we are traversing all of them. This is particularly bad on loops.

To solve this, we can remove the comment symbols right after reading the file:

let source: Vec<u8> = source
    .into_iter()
    .filter(|b| [b'+', b'-', b'.', b',', b'>', b'<', b'[', b']'].contains(b))
    .collect();

But we can do even better. One of the principals of Rust is to make invalid states unrepresentable. Here we have only 6 instructions, but it is being represented by a u8, that contains 256 valid states. So, it makes more sense to represent this instructions by an enum:

#[derive(PartialEq, Eq, Clone, Copy)]
enum Instruction {
    Increase,
    Decrease,
    MoveRight,
    MoveLeft,
    Input,
    Output,
    JumpRight,
    JumpLeft,
}

Then the previous code becomes:

let instructions: Vec<Instruction> = source
    .into_iter()
    .filter_map(|b| match b {
        b'+' => Some(Instruction::Increase),
        b'-' => Some(Instruction::Decrease),
        b'.' => Some(Instruction::Output),
        b',' => Some(Instruction::Input),
        b'>' => Some(Instruction::MoveRight),
        b'<' => Some(Instruction::MoveLeft),
        b'[' => Some(Instruction::JumpRight),
        b']' => Some(Instruction::JumpLeft),
        _ => None,
    })
    .collect();

Also, instead of keeping the variables of the program state in the main, we can encapsulate them in a struct.

struct Program {
    program_counter: usize,
    pointer: usize,
    instructions: Vec<Instruction>,
    memory: [u8; 30_000]
}

Then we add a method that takes care of transforming the source code in a program:

impl Program {
    fn new(source: &[u8]) -> Program {
        let instructions: Vec<Instruction> = source
            .iter()
            .filter_map(|b| match b {
                // ...
            })
            .collect();

        Program {
            program_counter: 0,
            pointer: 0,
            instructions,
            memory: [0; _],
        }
    }
}

A method for running the Program:

impl Program {
    // ...

    fn run(&mut self) -> std::io::Result<()> {
        'program: loop {
            use Instruction::*;
            match self.instructions[self.program_counter] {
                // ...
            }
            self.program_counter += 1;

            if self.instructions.len() == self.program_counter {
                break 'program;
            }
        }
    }
}

And the main becomes only:

fn main() -> std::io::Result<()> {
    let file_name = std::env::args().nth(1).unwrap();
    let source = std::fs::read(&file_name)?;

    Program::new(&source).run()
}

Then we can lock and flush stdout on output, handle stdin EOF, fix Windows issues, handle errors in a user-friendly manner and fix another Windows issue. And done! We have an idiomatic and well-written interpreter in 157 lines of Rust code.

Measuring the interpreter performance

The goal here is to progressively implement more complex versions of the interpreter/compiler, and that complexity aims to increase the runtime performance of our brainfuck programs. So it is important that we keep track of how fast our interpreter is getting.

Here we would like to have a suite of programs to run that could represent the real word distribution of brainfuck programs. But because doing so is far from trivial, I will follow the original blog post here, and measure the execution time when running a Mandelbrot generator, by Erik Bosma, and a factorization program, by Brian Raiter, with the input 179424691.

The current interpreter has run mandelbrot.bf in 96.1 seconds and factor.bf in 27.68 seconds¹².

Optimizing our interpreter

We can optimize our interpreter in multiple ways. That could be done by making some micro optimization to our interpreter code, like replacing indexing by raw pointers, looking at generated assembly for optimization opportunities, etc. But before that, we have even greater optimization opportunities.

Precomputing the paring brackets addresses.

Every time we encounter a [ or ], and pass the jump condition, we search for the paring bracket. But the pairing brackets are always the same, so there is no need to search them every time.

So the most obvious optimization that we can do now is to precompute the pairing branch of each bracket, and store them somewhere, and every time we execute a jump instruction, we look up the address of the pairing bracket.

In our case we will store the pairing bracket as a field of our Instruction enum, because it is the simpler thing to do, and later more instructions will have a field.

So our Instruction enum becomes³:

#[derive(PartialEq, Eq, Clone, Copy)]
enum Instruction {
    Increase,
    Decrease,
    MoveRight,
    MoveLeft,
    Input,
    Output,
    JumpRight(usize),
    JumpLeft(usize),
}

And the instructions’ execution will become much simpler:

use Instruction::*;
match self.instructions[self.program_counter] {
    // ...
    JumpRight(pair_address) => {
        if self.memory[self.pointer] == 0 {
            self.program_counter = pair_address;
        }
    }
    JumpLeft(pair_address) => {
        if self.memory[self.pointer] != 0 {
            self.program_counter = pair_address;
        }
    }
}
self.program_counter += 1;

Now, for computing the bracket pair we can use a stack: every time we encounter a [ we push its address to the stack (filling JumpRight with a dummy address), and every time we encounter a ] we pop from the stack, and fix both jump instructions’ addresses.

So in Program::new we replace the filter_map by:

let mut instructions = Vec::new();
let mut bracket_stack = Vec::new();

for b in source {
    let instr = match b {
        // ...
        b'[' => {
            let curr_address = instructions.len();
            bracket_stack.push(curr_address);
            // will be fixed at the pair ']'.
            Instruction::JumpRight(0)
        }
        b']' => {
            let curr_address = instructions.len();
            match bracket_stack.pop() {
                Some(pair_address) => {
                    instructions[pair_address] = Instruction::JumpRight(curr_address);
                    Instruction::JumpLeft(pair_address)
                }
                None => return Err(UnbalancedBrackets(']', curr_address)),
            }
        }
        _ => continue,
    };
    instructions.push(instr);
}

Also notice that there could be more ] than [, making the pop return None. In that case we are returning an error. Because this is the only possible error when parsing brainfuck code, I am defining the error as a simple tuple struct instead of a enum, as it is usually done:

struct UnbalancedBrackets(char, usize);

And the signature of Program::new becomes:

impl Program {
    fn new(source: &[u8]) -> Result<Program, UnbalancedBrackets> {
        // ...
        Ok(Program {
            program_counter: 0,
            pointer: 0,
            instructions,
            memory: [0; 30_000],
        })

    }
}

Ah, and we also cannot forget that there could also be more [ than ]:

let mut instructions = Vec::new();
let mut bracket_stack = Vec::new();

for b in source {
    // ...
    instructions.push(instr);
}

if let Some(unpaired_bracket) = bracket_stack.pop() {
    return Err(UnbalancedBrackets('[', unpaired_bracket));
}

And finally we need to handle that error in main:

fn main() -> ExitCode {
    let mut args = std::env::args();
    if args.len() != 2 {
        eprintln!("expected a single file path as argument");
        return ExitCode::from(1);
    }

    let file_name = args.nth(1).unwrap();
    let source = match std::fs::read(&file_name) {
        Ok(x) => x,
        Err(err) => {
            eprintln!("Error reading '{}': {}", file_name, err);
            return ExitCode::from(2);
        }
    };

    let mut program = match Program::new(&source) {
        Ok(x) => x,
        Err(UnbalancedBrackets(c, address)) => {
            eprintln!(
                "Error parsing file: didn't find pair for `{}` at instruction index {}",
                c, address
            );
            return ExitCode::from(3);
        }
    };

    if let Err(err) = program.run() {
        eprintln!("IO error: {}", err);
    }

    ExitCode::from(0)
}

It would be more useful to return the line and column of the bracket (at least the byte index would be better), but this is good enough for our needs.

Running the programs again:

A decrease of 25%! But we can do better.

Replace repeating operations by a single one

If you take a look at the source code of mandelbrot.bf, or of factor.bf, or of the example hello-word at the start of this page, you will see that there are a lot of instructions repeating in sequence.

Let start with increment. So instead of repeating an increment n times, we can replace them by a single add n instruction. Even better, we can replace both increment and decrement instructions by a single Add(u8) instruction, because a subtraction by n can be represented as an addition via the two-complement on n (or n.wrapping_neg() in our case).

So Instruction becomes:

#[derive(PartialEq, Eq, Clone, Copy)]
enum Instruction {
    Add(u8),
    MoveRight,
    MoveLeft,
    Input,
    Output,
    JumpRight(usize),
    JumpLeft(usize),
}

The instruction interpreting will become:

use Instruction::*;
match self.instructions[self.program_counter] {
    Add(n) => self.memory[self.pointer] = self.memory[self.pointer].wrapping_add(n),
    // ...
}
self.program_counter += 1;

And the parsing:

for b in source {
    let instr = match b {
        b'+' | b'-' => {
            let inc = if *b == b'+' { 1 } else { 1u8.wrapping_neg() };
            if let Some(Instruction::Add(value)) = instructions.last_mut() {
                *value = value.wrapping_add(inc);
                continue;
            }
            Instruction::Add(inc)
        }
        // ...
    };
    instructions.push(instr);
}

Running it:

Well… it was not as big of a reduction as I was expecting, only 0.26%. 0.26 ± 0.37 percent to be more precise. It is even uncertain that it really decreased.

But let’s do the same for the Move instructions:

#[derive(PartialEq, Eq, Clone, Copy)]
enum Instruction {
    Add(u8),
    Move(isize)
    Input,
    Output,
    JumpRight(usize),
    JumpLeft(usize),
}

use Instruction::*;
match self.instructions[self.program_counter] {
    // ...
    Move(n) => {
        let len = self.memory.len() as isize;
        let n = (len + n % len) as usize;
        self.pointer = (self.pointer + n) % len as usize;
    }
    // ...
}
self.program_counter += 1;

for b in source {
    let instr = match b {
        // ...
        b'>' | b'<' => {
            let inc = if *b == b'>' { 1 } else { -1 };
            if let Some(Instruction::Move(value)) = instructions.last_mut() {
                *value += inc;
                continue;
            }
            Instruction::Move(inc)
        }
        // ...
    };
    instructions.push(instr);
}

Running again:

Wow! This is what I was looking for, more than 3 times in performance increase!

But now, there is the question: Why does optimizing the increment/decrement result in such a small improvement, but optimizing the move instructions causes such a big change? Well, if I had paid more attention to the source code of mandelbrot.bf and of factor.bf, I could have noticed that there are much more ‘>’s and ‘<’s than ‘+’s and ‘-‘s:

	mandelbrot.bf	factor.bf
+	694	422
-	591	370
>	4,428	1,324
<	4,363	1,294
[	686	231
]	686	231
.	3	7
,	0	1

But what’s really important is the number of times these instructions are executed at runtime. In that case we can modify our interpreter to count the number of instructions being executed:

	mandelbrot.bf	factor.bf	mandelbrot.bf	factor.bf
+	179,053,599	212,428,900	351,460,869	396,481,127
-	177,623,022	212,328,376	0	0
>	4,453,036,023	1,220,387,724	1,408,648,727	484,719,480
<	4,453,036,013	1,220,387,704	0	0
[	422,534,152	118,341,127	422,534,152	118,341,127
]	835,818,921	242,695,606	835,818,921	242,695,606
,	6,240	21	6,240	21
,	0	10	0	10
Total	10,521,107,970	3,226,569,468	3,018,468,909	1,242,237,371

Now, with this data, everything becomes much clearer. The Add optimization has decreased the total number of operation by only 0.24%, and the Move optimization has decreased it by 69%!

But, of course, we can always do better.

Optimizing higher constructs

So we have already implemented an optimization that changes a single instruction, and one that replaces a repeated instruction by a single one. The next step is to replace a set of distinct instructions by a single one.

In brainfuck, there are a lot of basic operations that need to involve multiple instructions. For example, clearing a cell ([-] or [+]), adding a value of one cell to another (add to the third right would be [->>>+<<<], but this clears the original value), duplicating a cell ([->>>+>+<<<<]), etc.

Many of these operations are basically loops with instructions inside. In this case, we can profile our programs by counting how many times a given loop is run, which boils down to counting how many times each particular ] is executed. And then we present this data by deduplicating identical loops and displaying them in descending order.

Running the new profiling for factor.bf:

	count	loop
1	32,276,219	[-1<10+1>10]
2	28,538,377	[-1]
3	15,701,515	[-1<4+1>4]
4	12,581,941	[-1>3+1>1+1<4]
5	9,579,970	[-1>3+1>2+1<5]
6	9,004,028	[-1<3+1>3]
7	6,093,976	[-1<1-1>1]
8	6,085,735	[-1>3+1<3]
9	5,853,530	[-1<1+1<3+1>4]
10	5,586,229	[-1>3+2<3]

In first place comes the loop [-<<<<<<<<<<+>>>>>>>>>] (the annotation above is in compressed format), the “add cell to another” operation, running 32 million of times! In fact, positions 3, 6, 7, 8 and 10 are also the same operation. Position 2 is the “clear cell operation” ([-]), and position 4, 5 and 9 is the “duplicate cell” ([->>>+>+<<<<]).

For the mandelbrot.bf:

	count	loop
1	287,432,488	[>9]
2	200,272,618	[<9]
3	116,145,344	[>1[-1>9+1<9]<10]
4	32,021,044	[-1>9+1<9]
5	31,339,760	[>2[-1>9+1<9]<11]
6	12,637,333	[-1>2[-1<2+1>2]<2[-1>2+1>2+1<4]+1>9]
7	12,038,491	[-1]
8	11,813,904	[-1>2[-1<2+1>2]<2[-1>2+1>1+1<3]+1>9]
9	9,515,168	[>1+1>8]
10	9,017,333	[-1<4+1>1[<1-1>1-1<6+1>6]<1[-1>1+1<1]>4]

The mandelbrot one was less predictable. In the first and second place there is the “move 9 cell right/left repeatedly until it encounters a 0”. In 4, it is the “add cell to another”. 7 is the “clear cell” operation.

The other loops appear to be doing more specific operations (which may be giving a glimpse on how the mandelbrot algorithm works). 3, which is the outer loop of 4, appears to be shifting the elements of an array, where each element contains 9 cells. 6 and 8 appear to be duplicating some cell for each element of the array, using a temporary cell for that. 10 is more complicated but appears to be another operation in the array.

Well, each of the operations described above can be executed by the computer in a much more efficient way if they are implemented by a “real” programming language. So, we can identify some of these operations in the source code, and replace them by a single instruction implemented in Rust.

In this case, I will follow the original blog posts, and implement instructions for the “clear cell”, “add cell to another” and “move until 0” operators.

Let’s start with the “clear cell”:

#[derive(PartialEq, Eq, Clone, Copy)]
enum Instruction {
    // ..
    Clear,
}

Its implementation is simple enough:

'program: loop {
    use Instruction::*;
    match self.instructions[self.program_counter] {
        // ...
        Clear => self.memory[self.pointer] = 0,
    }
    // ...
}

And every time we parse a ] we check if the previous instructions match the “clear cell” operation, and replace it by the new instruction.

b']' => {
    let curr_address = instructions.len();
    match bracket_stack.pop() {
        Some(pair_address) => {
            instructions[pair_address] = Instruction::JumpRight(curr_address);

            use Instruction::*;
            match instructions.as_slice() {
                // could enter an infinite loop if n is even.
                [.., JumpRight(_), Add(n)] if n % 2 == 1 => {
                        let len = instructions.len();
                        instructions.drain(len - 2..);
                        Instruction::Clear
                    }
                _ => Instruction::JumpLeft(pair_address),
            }
        }
        None => return Err(UnbalancedBrackets(']', curr_address)),
    }
}

Running it:

A mean decrease of 4.5% in execution time. And it decreased the total number of instructions by 0.80% and 4.59%, for mandelbrot.bf and factor.bf, respectively.

Now for the “add to another”:

#[derive(PartialEq, Eq, Clone, Copy)]
enum Instruction {
    // ..
    AddTo(isize),
}

'program: loop {
    use Instruction::*;
    match self.instructions[self.program_counter] {
        // ...
        AddTo(n) => {
            let len = self.memory.len() as isize;
            let n = (len + n % len) as usize;
            let to = (self.pointer + n) % len as usize;

            self.memory[to] = self.memory[to].wrapping_add(self.memory[self.pointer]);
            self.memory[self.pointer] = 0
        }
    }
    // ...
}

match instructions.as_slice() {
    // ...
    &[.., JumpRight(_), Add(255), Move(x), Add(1), Move(y)]
        if x == -y =>
    {
        let len = instructions.len();
        instructions.drain(len - 5..);
        Instruction::AddTo(x)
    }
    _ => Instruction::JumpLeft(pair_address),
}

This one gives more substantial performance increase. A mean decrease in run-time by 15.3%. Only 2.56% on mandelbrot, but 28.13% on factor!

And finally the “move until 0” operator:

#[derive(PartialEq, Eq, Clone, Copy)]
enum Instruction {
    // ..
    MoveUntil(isize),
}

'program: loop {
    use Instruction::*;
    match self.instructions[self.program_counter] {
        // ...
        MoveUntil(n) => {
            let len = self.memory.len() as isize;
            let n = (len + n % len) as usize;
            loop {
                if self.memory[self.pointer] == 0 {
                    break;
                }

                self.pointer = (self.pointer + n) % len as usize;
            }
        }
    }
    // ...
}

match instructions.as_slice() {
    // ...
    &[.., JumpRight(_), Move(n)] => {
        let len = instructions.len();
        instructions.drain(len - 2..);
        Instruction::MoveUntil(n)
    }
    _ => Instruction::JumpLeft(pair_address),
}

As expected from the previous profiles, this optimization has greater impact on the mandelbrot.bf. Time goes down by 18.61% on mandelbrot and 3.79% on factor!

About Interpreters and Compilers

What we basically have done through this post, is to transform the program source code into instructions that can be interpreted in a more efficient way. We also looked for common patterns in the code and created specific instructions for them, increasing the performance even more.

These are already basic ideas of what interpreters and compilers do. Instead of interpreting/compiling the code directly from its source code (like how AST-walking interpreters, or single pass compilers do), they most of the time convert them to an intermediate representation (IR) that is more suitable for interpreting, compiling, applying optimizations, etc.

For example, in Lua an AST is not even created, the source code is directly compiled, in a single pass, to bytecode, which is the IR used by the Lua interpreter.

A more complex example would be Rust. In Rust, the code AST is first converted to a High-Level Intermediate Representation (process called HIR lowering), which is used for type inference, trait solving and type checking. After that, it is lowered to Mid-Level Intermediate Representation (MIR), which is used for borrow checking, applying optimizations and monomorphization collecting. And finally, MIR is converted to LLVM-IR, which is passed to LLVM to apply more optimizations and only then compiled down to machine code. More details can be found here.

Future work

From the first implementation to the last one we get a speedup of almost 7x. Of course, we can always do better. We could continue implementing specific instructions for common brainfuck operations, and we could do that on and on, without end (we could even make an instruction that prints the factors of a number!), but they would less and less generalize for all programs.

We could also start searching for micro optimizations, like inspecting the generated assembly for lost optimization opportunities, and modify the code to reach them, etc. We can also go unsafe and remove bound checks with get_unchecked and unreachable_unchecked⁴.

But this being a normal interpreter, there will always be a significant overhead introduced by the interpreter loop that cannot be removed (unless we manage to transform the source code to a single instruction).

So we will stop enhancing this interpreter for now, and, in the next part, we will build a JIT interpreter, that will convert our source code directly to machine code, and remove the interpreter loop altogether.

All the code developed during these series can be found here (may not be exactly equal to one showed here).

You can find all measures in detail, including method and environment, in this document. ↩
The implementation in the original blog post have taken, respectively, 64 seconds and 24 seconds on my machine. But my implementation is doing extra things like memory wrap-around and (unnecessary) bound checks, which may be causing this big difference. ↩
Note that here I am using a usize for the index of the pairing bracket, which has 8 bytes on a 64-bit system. This is causing the Instruction enum to have a size of 16 bytes (due to the discriminant and alignment). So a possible optimization could be to use a u32 or a u16 here, depending on how big brainfuck programs your interpreter needs to support. ↩
In fact, I have done some of these changes here, but didn’t carefully profile it. ↩