Brainfuck JIT 虚拟机教程

当我们谈到 JIT 时，通常会想到 V8、JVM 之类的庞然大物，然后望而生畏，觉得 JIT 是一种极其高深复杂的技术。

但 JIT 也可以变得非常简单，我们不需要做完善的优化和分析，只要输入源码，输出机器指令，再执行，这和普通的文本处理程序没什么区别。

在本教程中，我们将用 Rust 语言实现一个简单的 Brainfuck JIT 虚拟机，逐步理解 JIT 技术。

1. 背景知识

Brainfuck 简介

https://en.wikipedia.org/wiki/Brainfuck

Brainfuck 仅包含八个指令，却是图灵完备的，理论上它可以做到任何用图灵机能做到的事情。

机器模型：一个字节数组，一个数据指针，一个指令指针，输入流，输出流。

数组初始化为全零，数据指针初始时指向数组的第一个字节，指令指针指向第一条指令。

字符 '>'：将数据指针加一。

字符 '<'：将数据指针减一。

字符 '+'：将数据指针所指的单元加一。

字符 '-'：将数据指针所指的单元减一。

字符 ','：从输入流中读取一个字节，存入数据指针所指单元。

字符 '.'：输出数据指针所指单元的字节。

字符 '['：如果当前单元是 0，那么跳转到对应的 ']' 的下一条指令，否则继续执行。

字符 ']'：如果当前单元不是 0，那么跳转到对应的 '[' 的下一条指令，否则继续执行。

Brainfuck 可以直接对应到 C 代码，仅用几十行就能写个从 bf 到 c 的编译器。

初始化

char array[30000] = {0};
char *ptr=array;

指令	C 代码
>	++ptr;
<	--ptr;
+	++*ptr;
-	--*ptr;
.	putchar(*ptr);
,	*ptr=getchar();
[	while (*ptr) {
]	}

Brainfuck 可视化

https://ashupk.github.io/Brainfuck/brainfuck-visualizer-master/index.html

Rust 资源

Rust 官网 https://www.rust-lang.org/

Rust 一键安装 https://rustup.rs/

Rust 并不是一个能立即上手的语言，如果您是一位 Rust 新手，最好先全面了解 Rust 语言特性。

Rust 官网入门教程 https://www.rust-lang.org/zh-CN/learn

中文书《Rust编程之道》 https://book.douban.com/subject/30418895/

2. 启动项目

初始化项目

cargo new bfjit
cd bfjit
cargo run

空白项目默认为 Hello world

Finished dev [unoptimized + debuginfo] target(s) in 0.06s
Running `target/debug/bfjit`
Hello, world!

一切就绪。

3. 中间表示

为了方便后期处理，我们先将 Brainfuck 代码转换为一种中间表示 (IR).

在 src 目录下添加文件 bfir.rs

此时目录结构

.
├── Cargo.lock
├── Cargo.toml
└── src
   ├── bfir.rs
   └── main.rs

main.rs

mod bfir;

fn main() {
    println!("Hello, world!");
}

本节以下代码都将写入 bfir.rs

IR 定义

AddVal, SubVal 表示将当前单元加减某一数值。

AddPtr, SubPtr 表示将数据指针加减某一数值。

Jz，Jnz 表示跳转指令，带有跳转地址。

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum BfIR {
    AddVal(u8),  // +
    SubVal(u8),  // -
    AddPtr(u32), // >
    SubPtr(u32), // <
    GetByte,     // ,
    PutByte,     // .
    Jz,          // [
    Jnz,         // ]
}

错误处理

我们使用 thiserror 库，用来轻松定义错误类型。

https://github.com/dtolnay/thiserror

导入第三方库时：

若 Rust 版本 >= 1.62，可直接 cargo add 。

若 Rust 版本 < 1.62，推荐用 cargo-edit 插件。安装 cargo-edit 插件后，可以用命令导入第三方库。

https://github.com/killercup/cargo-edit

也可手动编辑 Cargo.toml。

