写一个 JIT 编译器

Posted on Jun 11, 2021

JIT

JIT 是什么？

JIT 是 Just In Time 的意思，按我的理解就是在运行时编译，把相关代码在运行的时候编译成汇编(机器码)，再申请一段可执行的内存，把这段机器码放入内存中，再去执行这段内存。

简单演示

来看一段代码

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23


#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <sys/mman.h>
#include <unistd.h>

int main(int argc, char* argv[])
{
    unsigned char code[] = "\xb8\x00\x00\x00\x00\xc3";
    if (argc < 2) {
        printf("Usage: %s <integer>\n", argv[0]);
        return -1;
    }
    int num = atoi(argv[1]);
    memcpy(&code[1], &num, 4);

    void* mem = mmap(NULL, sizeof(code), PROT_READ | PROT_WRITE, MAP_ANON | MAP_PRIVATE, -1, 0);
    memcpy(mem, code, sizeof(code));
    mprotect(mem, sizeof(code), PROT_READ | PROT_EXEC);
    int (*func)() = reinterpret_cast<int (*)()>(mem);
    printf("Your number was: %d\n", func());
    return 0;
}

代码里边的 code 就是我们的机器码

1
2


mov eax, 0
ret

我们得到程序第一个参数，把他转换为整数，然后将整数赋值到机器码中的 0 的那个部分。

容易看出，b8 是 mov eax, 的机器码，c3 则是 ret，而剩下的 4 组 0 就是我们填充的整数。

最后申请一块内存，把机器码复制进去，然后执行。

Barinfuck

brainfuck 是一个很小且图灵完备的语言，只有 8 条指令，分别是

下面是这八种状态的描述，其中每个状态由一个字符标识：

字符	含义
>(PtrAdd)	指针加 1
<(PtrSub)	指针减 1
+(ValAdd)	指针所指向的内存的值加 1
-(ValSub)	指针所指向的内存的值减 1
.(PutChar)	将指针指向的内存内容输出
,(ReadChar)	输入内容到指针指向的内存
[(LBracket)	如果指针指向的单元值为零，向后跳转到对应的`]`指令的下一条指令
](RBracket)	无条件跳转到 `[`

根据这几条简单的指令，我们可以很快的写出 Brainfuck 的解释器。

Brainfuck 解释器只需要一个几个字段即可完成：

一个 int8 的数组
一个程序计数器 pc
一个指向当前数组下标的计数器 sp

Hello World!

一个 Hello World! 的例子

1
2
3


++++++++++[>+++++++>++++++++++>+++>+<<<<-]
>++.>+.+++++++..+++.>++.<<+++++++++++++++.
>.+++.------.--------.>+.>.

把这段代码翻译一下就可以得到类似于下面的"指令集"

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14


ValAdd, ValAdd, ValAdd, ValAdd, ValAdd, ValAdd, ValAdd, ValAdd, 
ValAdd, ValAdd, LBracket, PtrAdd, ValAdd, ValAdd, ValAdd, ValAdd, 
ValAdd, ValAdd, ValAdd, PtrAdd, ValAdd, ValAdd, ValAdd, ValAdd, 
ValAdd, ValAdd, ValAdd, ValAdd, ValAdd, ValAdd, PtrAdd, ValAdd, 
ValAdd, ValAdd, PtrAdd, ValAdd, PtrSub, PtrSub, PtrSub, PtrSub, 
ValSub, RBracket, PtrAdd, ValAdd, ValAdd, PutChar, PtrAdd, ValAdd, 
PutChar, ValAdd, ValAdd, ValAdd, ValAdd, ValAdd, ValAdd, ValAdd, 
PutChar, PutChar, ValAdd, ValAdd, ValAdd, PutChar, PtrAdd, ValAdd, 
ValAdd, PutChar, PtrSub, PtrSub, ValAdd, ValAdd, ValAdd, ValAdd, 
ValAdd, ValAdd, ValAdd, ValAdd, ValAdd, ValAdd, ValAdd, ValAdd, 
ValAdd, ValAdd, ValAdd, PutChar, PtrAdd, PutChar, ValAdd, ValAdd, 
ValAdd, PutChar, ValSub, ValSub, ValSub, ValSub, ValSub, ValSub, 
PutChar, ValSub, ValSub, ValSub, ValSub, ValSub, ValSub, ValSub, 
ValSub, PutChar, PtrAdd, ValAdd, PutChar, PtrAdd, PutChar,

可以发现，这些代码存在很多的重复的部分，我们可以为这些指令集添加一个字段 times 用来表示每次同一个操作执行的次数，就像这样。

1
2
3
4
5
6


ValAdd(10), LBracket(13), PtrAdd(1), ValAdd(7), PtrAdd(1), ValAdd(10), PtrAdd(1), ValAdd(3), 
PtrAdd(1), ValAdd(1), PtrSub(4), ValSub(1), RBracket(1), PtrAdd(1), ValAdd(2), PutChar(), 
PtrAdd(1), ValAdd(1), PutChar(), ValAdd(7), PutChar(), PutChar(), ValAdd(3), PutChar(), 
PtrAdd(1), ValAdd(2), PutChar(), PtrSub(2), ValAdd(15), PutChar(), PtrAdd(1), PutChar(), 
ValAdd(3), PutChar(), ValSub(6), PutChar(), ValSub(8), PutChar(), PtrAdd(1), ValAdd(1), 
PutChar(), PtrAdd(1), PutChar(),

明显的，指令少了很多，执行速率也会加快。

这个优化器也很容易编写。

JIT

终于来到了我们的主题，JIT。

这里我们主要用到了以下几个寄存器：

rbp 就是上文的 sp
rdi，rsi，rdx，用于函数调用的参数
rax，存储系统调用的标号

为了方便，我直接把数组的长度申请了 3w，也就是 0x7530

汇编分为几个部分

保存现场
根据前面指令集生成的汇编
恢复现场

保存现场这个我们只用到了 rbp，所以保存 rbp 就够了

1
2
3


push rbp
mov rbp, rsp
sub rbp, 0x7530

恢复现场也很简单

1
2


pop rbp
ret

由于我是在 Linux 下编写的。

我们可以看看系统调用表，这虽然是 Chromium OS 的，但是对于我们用到的前两个系统调用 (read, write) 来说是一样的，可以凑合着看看。

>(PtrAdd)

add rbp, 0

<(PtrSub)

sub rbp, 0

+(ValAdd)

add byte ptr[rbp], 0

-(ValSub)

sub byte ptr[rbp], 0

.(PutChar)

1
2
3
4
5


mov rdi, 1
lea rsi, [rbp]
mov rax, 1
mov rdx, 1
syscall

,(ReadChar)

1
2
3
4
5


mov rax, 0
mov rdi, 0
lea rsi, [rbp]
mov rdx, 1
syscall

[(LBracket)

1
2


cmp byte ptr [rbp], 0
jne ...

](RBracket)

jmp ...

很简单对吧，我们只需要把这些汇编翻译成机器码，然后在对应的位置补上我们需要的值，就可以了！

我是使用这个网站在线将汇编转为机器码的

http://shell-storm.org/online/Online-Assembler-and-Disassembler/

上述是 x86_64 的汇编，在转换的时候记得选上。

值得注意的是，在写入参数的时候要注意字节序，我这里是小端序。