使用Rust实现Brainfuck的JIT编译器

作者: 端碗吹水 | 来源:发表于2022-06-13 08:41 被阅读0次

[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：

https://crates.io/crates/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)
}

使用Rust实现Brainfuck的JIT编译器

使用Rust实现Brainfuck的JIT编译器

x64汇编简介

Linux x64 汇编/Hello World

准备

NASM 语法

Hello World!

什么是JIT

Hello, JIT World: The Joy of Simple JITs

dynasm介绍

实现BrainfuckJIT

相关文章

网友评论

延伸阅读

深度阅读

栏目导航

热点阅读