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 可视化

http://fatiherikli.github.io/brainfuck-visualizer/

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 工作方式是这样的:

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

以下代码引用自 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