cargo install cargo-edit

cargo add thiserror

错误定义

#[derive(Debug, thiserror::Error)]
pub enum CompileErrorKind {
    #[error("Unclosed left bracket")]
    UnclosedLeftBracket,
    #[error("Unexpected right bracket")]
    UnexpectedRightBracket,
}

#[derive(Debug)]
pub struct CompileError {
    line: u32,
    col: u32,
    kind: CompileErrorKind,
}

为 CompileError 实现 Display 和 Error，Display 用于对人友好的信息，Error 表明它是一个错误类型。

#[derive(Debug)]
pub struct CompileError {
    line: u32,
    col: u32,
    kind: CompileErrorKind,
}

impl fmt::Display for CompileError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{} at line {}:{}", self.kind, self.line, self.col)
    }
}

编译为 IR

pub fn compile(src: &str) -> Result<Vec<BfIR>, CompileError> {
    let mut code: Vec<BfIR> = vec![];

    let mut stk: Vec<(u32, u32, u32)> = vec![];

    let mut line: u32 = 1;
    let mut col: u32 = 0;

compile 函数接收一个字符串，返回 IR 序列。

code 存储已解析的 IR，stk 作为栈，存储左括号的IR位置、源代码行位置、源代码列位置，line 和 col 分别记录行号和列号。

主循环结构

for ch in src.chars() {
    col += 1;
    match ch {
        ...
    }
}

处理换行

'\n' => {
    line += 1;
    col = 0;
}

处理普通指令

'+' => code.push(BfIR::AddVal(1)),
'-' => code.push(BfIR::SubVal(1)),
'>' => code.push(BfIR::AddPtr(1)),
'<' => code.push(BfIR::SubPtr(1)),
',' => code.push(BfIR::GetByte),
'.' => code.push(BfIR::PutByte),

处理左括号，将左括号位置入栈。

'[' => {
    let pos = code.len() as u32;
    stk.push((pos, line, col));
    code.push(BfIR::Jz)
}

处理右括号，从栈中弹出左括号位置，如果栈为空，说明没有匹配的左括号，生成一个编译错误并返回。如果有匹配的左括号，则正常生成 IR。

']' => {
    stk.pop().ok_or(CompileError {
        line,
        col,
        kind: CompileErrorKind::UnexpectedRightBracket,
    })?;

    code.push(BfIR::Jnz)
}

忽略其他字符

_ => {}

循环结束后，如果栈不为空，说明有左括号没有匹配到右括号，弹出左括号位置，生成编译错误。最后返回生成的IR.

if let Some((_, line, col)) = stk.pop() {
    return Err(CompileError {
        line,
        col,
        kind: CompileErrorKind::UnclosedLeftBracket,
    });
}
Ok(code)

完整代码

pub fn compile(src: &str) -> Result<Vec<BfIR>, CompileError> {
    let mut code: Vec<BfIR> = vec![];

    let mut stk: Vec<(u32, u32, u32)> = vec![];

    let mut line: u32 = 1;
    let mut col: u32 = 0;

    for ch in src.chars() {
        col += 1;
        match ch {
            '\n' => {
                line += 1;
                col = 0;
            }
            '+' => code.push(BfIR::AddVal(1)),
            '-' => code.push(BfIR::SubVal(1)),
            '>' => code.push(BfIR::AddPtr(1)),
            '<' => code.push(BfIR::SubPtr(1)),
            ',' => code.push(BfIR::GetByte),
            '.' => code.push(BfIR::PutByte),
            '[' => {
                let pos = code.len() as u32;
                stk.push((pos, line, col));
                code.push(BfIR::Jz)
            }
            ']' => {
                stk.pop().ok_or(CompileError {
                    line,
                    col,
                    kind: CompileErrorKind::UnexpectedRightBracket,
                })?;

                code.push(BfIR::Jnz)
            }
            _ => {}
        }
    }
    if let Some((_, line, col)) = stk.pop() {
        return Err(CompileError {
            line,
            col,
            kind: CompileErrorKind::UnclosedLeftBracket,
        });
    }
    Ok(code)
}

简单优化

brainfuck 代码中经常有连续的算术加减和指针移动，这些操作完全可以折叠起来。

优化函数直接操作 IR 序列，用一次遍历完成折叠，原地操作。时间复杂度 O(n)，空间复杂度 O(1)。

这里定义两个宏来避免复制粘贴大量同样的代码。

    let mut i = 0;
    let mut pc = 0;

    macro_rules! _fold_ir {
        ($variant:ident, $x:ident) => {{
            let mut j = i + 1;
            while j < len {
                if let $variant(d) = code[j] {
                    $x = $x.wrapping_add(d);
                } else {
                    break;
                }
                j += 1;
            }
            i = j;
            code[pc] = $variant($x);
            pc += 1;
        }};
    }

    macro_rules! _normal_ir {
        () => {{
            code[pc] = code[i];
            pc += 1;
            i += 1;
        }};
    }

    use BfIR::*;
    while i < len {
        match code[i] {
            AddPtr(mut x) => _fold_ir!(AddPtr, x),
            SubPtr(mut x) => _fold_ir!(SubPtr, x),
            AddVal(mut x) => _fold_ir!(AddVal, x),
            SubVal(mut x) => _fold_ir!(SubVal, x),
            GetByte => _normal_ir!(),
            PutByte => _normal_ir!(),
            Jz => _normal_ir!(),
            Jnz => _normal_ir!(),
        }
    }

    code.truncate(pc);
    code.shrink_to_fit();
}

