用Rust实现Brainfuck的JIT编译器

2022-06-06 13:01:28 浏览数 (1)

x64汇编简介

Linux x64 汇编/Hello World

我们每天产出大量的垃圾代码,我们每个人都可以像这样简单地编写最简单的代码:

代码语言:javascript复制
#include <stdio.h>

int main() {
  int x = 10;
  int y = 100;
  printf("x   y = %dn",x   y);
  return 0;
}

希望读者们都可以理解上述 C 代码的作用。但是,此代码在底层如何工作?我认为并非所有人都能回答这个问题,我也是。我可以用Haskell,Erlang,Go 等高级编程语言编写代码,但是在它们编译后我并不知道它在底层是如何工作的。因此,我决定采取一些更深入的步骤,进行记录,并描述我对此的学习过程。希望这个过程不仅仅只是对我来说很有趣。让我们开始吧。

准备

开始之前,我们必须准备一些事情,如我所写的那样,我目前使用 Ubuntu 18.04,因此我的文章将针对该操作系统和体系结构。不同的 CPU 支持不同的指令集,目前我使用 Intel 的 64 位 CPU。同时我也将使用 NASM 语法。你可以使用以下方法安装它:

代码语言:javascript复制
$ apt install nasm

记住,Netwide Assembler(简称 NASM)是一款基于英特尔 x86 架构的汇编与反汇编工具。这就是我们目前需要的工具。

NASM 语法

在这里,我将不介绍完整的汇编语法,我们仅提及其庞大语法的一小部分,也是那些我们将在本文中使用到的部分。通常, NASM 程序分为几个段(section),在这篇文章中,我们将遇到以下两个段:

  • 数据段:data section
  • 文本段:text section

数据段部分用于声明常量,此数据在运行时不会更改。声明数据部分的语法为:

代码语言:javascript复制
section .data

文本段部分用于代码。该部分必须以全局声明 _start 开头,该声明告诉内核程序从何处开始执行。

代码语言:javascript复制
section .text
global _start
_start:

注释以符号 ; 开头。每条 NASM 源代码行都包含以下四个字段的某种组合:

代码语言:javascript复制
[label:] instruction [operands] [; comment]

方括号中的字段是可选的。基本的 NASM 指令由两部分组成,第一部分是要执行的指令的名称,第二部分是该命令的操作数。例如:

代码语言:javascript复制
mov rax, 48 ; put value 48 in the rax register

Hello World!

让我们用 NASM 编写第一个程序。当然,这将是经典的 Hello World! 程序。这是它的代码:

代码语言:javascript复制
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 寄存器。现在我们知道 raxrdirsi 等代表了什么了,但是还需要知道什么时候该使用 rax ,什么时候使用 rdi 等。

  • rax:临时寄存器,当我们调用 syscall 时,rax 必须包含 syscall 号码,所以后面的数字就是 syscall 的号码
  • rdi:用于将第 1 个参数传递给函数
  • rsi:用于将第 2 个参数传递给函数
  • rdx:用于将第 3 个参数传递给函数

换句话说,我们只是在调用 sys_write 的 syscall。看看 sys_write 的定义:

代码语言:javascript复制
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。

最后,让我们来构建这个程序,我们需要执行以下命令:

代码语言:javascript复制
$ 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 编译器只需要四步,就像把大象装到冰箱里一样。

  1. 申请一段可写和可执行的内存
  2. 将源码翻译为机器码(通常经过汇编)
  3. 将机器码写入第一步申请的内存
  4. 执行这部分内存

Hello, JIT World: The Joy of Simple JITs

事不宜迟,让我们跳进我们的第一个 JIT 程序。该代码是特定于 64 位 Unix 的,因为它使用了 mmap()。鉴于此读者需要拥有支持该代码的处理器和操作系统。

代码语言:javascript复制
#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。它动态生成一个函数,该函数返回运行时指定的整数,然后运行该函数。读者可以验证其是否正常运行:

代码语言:javascript复制
$ 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 文件里添加相应的依赖项:

代码语言:javascript复制
[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 个步骤:

  1. 申请一段可写和可执行的内存
  2. 将源码翻译为机器码(通常经过汇编)
  3. 将机器码写入第一步申请的内存
  4. 执行这部分内存

首先申请一段内存作为 Brainfuck 解释器的纸带,并取得该段纸带在内存中的起始地址和终止地址:

代码语言:javascript复制
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 寄存器中。我们将它们复制到 r12r13 这两个寄存器内持久化存储。同时 rcx 寄存器被用作为纸带的指针 SP,赋予其初始值为纸带起始地址。

代码语言:javascript复制
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 函数,使之后的汇编代码中可以调用这两个函数。

代码语言:javascript复制
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 个操作符最为简单,它们均只使用一行汇编代码即可完成翻译。之后是 PUTCHARGETCHAR,们遵循汇编中函数调用的逻辑,的参数与地址按照规则写入指定寄存器,然后,用 call 指令调用该函数。最后是 JIZ(x)JNZ(x),它们,流程控制,其对应,编代码通过组合使用标签,比较运算,jump 指令完成。

代码语言:javascript复制
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

代码语言:javascript复制
dynasm!(ops
    ; ret
);

最后,通过强制类型换将这段内存标记为一个合法的 rust 函数的函数体,这可以通过 std::mem::transmute 函数实现。

代码语言:javascript复制
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,出等方面的错误处理,一项复杂的工程,并且特意不被加入到代码中以便读者只关心其核心逻辑。你可以尝试自己去实现。

完整代码如下:

代码语言:javascript复制
#![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)
}

0 人点赞