Brainfuck JIT Compiler in Rust

Hello JIT


JIT不是一个神秘的玩意。

—— Tondbal ik Ni

我们都知道,对于解释型的语言实现来说,性能是大家关注的焦点。比如,这位 Tondbal ik Ni 曾经还说过:

P*没上JIT,慢的一逼!

—— Tondbal ik Ni

似乎这句话总是隐含着另一层意思:实现JIT,难!而当我们再一联想到JVM这种庞然大东西的时候,很自然的就

然而!JIT原理并不复杂,做出一个玩具JIT Compiler更是非常轻松。之所以JVMs那么庞大而复杂,原因之一在于它们做了大大大量的优化工作。

我们今天就要来看看JIT究竟是个什么东西!

Just-in-Time Compiler


JIT Compiler,究其根本还是一个Compiler。一个Ahead-of-Time Compiler的编译过程往往会有这些(既不充分也不必要的)步骤:

  • 词法分析
  • 语法分析
  • 语法制导定义或翻译
  • 中间代码生成
  • 代码优化
  • 目标代码生成

对于解释器来说,往往将编译工作进行到中间某一步后就直接进行解释执行了,并不生成目标代码。而JIT Compiler却是要生成目标代码的,最终执行的是编译好后的Native Code。只不过,它将目标代码生成的部分推迟到了执行期才进行。这样的好处有:

  • 无需重新编译就可以实现跨平台
    参考Java,它将平台差异隔离在了中间代码部分(指Java ByteCode)。
  • 运行时优化
    当年大欢还在用Gentoo的时候曾经开过嘲讽:本地编译,优化开的细,比你Arch强多了(然而后来还是倒戈到了Arch 666666)而一个JIT Compiler不仅知道目标平台的信息,更知道程序的运行时信息,因此往往可以有更多的优化。
  • 解释/编译混合
    这其实也可以看作是一种优化措施,即执行次数多的代码JIT编译后执行,执行次数少的代码解释执行。因为JIT还是需要一步编译的过程,如果代码执行次数少,很可能抵消不了编译过程带来的时间开销。

所以,其实优化是JIT Compiler中相当重要的一部分。如果我们不要优化,那可是简单了很多哟。

Core of JIT


如果你能看懂这段代码,那就说明你已经掌握了JIT的精髓了:

#include <stdio.h> 
#include <string.h>
#include <sys/mman.h> 

char f[] = {0x48, 0xC7, 0xC0, 0x00, 0x00, 0x00, 0x00, 0xC3}; 

int main(){  
    int a = 5; memcpy(&f[3], &a, 4); 
    char* mem = mmap(NULL, sizeof(f), PROT_WRITE | PROT_EXEC, MAP_ANON | MAP_PRIVATE, -1, 0); 
    memcpy(mem, f, sizeof(f)); 
    printf("%dn", ((int (*)())mem)()); munmap(mem, sizeof(f)); 
    return 0;
 }

其中mmap的作用是申请按照PageSize对齐的内存,并可以修改页面保护属性。

所以一个JIT Compiler编译的主要步骤就是:

  • 申请一块可写、并且将来可以改成可执行的内存。
  • 生成机器代码,写入内存。
  • 修改内存页面属性,Call it!

So easy,下面进入脑残环节。

An interpreter for Brainf*ck


我们将实现一个Brainfuck的解释器,随后再实现一个JIT编译器。之所以选择Brainfuck,自然是因为它相当简单,完全可以当做中间代码进行处理,省去了词法语法分析、中间代码生成等与编译原理直接相关的部分。