简单测试

Rust 内置一套单元测试框架，在模块内随手写个函数，标上 #[test]，然后运行命令 cargo test。cargo 会收集所有测试，逐个运行，并报告结果。

#[test]
fn test_compile() {
    assert_eq!(
        compile("+[,.]").unwrap(),
        vec![
            BfIR::AddVal(1),
            BfIR::Jz,
            BfIR::GetByte,
            BfIR::PutByte,
            BfIR::Jnz,
        ]
    );

    match compile("[").unwrap_err().kind {
        CompileErrorKind::UnclosedLeftBracket => {}
        _ => panic!(),
    };

    match compile("]").unwrap_err().kind {
        CompileErrorKind::UnexpectedRightBracket => {}
        _ => panic!(),
    };

    let mut code = compile("[+++++]").unwrap();
    optimize(&mut code);
    assert_eq!(code, vec![BfIR::Jz, BfIR::AddVal(5), BfIR::Jnz]);
}

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,
    }
}
}

5. 来点错误

添加一个错误模块

src/main.rs

mod error;

src/error.rs

#[derive(Debug, thiserror::Error)]
pub enum RuntimeError {
    #[error("IO: {0}")]
    IO(#[from] std::io::Error),

    #[error("Pointer overflow")]
    PointerOverflow,
}

#[derive(Debug, thiserror::Error)]
pub enum VMError {
    #[error("IO: {0}")]
    IO(#[from] std::io::Error),

    #[error("Compile: {0}")]
    Compile(#[from] crate::bfir::CompileError),

    #[error("Runtime: {0}")]
    Runtime(#[from] RuntimeError),
}

pub type Result<T> = std::result::Result<T, VMError>;

虚拟机运行时可能出现两种错误，IO 错误和指针溢出。

虚拟机首先需要读取源代码，可能出现IO错误，然后将源代码编译为 IR，可能发生编译错误，最后运行程序，可能有运行时错误。

6. 实现虚拟机

添加模块 src/bfjit.rs

导入必要的类型

use crate::bfir::{self, BfIR};
use crate::error::{Result, RuntimeError, VMError};

use std::io::{Read, Write};
use std::path::Path;
use std::ptr;

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

虚拟机定义与 Brainfuck 机器模型一致。

机器码缓冲区，起始偏移，字节数组，输入输出流。

字节数组大小至少为 30000，这里设置为 4 MiB.

pub struct BfVM<'io> {
    code: dynasmrt::ExecutableBuffer,
    start: dynasmrt::AssemblyOffset,
    memory: Box<[u8]>,
    input: Box<dyn Read + 'io>,
    output: Box<dyn Write + 'io>,
}

准备汇编可调用的函数，brainfuck 通过 unsafe 函数与虚拟机交互。

为了传出可能的错误，把错误移动到堆上，返回裸指针。您也可以选择其他方式来传出错误。

在 FFI 时完善地处理 panic 和 backtrace 是一个复杂的问题。为了不增加过多处理代码，这里选择忽略 unsafe 函数中的 panic 问题。在代码正确的情况下，panic 应该不会发生。

请注意：跨语言的栈展开 (stack unwind) 是未定义行为，有可能引发段错误，您需要仔细研究 ABI 才能解决它。

#[inline(always)]
fn vm_error(re: RuntimeError) -> *mut VMError {
    let e = Box::new(VMError::from(re));
    Box::into_raw(e)
}

impl<'io> BfVM<'io> {
    unsafe extern "sysv64" fn getbyte(this: *mut Self, ptr: *mut u8) -> *mut VMError {
        let mut buf = [0_u8];
        let this = &mut *this;
        match this.input.read(&mut buf) {
            Ok(0) => {}
            Ok(1) => *ptr = buf[0],
            Err(e) => return vm_error(RuntimeError::IO(e)),
            _ => unreachable!(),
        }
        ptr::null_mut()
    }

    unsafe extern "sysv64" fn putbyte(this: *mut Self, ptr: *const u8) -> *mut VMError {
        let buf = std::slice::from_ref(&*ptr);
        let this = &mut *this;
        match this.output.write_all(buf) {
            Ok(()) => ptr::null_mut(),
            Err(e) => vm_error(RuntimeError::IO(e)),
        }
    }

    unsafe extern "sysv64" fn overflow_error() -> *mut VMError {
        vm_error(RuntimeError::PointerOverflow)
    }
}

用输入流、输出流、源文件路径初始化虚拟机，优化选项也由外部提供。

compile 方法暂时留空。

impl<'io> BfVM<'io> {
    fn compile(code: &[BfIR]) -> Result<(dynasmrt::ExecutableBuffer, dynasmrt::AssemblyOffset)> {
        todo!()
    }

impl<'io> BfVM<'io> {
    pub fn new(
        file_path: &Path,
        input: Box<dyn Read + 'io>,
        output: Box<dyn Write + 'io>,
        optimize: bool,
    ) -> Result<Self> {
        let src = std::fs::read_to_string(file_path)?;
        let mut ir = bfir::compile(&src)?;
        drop(src);

        if optimize {
            bfir::optimize(&mut ir);
        }
        let (code, start) = Self::compile(&ir)?;
        drop(ir);

        let memory = vec![0; MEMORY_SIZE].into_boxed_slice();
        Ok(Self {
            code,
            start,
            memory,
            input,
            output,
        })
    }

即时编译出的裸函数接收虚拟机指针和字节数组的边界指针，返回错误指针。

执行函数后检查错误指针，如果非空，就把错误从堆上移动到栈上再返回。

    pub fn run(&mut self) -> Result<()> {
        type RawFn = unsafe extern "sysv64" fn(
            this: *mut BfVM<'_>,
            memory_start: *mut u8,
            memory_end: *const u8,
        ) -> *mut VMError;

        let raw_fn: RawFn = unsafe { std::mem::transmute(self.code.ptr(self.start)) };

        let this: *mut Self = self;
        let memory_start = self.memory.as_mut_ptr();
        let memory_end = unsafe { memory_start.add(MEMORY_SIZE) };

        let ret: *mut VMError = unsafe { raw_fn(this, memory_start, memory_end) };

        if ret.is_null() {
            Ok(())
        } else {
            Err(*unsafe { Box::from_raw(ret) })
        }
    }

7. 生成机器码

完成上一节留空的 compile 方法。

首先初始化汇编器。整个 brainfuck 程序编译为一个大函数，函数起始地址就是最开始的偏移。

loops 作为栈，用来存放动态标签，指引跳转。

impl<'io> BfVM<'io> {
    #[allow(clippy::fn_to_numeric_cast)]
    fn compile(code: &[BfIR]) -> Result<(dynasmrt::ExecutableBuffer, dynasmrt::AssemblyOffset)> {
        let mut ops = dynasmrt::x64::Assembler::new()?;
        let start = ops.offset();

        let mut loops = vec![];

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

在注释中约定各个寄存器存放的参数，虚拟机调用裸函数时将把参数放入寄存器。

        // this:         rdi r12
        // memory_start: rsi r13
        // memory_end:   rdx r14
        // ptr:          rcx r15

程序开始，首先把 rax 压栈。x86-64-psABI 规定参数区域的结尾按16字节对齐。函数开始时返回地址压栈，此时栈指针+8 是 16 的倍数，因此再把 rax 压栈，使栈指针对齐，以便之后的函数调用，rax 的内容没有意义。

把寄存器中的各参数存入非易失性寄存器，调用其他函数时参数寄存器的值可能丢失。rbp, rbx, r12 ~ r15 是非易失性寄存器。

        dynasm!(ops
            ; push rax
            ; mov r12, rdi   // save this
            ; mov r13, rsi   // save memory_start
            ; mov r14, rdx   // save memory_end
            ; mov rcx, rsi   // ptr = memory_start
        );

每个 IR 依次映射到汇编。

指针移动，需要检查算术溢出和数组边界溢出，出错即跳转到全局标签所指的错误处理区域。

        use BfIR::*;
        for &ir in code {
            match ir {
                AddPtr(x) => dynasm!(ops
                    ; add rcx, x as i32     // ptr += x
                    ; jc  ->overflow        // jmp if overflow
                    ; cmp rcx, r14          // ptr - memory_end
                    ; jnb ->overflow        // jmp if ptr >= memory_end
                ),
                SubPtr(x) => dynasm!(ops
                    ; sub rcx, x as i32     // ptr -= x
                    ; jc  ->overflow        // jmp if overflow
                    ; cmp rcx, r13          // ptr - memory_start
                    ; jb  ->overflow        // jmp if ptr < memory_start
                ),

单个字节的算术加减，允许溢出。

                AddVal(x) => dynasm!(ops
                    ; add BYTE [rcx], x as i8    // *ptr += x
                ),
                SubVal(x) => dynasm!(ops
                    ; sub BYTE [rcx], x as i8    // *ptr -= x
                ),

IO 操作。首先保存当前数据指针寄存器，将虚拟机函数所需的各参数和函数地址放入寄存器，调用函数。

如果函数返回的不是空指针，说明出错，应该跳转到IO错误处理区域。

最后恢复数据指针寄存器。

                GetByte => dynasm!(ops
                    ; mov  r15, rcx         // save ptr
                    ; mov  rdi, r12
                    ; mov  rsi, rcx         // arg0: this, arg1: ptr
                    ; mov  rax, QWORD BfVM::getbyte as _
                    ; call rax              // getbyte(this, ptr)
                    ; test rax, rax
                    ; jnz  ->io_error       // jmp if rax != 0
                    ; mov  rcx, r15         // recover ptr
                ),
                PutByte => dynasm!(ops
                    ; mov  r15, rcx         // save ptr
                    ; mov  rdi, r12
                    ; mov  rsi, rcx         // arg0: this, arg1: ptr
                    ; mov  rax, QWORD BfVM::putbyte as _
                    ; call rax              // putbyte(this, ptr)
                    ; test rax, rax
                    ; jnz  ->io_error       // jmp if rax != 0
                    ; mov  rcx, r15         // recover ptr
                ),

跳转指令。利用 dynasm 提供的 api，创建两个动态标签，分别生成跳转汇编。由于编译到 IR 时已经验证过跳转指令的对应关系，这里的栈可以直接弹出。

                Jz => {
                    let left = ops.new_dynamic_label();
                    let right = ops.new_dynamic_label();
                    loops.push((left, right));

                    dynasm!(ops
                        ; cmp BYTE [rcx], 0
                        ; jz => right       // jmp if *ptr == 0
                        ; => left
                    )
                }
                Jnz => {
                    let (left, right) = loops.pop().unwrap();
                    dynasm!(ops
                        ; cmp BYTE [rcx], 0
                        ; jnz => left       // jmp if *ptr != 0
                        ; => right
                    )
                }
            }
        }

正常退出函数时应该返回空指针，表示没有错误。

溢出时生成一个溢出错误，IO错误时错误指针已经存入 rax，无需处理。

最后退栈，与函数开始时的压栈对应，维持栈平衡。

        dynasm!(ops
            ; xor rax, rax
            ; jmp >exit
            ; -> overflow:
            ; mov rax, QWORD BfVM::overflow_error as _
            ; call rax
            ; jmp >exit
            ; -> io_error:
            ; exit:
            ; pop rdx
            ; ret
        );

完成汇编，取出可执行缓冲区，返回。

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

        Ok((code, start))
    }

命令行界面

虚拟机已经完成，剩下的就是把它包装成命令行界面 (CLI) 了。

我们使用 clap 库，用结构体定义命令行参数解析。

https://github.com/clap-rs/clap

cargo add clap --features derive

这就是 main.rs 的全部代码。

mod bfir;
mod bfjit;
mod error;

use crate::bfjit::BfVM;

use std::io::{stdin, stdout};
use std::path::PathBuf;

use clap::Parser;

#[derive(Debug, clap::Parser)]
#[clap(version)]
struct Opt {
    #[clap(name = "FILE")]
    file_path: PathBuf,

    #[clap(short = 'o', long = "optimize", help = "Optimize code")]
    optimize: bool,
}

fn main() {
    let opt = Opt::parse();

    let stdin = stdin();
    let stdout = stdout();

    let ret = BfVM::new(
        &opt.file_path,
        Box::new(stdin.lock()),
        Box::new(stdout.lock()),
        opt.optimize,
    )
    .and_then(|mut vm| vm.run());

    if let Err(e) = &ret {
        eprintln!("bfjit: {}", e);
    }

    std::process::exit(ret.is_err() as i32)
}

把写好的命令行应用安装到系统。

$ cargo install --path .
$ bfjit --help
Usage: bfjit [OPTIONS] <FILE>

Arguments:
  <FILE>  

Options:
  -o, --optimize  Optimize code
  -h, --help      Print help
  -V, --version   Print version

从 github 上找一些 brainfuck 程序运行，观察 JIT 与解释器的效率差别。

$ bfjit bf/hello.bf
Hello World!

观察折叠操作对速度的影响。

$ time bfjit bf/mendelbrot.bf > m.txt

real    0m5.858s
user    0m5.840s
sys     0m0.012s

$ time bfjit bf/mendelbrot.bf -o > m.txt

real    0m1.921s
user    0m1.903s
sys     0m0.010s

甚至可以运行一层 brainfuck 自解释器。

注意：多层自解释器的效率会严重下降，短时间内无法得出结果。

https://github.com/cagataycali/awesome-brainfuck

http://homepages.xnet.co.nz/~clive/eigenratios/cgbfi2.b

结语

本教程基于 dynasm 实现了一个 Brainfuck JIT 虚拟机，主要功能有 Brainfuck 解析编译、简单优化、动态生成机器码，并提供了友好的命令行界面。

希望本项目对有兴趣深入研究 JIT 技术的人有所帮助。

欢迎 PR 来进一步改进本项目。

第三方依赖：

https://crates.io/crates/dynasm

https://github.com/dtolnay/thiserror

https://github.com/clap-rs/clap