使用Rust实现Brainfuck的JIT编译器
[TOC]
使用Rust实现Brainfuck的JIT编译器
x64汇编简介
Linux x64 汇编/Hello World
我们每天产出大量的垃圾代码，我们每个人都可以像这样简单地编写最简单的代码：
#include <stdio.h>
int main() {
int x = 10;
int y = 100;
printf("x + y = %d\n"，x + y);
return 0;
}
希望读者们都可以理解上述 C 代码的作用。但是，此代码在底层如何工作？我认为并非所有人都能回答这个问题，我也是。我可以用Haskell，Erlang，Go 等高级编程语言编写代码，但是在它们编译后我并不知道它在底层是如何工作的。因此，我决定采取一些更深入的步骤，进行记录，并描述我对此的学习过程。希望这个过程不仅仅只是对我来说很有趣。让我们开始吧。
准备
开始之前，我们必须准备一些事情，如我所写的那样，我目前使用 Ubuntu 18.04，因此我的文章将针对该操作系统和体系结构。不同的 CPU 支持不同的指令集，目前我使用 Intel 的 64 位 CPU。同时我也将使用 NASM 语法。你可以使用以下方法安装它：
$ apt install nasm
记住，Netwide Assembler（简称 NASM）是一款基于英特尔 x86 架构的汇编与反汇编工具。这就是我们目前需要的工具。
NASM 语法
在这里，我将不介绍完整的汇编语法，我们仅提及其庞大语法的一小部分，也是那些我们将在本文中使用到的部分。通常， NASM 程序分为几个段（section），在这篇文章中，我们将遇到以下两个段：
- 数据段：data section
- 文本段：text section
数据段部分用于声明常量，此数据在运行时不会更改。声明数据部分的语法为：
section .data
文本段部分用于代码。该部分必须以全局声明 _start 开头，该声明告诉内核程序从何处开始执行。
section .text
global _start
_start:
注释以符号 ; 开头。每条 NASM 源代码行都包含以下四个字段的某种组合：
[label:] instruction [operands] [; comment]
方括号中的字段是可选的。基本的 NASM 指令由两部分组成，第一部分是要执行的指令的名称，第二部分是该命令的操作数。例如：
mov rax, 48 ; put value 48 in the rax register
Hello World!
让我们用 NASM 编写第一个程序。当然，这将是经典的 Hello World! 程序。这是它的代码：
section .data
msg db "Hello World!", 0x0A
section .text
global _start
_start:
mov rax, 1
mov rdi, 1
mov rsi, msg
mov rdx, 13
syscall
mov rax, 60
mov rdi, 0
syscall
是的，它看起来一点都不像 printf("Hello World!\n")。让我们尝试了解它是什么以及它如何工作。首先看第一和第二行，我们定义了数据段部分，并将 msg 常量与 “Hello, World!” 值放在一起。现在，我们可以在代码中使用此常量。接下来是声明文本段部分和程序的入口。程序将从第 7 行开始执行。现在开始最有趣的部分，我们已经知道 mov 指令是什么，它获得 2 个操作数，并将第二个的值放在第一位。但是这些 rax, rdi 等是什么？正如我们在 Wikipedia 中可以看到的：
中央处理器（CPU）是计算机中的硬件，它通过执行系统的基本算术，逻辑和输入/输出操作来执行计算机程序的指令。
好的，CPU 会执行一些运算。但是，在哪里可以获取该运算的数据，是内存吗？从内存中读取数据并将数据写回到内存中会减慢处理器的速度，因为它涉及通过控制总线发送数据请求的复杂过程。因此，CPU 具有自己的内部存储器，称为寄存器。
因此，当我们编写 mov rax, 1 时，意味着将 1 放入 rax 寄存器。现在我们知道 rax，rdi，rsi 等代表了什么了，但是还需要知道什么时候该使用 rax ，什么时候使用 rdi 等。
-
rax：临时寄存器，当我们调用syscall时，rax必须包含syscall号码，所以后面的数字就是syscall的号码 -
rdi：用于将第 1 个参数传递给函数 -
rsi：用于将第 2 个参数传递给函数 -
rdx：用于将第 3 个参数传递给函数
换句话说，我们只是在调用 sys_write 的 syscall。看看 sys_write 的定义：
size_t sys_write(unsigned int fd, const char * buf, size_t count);
它具有3个参数：
-
fd：文件描述符。对于 stdin，stdout 和 stderr 来说，其值分别为 0，1 和 2 -
buf：指向字符数组 -
count：指定要写入的字节数
我们将 1 写入 rax，这意味我们要调用 sys_write。完整的 syscall 列表可以在 https://github.com/torvalds/linux/blob/master/arch/x86/entry/syscalls/syscall_64.tbl 找到。在完成该调用之后，将 60 写入 rax，这意味着我们要调用 sys_exit 退出程序，且退出码为 0。
最后，让我们来构建这个程序，我们需要执行以下命令：
$ nasm -f elf64 -o main.o main.asm
$ ld -o main main.o
尝试运行这个程序吧！
什么是JIT
本文部分翻译自：https://blog.reverberate.org/2012/12/hello-jit-world-joy-of-simple-jits.html.
“JIT” 一词往往会唤起工程师内心最深处的恐惧和崇拜，通常这并没有什么错，只有最核心的编译器团队才能梦想创建这种东西。它会使你联想到 JVM 或 .NET，这些家伙都是具有数十万行代码的超大型运行时。你永远不会看到有人向你介绍 “Hello World!” 级别的 JIT 编译器，但事实上只需少量代码即可完成一些有趣的工作。本文试图改变这一点。
编写一个 JIT 编译器只需要四步，就像把大象装到冰箱里一样。
- 申请一段可写和可执行的内存
- 将源码翻译为机器码（通常经过汇编）
- 将机器码写入第一步申请的内存
- 执行这部分内存
Hello, JIT World: The Joy of Simple JITs
事不宜迟，让我们跳进我们的第一个 JIT 程序。该代码是特定于 64 位 Unix 的，因为它使用了 mmap()。鉴于此读者需要拥有支持该代码的处理器和操作系统。
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/mman.h>
int main(int argc, char *argv[]) {
// Machine code for:
// mov eax, 0
// ret
unsigned char code[] = {0xb8, 0x00, 0x00, 0x00, 0x00, 0xc3};
if (argc < 2) {
fprintf(stderr, "Usage: jit1 <integer>\n");
return 1;
}
// Overwrite immediate value "0" in the instruction
// with the user's value. This will make our code:
// mov eax, <user's value>
// ret
int num = atoi(argv[1]);
memcpy(&code[1], &num, 4);
// Allocate writable/executable memory.
// Note: real programs should not map memory both writable
// and executable because it is a security risk.
void *mem = mmap(NULL, sizeof(code), PROT_WRITE | PROT_EXEC,
MAP_ANON | MAP_PRIVATE, -1, 0);
memcpy(mem, code, sizeof(code));
// The function will return the user's value.
int (*func)() = mem;
return func();
}
似乎很难相信上面的 33 行代码是一个合法的 JIT。它动态生成一个函数，该函数返回运行时指定的整数，然后运行该函数。读者可以验证其是否正常运行：
$ gcc -o jit jit.c
$ ./jit 42
$ echo $?
# 42
您会注意到，代码中使用 mmap() 分配内存，而不是使用 malloc() 从堆中获取内存的常规方法。这是必需的，因为我们需要内存是可执行的，因此我们可以跳转到它而不会导致程序崩溃。在大多数系统上，栈和堆都配置为不允许执行，因为如果你的代码跳转到了栈或堆，则意味着程序发生了很大的错误，这是由操作系统的内存结构决定的。更糟糕的是，利用缓冲区溢出的黑客可以使用可执行堆栈来更轻松地利用该漏洞。因此，通常我们希望避免映射任何可写和可执行的内存，这也是在你自己的程序中遵循此规则的好习惯。我在上面打破了这个规则，但这只是为了使我们的第一个程序尽可能简单。
dynasm介绍
DynASM 是为 LuaJIT 编写的 JIT 汇编预处理器和微型运行时库。简单来讲，DynASM完成两个工作，一个是预处理，把你写的汇编指令（对，没有Elixir，DynASM并不能直接把逻辑变成汇编，需要你手动把你的逻辑用汇编语言重写一遍，因此性能也取决于你的汇编代码写的好坏）变成真正的二进制机器码，另一个是提供一个微型运行时，来处理那些必须推迟到运行时才能执行的代码。
而 Rust 生态中也有一个参照 DynASM 所开发的项目，也叫 dynasm：
为了在 Rust 中编写汇编代码，我们将使用这个名为 dynasm-rs 的扩展库。由于 dynasm-rs 是对著名 C/C++ 库 dynasm 的模仿，后者则是 LuaJIT 项目的重要组成部分。因此，其作用与 Lua 的 DynASM 是一样的，dynasm-rs 是一个汇编语言编译器，它可以将汇编代码编译为机器码。
在项目中的 Cargo.toml 文件里添加相应的依赖项：
[dependencies]
dynasm = "1.2.3"
dynasmrt = "1.2.3"
实现BrainfuckJIT
在生活中，如果有两种合理但不同的方法时，你应该总是研究两者的结合，看看能否找到两全其美的方法。我们称这种组合为杂合（hybrid）。例如，为什么只吃巧克力或简单的坚果，而不是将两者结合起来，成为一块可爱的坚果巧克力呢？
在 1960 年约翰·麦卡锡偶然发现了此方法。在他的重要论文《符号表达式的递归函数及其在机器上的计算》（Recursive functions of symbolic expressions and their computation by machine, Part I）第一部分中，他提到了在运行时被转换的函数，因此不需要编译并执行系统。JIT 编译是两种传统的机器代码翻译方法：提前编译（AOT）和解释（Interpreter）的结合，它结合了两者的优点和缺点。
让我们回忆一下第一篇文章的内容，是关于编写 JIT 所需要的 4 个步骤：
- 申请一段可写和可执行的内存
- 将源码翻译为机器码（通常经过汇编）
- 将机器码写入第一步申请的内存
- 执行这部分内存
首先申请一段内存作为 Brainfuck 解释器的纸带，并取得该段纸带在内存中的起始地址和终止地址：
let mut tape: Box<[u8]> = vec![0; 65536].into_boxed_slice();
let tape_addr_from = tape.as_mut_ptr();
let tape_addr_to = unsafe { tape_addr_from.add(tape.len()) };
我们将整个 Brainfuck 程序看作一个函数，它接收两个参数：纸带起始地址和终止地址。根据 nasm 规范，函数的第一个参数被存在 rdi 寄存器中，第二个参数被存在 rsi 寄存器中。我们将它们复制到 r12 和 r13 这两个寄存器内持久化存储。同时 rcx 寄存器被用作为纸带的指针 SP，赋予其初始值为纸带起始地址。
let mut ops = dynasmrt::x64::Assembler::new()?;
let entry_point = ops.offset();
dynasm!(ops
; .arch x64
; mov r12, rdi // arg tape_addr_from
; mov r13, rsi // arg tape_addr_to
; mov rcx, rdi // stack pointer
);
编写 sysv64 格式的 getchar/putchar 函数，使之后的汇编代码中可以调用这两个函数。
unsafe extern "sysv64" fn putchar(char: *mut u8) {
std::io::stdout()
.write_all(std::slice::from_raw_parts(char, 1))
.unwrap();
}
unsafe extern "sysv64" fn getchar(char: *mut u8) {
std::io::stdout().flush().unwrap();
std::io::stdin()
.read_exact(std::slice::from_raw_parts_mut(char, 1))
.unwrap();
}
之后, 将每个 IR 翻译为语义等价的汇编代码如下。首先 SHL(x)，SHR(x)，ADD(x) 和 SUB(x) 4 个操作符最为简单，它们均只使用一行汇编代码即可完成翻译。之后是 PUTCHAR 与 GETCHAR，们遵循汇编中函数调用的逻辑，的参数与地址按照规则写入指定寄存器，然后，用 call 指令调用该函数。最后是 JIZ(x) 与 JNZ(x)，它们，流程控制，其对应，编代码通过组合使用标签，比较运算，jump 指令完成。
for ir in code.instrs {
match ir {
ir::IR::SHL(x) => dynasm!(ops
; sub rcx, x as i32 // sp -= x
),
ir::IR::SHR(x) => dynasm!(ops
; add rcx, x as i32 // sp += x
),
ir::IR::ADD(x) => dynasm!(ops
; add BYTE [rcx], x as i8 // *sp += x
),
ir::IR::SUB(x) => dynasm!(ops
; sub BYTE [rcx], x as i8 // *sp -= x
),
ir::IR::PUTCHAR => dynasm!(ops
; mov r15, rcx
; mov rdi, rcx
; mov rax, QWORD putchar as _
; sub rsp, BYTE 0x28
; call rax
; add rsp, BYTE 0x28
; mov rcx, r15
),
ir::IR::GETCHAR => dynasm!(ops
; mov r15, rcx
; mov rdi, rcx
; mov rax, QWORD getchar as _
; sub rsp, BYTE 0x28
; call rax
; add rsp, BYTE 0x28
; mov rcx, r15
),
ir::IR::JIZ(_) => {
let l = ops.new_dynamic_label();
let r = ops.new_dynamic_label();
loops.push((l, r));
dynasm!(ops
; cmp BYTE [rcx], 0
; jz => r
; => l
)
}
ir::IR::JNZ(_) => {
let (l, r) = loops.pop().unwrap();
dynasm!(ops
; cmp BYTE [rcx], 0
; jnz => l
; => r
)
}
}
}
同时不要忘记为该函数写上 return：
dynasm!(ops
; ret
);
最后，通过强制类型换将这段内存标记为一个合法的 rust 函数的函数体，这可以通过 std::mem::transmute 函数实现。
let exec_buffer = ops.finalize().unwrap(); // compile the asm
let fun: extern "sysv64" fn(tape_addr_from: *mut u8, tape_addr_to: *mut u8) =
unsafe { std::mem::transmute(exec_buffer.ptr(entry_point)) };
fun(tape_addr_from, tape_addr_to);
至此，我们成功将一个 Brainfuck 程序动态编译为一个函数。上面的汇编代码中没有进行包括 I/O，出等方面的错误处理，一项复杂的工程，并且特意不被加入到代码中以便读者只关心其核心逻辑。你可以尝试自己去实现。
完整代码如下：
#![feature(proc_macro_hygiene)]
use std::io::prelude::*;
use dynasm::dynasm;
use dynasmrt::{DynasmApi, DynasmLabelApi};
mod opcode;
mod ir_code;
use ir_code::IR;
unsafe extern "sysv64" fn putchar(char: u8) {
std::io::stdout().write_all(&[char]).unwrap()
}
#[derive(Default)]
struct Interpreter {}
impl Interpreter {
fn run(&mut self, data: Vec<u8>) -> Result<(), Box<dyn std::error::Error>> {
// 将源文件中的内容转换为 Opcode 数组
let opcode_code = opcode::Code::from(data)?;
// 将 opcode 转换为 ir code
let code = ir_code::Code::from(opcode_code.instrs)?;
// 用来当匹配到 [ 和 ] 时执行跳转的
let mut loops = vec![];
// 通过此对象分配可写、可执行的内存，并且需要将代码都写入到此内存中
let mut ops = dynasmrt::x64::Assembler::new()?;
// 代码的入口点
let entry_point = ops.offset();
dynasm!(ops
// 声明要使用的汇编语法
; .arch x64
// 将内存的起始地址放到 rcx，rdi存储的是我们生成的函数的第一个参数
; mov rcx, rdi
);
for ir in code.instrs {
match ir {
IR::SHL(x) => dynasm!(ops
; sub rcx, x as i32 // sp -= x
),
IR::SHR(x) => dynasm!(ops
; add rcx, x as i32 // sp += x
),
IR::ADD(x) => dynasm!(ops
; add BYTE [rcx], x as i8 // sp* += x
),
IR::SUB(x) => dynasm!(ops
; sub BYTE [rcx], x as i8 // sp* -= x
),
IR::PUTCHAR => dynasm!(ops
; mov r12, rcx
; mov rdi, [rcx]
; mov rax, QWORD putchar as _
; sub rsp, BYTE 0x28 // Make stack 16 byte aligned
; call rax
; add rsp, BYTE 0x28
; mov rcx, r12
),
IR::GETCHAR => {
// 用的少，省略
}
IR::JIZ(_) => {
let l = ops.new_dynamic_label();
let r = ops.new_dynamic_label();
loops.push((l, r));
// 如果指针指向的单元值为零，向后跳转到对应的 ] 指令的次一指令处
dynasm!(ops
; cmp BYTE [rcx], 0
; jz => r
; => l
)
}
IR::JNZ(_) => {
let (l, r) = loops.pop().unwrap();
// 如果指针指向的单元值不为零，向前跳转到对应的 [ 指令的次一指令处
dynasm!(ops
; cmp BYTE [rcx], 0
; jnz => l
; => r
)
}
}
}
dynasm!(ops
; ret
);
// 对 ops 对象调用 finalize 后才会实际分配内存
let exec_buffer = ops.finalize().unwrap();
// 给 Brainfuck 源码执行时所使用的内存
let mut memory: Box<[u8]> = vec![0; 65536].into_boxed_slice();
// 内存起始地址
let memory_addr_from = memory.as_mut_ptr();
// 内存结束地址
let memory_addr_to = unsafe { memory_addr_from.add(memory.len()) };
// 将 exec_buffer 这块可写、可执行的内存转换成一个函数
let fun: fn(memory_addr_from: *mut u8, memory_addr_to: *mut u8) =
unsafe { std::mem::transmute(exec_buffer.ptr(entry_point)) };
// 执行该函数
fun(memory_addr_from, memory_addr_to);
Ok(())
}
}
fn main() -> Result<(), Box<dyn std::error::Error>> {
let args: Vec<String> = std::env::args().collect();
assert!(args.len() >= 2);
let mut f = std::fs::File::open(&args[1])?;
let mut c: Vec<u8> = Vec::new();
f.read_to_end(&mut c)?;
let mut i = Interpreter::default();
i.run(c)
}
