4. Hello, JIT

通常来说，Just In Time (JIT) 编译器是指在某段高阶代码即将运行时将其编译到机器码再执行的程序。下文中，我们把这样的程序叫做 "JIT"。

一个最简单的 JIT 工作方式是这样的：

将源代码编译为机器码。
申请一段可写可执行内存，写入机器码。
跳转到这段内存，执行机器码。
执行完毕，稍作清理，退出。

以下代码引用自 Hello, JIT World: The Joy of Simple JITs

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

int main(int argc, char *argv[]) {
  // 机器码：
  //   mov eax, 0
  //   ret
  unsigned char code[] = {0xb8, 0x00, 0x00, 0x00, 0x00, 0xc3};

  if (argc < 2) {
    fprintf(stderr, "Usage: jit1 <integer>\n");
    return 1;
  }

  // 把用户给出的数值写入机器码，覆盖立即数 "0"
  //   mov eax, <user's value>
  //   ret
  int num = atoi(argv[1]);
  memcpy(&code[1], &num, 4);

  // 分配可写可执行内存
  // 注意：真实的程序不应该映射同时可写可执行的内存，
  // 这里有安全风险。
  void *mem = mmap(NULL, sizeof(code), PROT_WRITE | PROT_EXEC,
                   MAP_ANON | MAP_PRIVATE, -1, 0);
  memcpy(mem, code, sizeof(code));

  // 定义一个函数指针指向机器码内存，再执行函数
  int (*func)() = mem;
  return func();
}

编写 JIT 的一大难点是如何生成机器码，这里通常有跨平台问题、可读性消失问题。

最笨的方法：写一段汇编，用汇编器生成机器码，再复制到高级代码里。但这样不具有通用性，开发效率非常低。

dynasm

DynAsm 是 LuaJIT 的一部分，它用预处理器把混合汇编的 C 文件转换成纯 C 文件，还包含一个微型运行时，用来执行运行时工作。

dynasm-rs 是对应的 Rust 实现，用过程宏在编译期解析汇编语法，也包含微型运行时。

Rust 过程宏作为编译器插件几乎是万能的，不光是汇编，html、shell、cpp 等语言都能嵌入，转换成对应的 Rust 结构，给人一种相当纯粹的感觉。过程宏还有很多用途，感兴趣的可以自行研究。

cargo add dynasm dynasmrt

修改 main.rs，导入 dynasm.

mod bfir;

use dynasm::dynasm;
use dynasmrt::{DynasmApi, DynasmLabelApi};

use std::io::{stdout, Write};

编写 print 函数，使用 "sysv64" ABI。

x86-64 Linux 系统上默认为 System V ABI. 相关文档

unsafe extern "sysv64" fn print(buf: *const u8, len: u64) -> u8 {
    let buf = std::slice::from_raw_parts(buf, len as usize);
    stdout().write_all(buf).is_err() as u8
}

首先初始化汇编器，指定架构为 x64，全局标签 hello 指向字符串。

fn main() {
    let mut ops = dynasmrt::x64::Assembler::new().unwrap();
    let string = b"Hello, JIT!\n";

    dynasm!(ops
        ; .arch x64
        ; ->hello:
        ; .bytes string
    );

dynasm 使用 nasm 的方言，左操作数为目标，右操作数为源。

sysv64 调用约定规定 rdi, rsi, rdx, rcx 存放前四个整数参数，rax 存放返回值。

    let hello = ops.offset();
    dynasm!(ops
        ; lea rdi, [->hello]                // 将 hello 字符串地址放入 rdi
        ; mov rsi, QWORD string.len() as _  // 将 字符串长度放入 rsi
        ; mov rax, QWORD print as _         // 将 print 函数地址放入 rax
        ; call rax                          // 调用函数
        ; ret                               // 返回
    );

完成汇编，取得可执行缓冲区。根据偏移拿到函数地址，强制转换为函数指针。最后调用机器码，得到结果。

    let buf = ops.finalize().unwrap();

    let hello_fn: unsafe extern "sysv64" fn() -> u8 =
        unsafe { std::mem::transmute(buf.ptr(hello)) };

    let ret = unsafe { hello_fn() };

    assert_eq!(ret, 0);
}

运行结果

$ cargo run
    Finished dev [unoptimized + debuginfo] target(s) in 1.30s
     Running `target/debug/bfjit`
Hello, JIT!

完整代码

mod bfir;

use dynasm::dynasm;
use dynasmrt::{DynasmApi, DynasmLabelApi};

use std::io::{stdout, Write};

unsafe extern "sysv64" fn print(buf: *const u8, len: u64) -> u8 {
    let buf = std::slice::from_raw_parts(buf, len as usize);
    stdout().write_all(buf).is_err() as u8
}

fn main() {
    let mut ops = dynasmrt::x64::Assembler::new().unwrap();
    let string = b"Hello, JIT!\n";

    dynasm!(ops
        ; .arch x64
        ; ->hello:
        ; .bytes string
    );

    let hello = ops.offset();
    dynasm!(ops
        ; lea rdi, [->hello]
        ; mov rsi, QWORD string.len() as _
        ; mov rax, QWORD print as _
        ; call rax
        ; ret
    );

    let buf = ops.finalize().unwrap();

    let hello_fn: unsafe extern "sysv64" fn() -> u8 =
        unsafe { std::mem::transmute(buf.ptr(hello)) };

    let ret = unsafe { hello_fn() };

    assert_eq!(ret, 0);
}

如何处理 panic

跨越 Rust 边界的 panic 是未定义行为，我们很难让汇编去匹配 unwind ABI。

暴露给外部调用的 Rust 函数最好捕获 panic，用其他方式去处理。

例如这样

#![allow(unused)]
fn main() {
unsafe extern "sysv64" fn print(buf: *const u8, len: u64) -> u8 {
    let ret = std::panic::catch_unwind(|| {
        let buf = std::slice::from_raw_parts(buf, len as usize);
        stdout().write_all(buf).is_err()
    });
    match ret {
        Ok(false) => 0,
        Ok(true) => 1,
        Err(_) => 2,
    }
}
}

Brainfuck JIT 虚拟机教程

4. Hello, JIT

dynasm

如何处理 panic