解释器写起来就太简单了。Brainfuck预设我们有一块无限大的内存,初值均为0。数据指针指向第一个字节。它支持8种操作:

  • >
    数据指针右移一格
  • <
    数据指针左移一格
  • +
    数据指针指向的内存内容+1
  • -
    数据指针指向的内存内容-1
  • .
    putchar
  • ,
    getchar
  • [
    如果当前内存内容 == 0,则向后跳转至对应的]
  • ]
    如果当前内存内容 != 0,则向前跳转至对应的[

翻译器部分可以作为大一的C语言实验哈哈哈哈

A JIT Compiler for Brainf*ck


如果要手撸JIT Compiler,则需要对目标平台有一定的了解。我们这里的目标平台是x86_64,这个网址 可以在线将汇编生成为x86或x64的机器代码。

第一步:申请PageSize对齐的内存
这一步除了使用mmap,还可以使用posix_memalign。

第二步:生成函数框架
在这里我们将一整个Brainfuck程序编译成一个函数。这个函数接受一个参数,即我们事先申请好的一块内存作为数据区域。对于x64来说,Linux等类Unix系统遵循的调用约定是System V AMD64 ABI,函数的第一个参数由Rdi传递。因此我们生成的函数的开始与结束部分如下:

pub fn compile_and_run(&mut self, code: &str) {  
    self.jit_code.emit_push_r(Reg::Rbp); 
    self.jit_code.emit_mov_rr(Reg::Rbp, Reg::Rsp); 
    self.emit_push_regs(); //Save regs 
    self.jit_code.emit_mov_ri(Reg::Rbx, 0); //Rbx = data_pointer 
    self.jit_code.emit_mov_rr(Reg::Rdx, Reg::Rdi); //Rdx = memory_base 

    //more code here... 

    self.emit_pop_regs(); //Load regs that saved before 
    self.jit_code.emit_mov_rr(Reg::Rsp, Reg::Rbp);
    self.jit_code.emit_pop_r(Reg::Rbp);     self.jit_code.emit_ret(); 
}

上面的代码中,各个emit函数的作用是生成相应的机器代码,放入内存中。例如emit_push_r(Reg::Rbp)将生成机器码0x55,它对应的汇编为push rbp

接下来就是根据Brainfuck各个操作生成机器码了。例如>操作:

self.jit_code.emit_inc_r(Reg::Rbx);  

So easy. 需要额外说明的只有两点:

一是我们可能需要重定位操作。当我们的buffer空间不够的时候,需要对其进行扩大,这样的话我们代码所在的 地址就会变动。而有一些指令(比如Relative跳转、RelativeCall等)它的操作数是当前RIP(即程序计数器PC)与目标地址的 Offset,这就需要当我们最终结束生成这个函数时,再对这些指令的操作数进行计算。

二是对于[操作来说,需要一个patch back的过程。当我们在编译过程中遇到它的时候,我们并不知道跳转的目的地址是哪里。只有在遇到对应的]时,才能更新它的跳转地址。

第三步:Call it!

let mut memory: Vec<u8>= vec![0; 10000];  
let func: extern "C" fn(*mut u8) = unsafe { mem::transmute(self.jit_code.function()) };  
func(memory.as_mut_ptr());  

Performance Comparison


我找了一个用Brainfuck写的Mendelbrot程序进行测试,它会在控制台输出Mendelbrot的ASCII图(大神请受我一拜 )。除了上面自己实现的解释器和JIT编译器外,我还找了一个Brainfuck的编译器bfc进行测试。

测试结果大致如下:

Interpreter: ~1min
JIT: ~4.31s
bfc: ~4.21s

代码请参 https://github.com/dontpanic92/rainfuck

最后,由于将汇编翻译为机器码是一件体力活,我们还可以使用一些现成的工具。例如DynASM,通过预处理的方式将C + ASM混合代码处理为C语言代码(即省去了我们显示emit的部分)。或者,也可以考虑使用LibJIT或LibGCCJIT等库。

Reference


  1. http://blog.reverberate.org/2012/12/hello-jit-world-joy-of-simple-jits.html
  2. http://www.hydrocodedesign.com/2014/01/17/jit-just-in-time-compiler-rust/
  3. http://www.jonathanturner.org/2015/12/building-a-simple-jit-in-rust.html