5微秒内完成JIT编译：从零构建高性能正则引擎
1. 引言：JIT编译的挑战与机遇
历史上，快速JIT编译（Just-In-Time compilation）被视为一项黑魔法。要编写一个快速的JIT编译器，你需要精通汇编语言。一个典型的例子是：目前没有任何生产级数据库拥有自己的JIT编译器，它们要么使用LLVM，要么生成C/C++代码。这两种方案都存在编译时间过长的问题，限制了其适用性。如今，借助AI，通过直接针对汇编语言编写JIT编译器，实现快速编译变得前所未有的简单。这也是新数据库超越旧数据库的一个机会领域。在构建pgrust时，我最初认为实现JIT编译器会非常困难。最终，在AI的辅助下，我发现它比预期的要容易得多，这也成为了pgrust性能卓越的部分原因。pgrust的JIT编译器在约5微秒内完成代码编译，这使我们能够对每条SQL查询进行JIT编译，而不仅仅是其中的一部分。在这篇文章中，我将带你了解如何构建自己的快速JIT编译器。我们将以一个使用JIT编译的简单正则表达式引擎为例。
2. 为什么需要JIT编译
JIT编译是指在运行时（即“即时”）生成编译代码的实践。如果做得好，它可以带来巨大的性能提升，通常是2-5倍，有时甚至更多。JIT编译的主要应用场景是当你在运行时获得的信息会显著改变程序行为时。这在编程语言解释器中尤为常见，因为它们接收要执行的代码是在运行时。JIT编译器在编程语言之外的领域也很有用，例如数据解析。有时，直到运行时你才知道要解析的数据的模式（schema），JIT可以帮助解决这个问题。让我们先实现一个玩具正则表达式引擎。为了简单起见，我们只支持两个特性：字面量字符串和重复（即正则表达式中的*）。我们还将跳过解析器，直接将正则表达式表示为已解析的Rust结构。这意味着我们将能够支持诸如“apples”和“b(an)*”这样的字符串，但不支持交替（alternation）或后顾（lookbehind）等。在代码中，这非常简单。我们将有三种节点类型：字面量字符串节点、重复节点和连接节点（即两个节点的组合）。代码如下：enum Node { Literal( 'static str), Concatenation(Box<Node>, Box<Node>), Repetition(Box<Node>) } fn literal(text: 'static str) -> Node { Node::Literal(text) } fn concatenation(left: Node, right: Node) -> Node { Node::Concatenation(Box::new(left), Box::new(right)) } fn repetition(body: Node) -> Node { Node::Repetition(Box::new(body)) }
3. 解释器实现与性能瓶颈
为我们的正则表达式引擎编写解释器也很直接：fn match_node(node: Node, input: [u8], pos: usize, next: dyn Fn(usize) -> bool) -> bool { match node { Node::Literal(text) => { let literal = text.as_bytes(); input[pos..].starts_with(literal) && next(pos + literal.len()) } Node::Concatenation(left, right) => { match_node(left, input, pos, |left_end| { match_node(right, input, left_end, next) }) } Node::Repetition(body) => { match_node(body, input, pos, |body_end| { match_node(node, input, body_end, next) }) || next(pos) } } } fn interp_match(regex: Node, input: str) -> bool { let bytes = input.as_bytes(); match_node(regex, bytes, 0, |pos| pos == bytes.len()) } 这个正则表达式引擎非常简单，代码不到20行，但让我们看看它的性能表现。为了对比，我们将代码与针对该正则表达式专门实现的手写代码进行比较。我们的示例使用正则表达式“b(an)*”。手写代码看起来像这样：fn handwritten_b_an_star(input: str) -> bool { let bytes = input.as_bytes(); let mut pos = 0; if pos == bytes.len() || bytes[pos] != b'b' { return false; } pos += 1; while pos < bytes.len() { if bytes[pos] != b'a' { return false; } pos += 1; if pos == bytes.len() || bytes[pos] != b'n' { return false; } pos += 1; } true } （有方法可以优化这段代码并使其更快，但为了我们的目的，它作为一个很好的比较基准。）当我针对这两个实现进行基准测试时，手写版本比解释器快10-20倍。显然，还有很大的改进空间。现在，让我们看看如何使用JIT编译来获得一个性能与手写版本相当的正则表达式引擎。
4. JIT编译的步骤与设计决策
JIT编译代码有两个步骤。首先，为要运行的代码生成汇编。一旦有了代码，就将汇编代码打包成一个函数，你可以像调用程序中的任何其他代码一样调用它。为了生成汇编，我们将使用一种称为“复制并修补”（copy-and-patch）的方法的变体。其思想是，我们有一系列针对要JIT编译的不同操作的汇编模板。这些模板被称为“模板”（stencils）。当我们想要JIT编译一个操作时，我们取出相关的模板，并根据操作的具体情况做小的调整。这与填充真实的模板非常相似。通过将几个填充好的模板串在一起，我们可以在运行时构建一个与手写版本性能相似的程序。我们将采取以下路径：首先，我们查看为“b(an)”生成的ARM64代码。然后，我们将重复的指令序列转换为可重用的模板，编写一个发射器（emitter），从正则表达式AST填充和组合这些模板，最后将生成的指令复制到可执行内存中，以便Rust可以像普通函数一样调用它们。为了让你了解这是如何工作的，最容易的是从生成的代码开始，然后回溯到JIT编译器本身。再次强调，我们处理的是正则表达式“b(an)”。为了布局一些设计决策：我们将使用一个栈进行回溯。栈将跟踪如果正则表达式遇到死胡同我们应该去的状态。我们匹配的字符串将以空字节结尾。这意味着，如果我们到达字符串末尾，任何字符比较都会自动失败。这意味着我们不需要在任何点进行长度比较。对于程序的状态，我们将使用以下寄存器：x0 – 当前字符串位置和返回值；x1 – 用于回溯的栈顶；x2 – 用于回溯的栈底（这用于确定栈是否为空）；x9 – 用作临时变量。对于程序的输入，我们将传递：x0 – 指向字符串开头的指针；x1 – 指向我们将用作栈的位置的指针。
5. 生成的ARM64代码详解
现在我们已经处理了这些，让我们逐部分讲解生成的汇编。这特别针对macOS上的ARM64。首先，我们有初始化程序的序言（prologue）。它所做的只是通过将栈顶和栈底设置为传入的值来初始化栈：0: aa0103e2 mov x2, x1 接下来，我们有检查字符“b”的代码。如果它看到一个不是“b”的字符，我们跳转到一个处理回退逻辑的代码块。否则，我们推进字符串中的位置：; CHAR 'b' 4: 39400009 ldrb w9, [x0] ; 加载当前输入字节 8: 7101893f cmp w9, #0x62 ; 是'b'吗？ c: 54000281 b.ne 0x5c ; 否 - 回退块 10: 91000400 add x0, x0, #1 ; 是 - 推进输入 接下来，我们有重复部分“(an)*”。对于重复，我们需要进行回溯。如果我们在这里回溯，那意味着我们立即跳转到循环的末尾。这意味着我们需要将循环后的指令地址和我们在字符串中的位置都存储在栈上。14: d2800989 movz x9, #0x004c ; 构建恢复地址 18: f2a00009 movk x9, #0x0000, lsl #16 ; = 0x1_0000_004c 1c: f2c00029 movk x9, #0x0001, lsl #32 ; (循环出口) 20: f2e00009 movk x9, #0x0000, lsl #48 ; 24: a8810029 stp x9, x0, [x1], #16 ; 将(exit, pos)压入栈 有了这些，我们现在可以执行重复的主体。这将检查字符'a'和'n'，如果看到它们，就回到重复的顶部，但使用新的字符串位置。; CHAR 'a' 28: 39400009 ldrb w9, [x0] 2c: 7101853f cmp w9, #0x61 ; 'a'? 30: 54000161 b.ne 0x5c ; 否 - 回退块 34: 91000400 add x0, x0, #1 ; CHAR 'n' 38: 39400009 ldrb w9, [x0] 3c: 7101b93f cmp w9, #0x6e ; 'n'? 40: 540000e1 b.ne 0x5c ; 否 - 回退块 44: 91000400 add x0, x0, #1 ; JMP 48: 17fffff3 b 0x14 ; 回到循环顶部 现在我们过了循环。这是回溯时跳转的地方。一旦我们完成重复，我们就到了正则表达式的末尾。我们现在要做的就是检查是否在字符串的末尾。如果在字符串末尾，我们返回1表示成功。如果不是，那意味着正则表达式匹配失败，我们需要运行失败逻辑来进行回退。4c: 39400009 ldrb w9, [x0] 50: 35000069 cbnz w9, 0x5c ; 不在NUL - 回退块 54: d2800020 mov x0, #1 ; 成功 58: d65f03c0 ret 最后，我们有回退逻辑。这检查栈是否为空。如果为空，我们返回0。如果不为空，我们从栈中弹出回退地址和回退字符串位置，然后跳转到回退地址。5c: eb02003f cmp x1, x2 ; 还有帧吗？ 60: 54000060 b.eq 0x6c ; 否 - 放弃 64: a9ff0029 ldp x9, x0, [x1, #-16]! ; 弹出(resume, pos) 68: d61f0120 br x9 ; 跳转到那里 6c: d2800000 mov x0, #0 ; 不匹配 70: d65f03c0 ret
6. 构建模板（Stencils）
既然你已经看到了编译后的代码，你应该开始理解复制并修补编译器是如何工作的。我们有共同的指令集，它们之间只有微小的差异。对于这些功能块中的每一个，我们可以编写一个函数来生成相应的代码。每个函数将接受用于修改代码的值。例如，stencil_char的参数之一将是正则表达式中要比较的字符。我们将直接将该字符插入机器代码中。序言很简单，因为它只是一个代码块：const PROLOGUE_WORDS: usize = 1; fn stencil_prologue() -> [u32; PROLOGUE_WORDS] { [0xAA0103E2] // mov x2, x1 } 对于字符比较，我们需要插入我们正在比较的字符以及跳转到回退逻辑的位置：const CHAR_WORDS: usize = 4; fn stencil_char(byte: u8, stencil_pos: usize, fail_pos: usize) -> [u32; CHAR_WORDS] { [ 0x39400009, // ldrb w9, [x0] 0x7100013F | ((byte as u32) << 10), // cmp w9, #byte 0x54000001 | cond_branch_offset(stencil_pos + 2, fail_pos), // b.ne fail 0x91000400, // add x0, x0, #1 ] } 对于重复，我们有循环开始处压栈的代码和跳转到末尾的代码：const SPLIT_WORDS: usize = 5; fn stencil_split(resume_addr: u64) -> [u32; SPLIT_WORDS] { [ 0xD2800009 | addr_bits(resume_addr, 0), // movz x9, #addr[0..16] 0xF2A00009 | addr_bits(resume_addr, 1), // movk x9, #addr[16..32], lsl 16 0xF2C00009 | addr_bits(resume_addr, 2), // movk x9, #addr[32..48], lsl 32 0xF2E00009 | addr_bits(resume_addr, 3), // movk x9, #addr[48..64], lsl 48 0xA8810029, // stp x9, x0, [x1], #16 ] } const JMP_WORDS: usize = 1; fn stencil_jmp(stencil_pos: usize, target_pos: usize) -> [u32; JMP_WORDS] { [0x14000000 | branch_offset(stencil_pos, target_pos)] // b target } 然后我们有匹配和失败块，它们相当简洁：const MATCH_WORDS: usize = 4; fn stencil_match(stencil_pos: usize, fail_pos: usize) -> [u32; MATCH_WORDS] { [ 0x39400009, // ldrb w9, [x0] 0x35000009 | cond_branch_offset(stencil_pos + 1, fail_pos), // cbnz w9, fail 0xD2800020, // mov x0, #1 0xD65F03C0, // ret ] } const FAIL_WORDS: usize = 6; fn stencil_fail() -> [u32; FAIL_WORDS] { [ 0xEB02003F, // cmp x1, x2 0x54000060, // b.eq +3 (到下面的mov) 0xA9FF0029, // ldp x9, x0, [x1, #-16]! 0xD61F0120, // br x9 0xD2800000, // mov x0, #0 0xD65F03C0, // ret ] } 为了完整起见，这里是我们使用的辅助函数，它们帮助我们向指令中插入特定数据：// 计算条件分支（b.ne / cbnz）的分支偏移字段： // 从分支到目标的指令数，存储在位5..24。 fn cond_branch_offset(branch_pos: usize, target_pos: usize) -> u32 { let instr_count = target_pos as i64 - branch_pos as i64; // 可能为负 (((instr_count as u64) & 0x7FFFF) << 5) as u32 } // 计算无条件分支（b）的分支偏移字段： // 同样的想法，但存储在位0..26。 fn branch_offset(branch_pos: usize, target_pos: usize) -> u32 { let instr_count = target_pos as i64 - branch_pos as i64; // 可能为负 ((instr_count as u64) & 0x3FF_FFFF) as u32 } // 提取绝对地址的16位，为movz/movk立即数定位。 fn addr_bits(addr: u64, part: usize) -> u32 { (((addr >> (16 * part)) & 0xFFFF) as u32) << 5 }
7. 发射代码与加载机器码
现在驱动它的代码：// 计算一个节点编译成多少条指令。 fn node_words(node: Node) -> usize { match node { Node::Literal(text) => text.len() * CHAR_WORDS, Node::Concatenation(left, right) => node_words(left) + node_words(right), Node::Repetition(body) => SPLIT_WORDS + node_words(body) + JMP_WORDS, } } struct Emitter { code: Vec<u32>, fail: usize, // 共享失败块的单词偏移 base: u64, // code[0]的运行时地址，用于绝对地址空洞 } impl Emitter { // 返回下一条指令将被放置的偏移量。 fn pos(&self) -> usize { self.code.len() } // 将填充的模板追加到代码缓冲区。 fn emit(&mut self, stencil: &[u32]) { self.code.extend_from_slice(stencil); } // 为一个节点发射代码，递归到子节点。 fn emit_node(&mut self, node: Node) { match node { Node::Literal(text) => { for byte in text.as_bytes() { self.emit( &stencil_char(byte, self.pos(), self.fail)); } } Node::Concatenation(left, right) => { self.emit_node(left); self.emit_node(right); } Node::Repetition(body) => { let split_at = self.pos(); let exit = split_at + SPLIT_WORDS + node_words(body) + JMP_WORDS; self.emit( &stencil_split(self.base + exit as u64 * 4)); self.emit_node(body); self.emit( &stencil_jmp(self.pos(), split_at)); } } } } // 生成完整程序：序言、编译后的AST、MATCH、失败块。 fn generate_code(regex: Node, base: u64) -> Vec<u32> { let nwords = PROLOGUE_WORDS + node_words(regex) + MATCH_WORDS + FAIL_WORDS; let mut emitter = Emitter { code: Vec::with_capacity(nwords), fail: nwords - FAIL_WORDS, base, }; emitter.emit( &stencil_prologue()); emitter.emit_node(regex); let match_at = emitter.pos(); emitter.emit( &stencil_match(match_at, emitter.fail)); emitter.emit( &stencil_fail()); assert_eq!(emitter.pos(), nwords); emitter.code } 这就是困难的部分！就我个人而言，我发现AI在编写汇编时最有帮助。我的主要汇编经验是完成microcorruption CTF。我实际上从未自己写过汇编。我很难弄清楚需要的确切指令以及如何修改它们以获得我想要的输出。有了AI，我可以给我的编码代理一个关于JIT编译器如何工作的大致形状，它可以为我处理很多这些细节。为了完成我们的编译器，我们需要实际加载代码。为此，我们将使用mmap分配一块可读、可写且可执行的内存。然后我们将代码复制到该内存中，并将该内存块转换为一个函数，然后调用它：const BSTACK_MAX: usize = 4096; // 这些函数包含在mac系统库中 unsafe extern "C" { fn pthread_jit_write_protect_np(enabled: libc::c_int); fn sys_icache_invalidate(start: *mut libc::c_void, len: libc::size_t); } type MatchFn = unsafe extern "C" fn(input: *const u8, bstack: *mut u64) -> u64; struct Jit { buf: *mut u32, nbytes: usize, bstack: Vec<u64>, } impl Jit { fn compile(regex: Node) -> Jit { let nwords = PROLOGUE_WORDS + node_words(regex) + MATCH_WORDS + FAIL_WORDS; let nbytes = nwords * 4; unsafe { let buf = libc::mmap( std::ptr::null_mut(), nbytes, libc::PROT_READ | libc::PROT_WRITE | libc::PROT_EXEC, libc::MAP_PRIVATE | libc::MAP_ANON | libc::MAP_JIT, -1, 0, ) as *mut u32; assert!(buf as *mut libc::c_void != libc::MAP_FAILED, "mmap failed"); let code = generate_code(regex, buf as u64); pthread_jit_write_protect_np(0); // 使区域可写（Apple W^X） std::slice::from_raw_parts_mut(buf, code.len()).copy_from_slice( &code); pthread_jit_write_protect_np(1); // 恢复可执行 sys_icache_invalidate(buf as *mut libc::c_void, nbytes); Jit { buf, nbytes, bstack: vec![0; BSTACK_MAX * 2] } } } // 运行生成的代码。输入必须以NUL字节结尾。 fn is_match(&mut self, nul_terminated: &[u8]) -> bool { debug_assert_eq!(nul_terminated.last(), Some(&0)); unsafe { let matcher: MatchFn = std::mem::transmute(self.buf); matcher(nul_terminated.as_ptr(), self.bstack.as_mut_ptr()) != 0 } } } impl Drop for Jit { fn drop(&mut self) { unsafe { libc::munmap(self.buf as *mut libc::c_void, self.nbytes); } } }
8. 性能对比与结论
完成所有这些后，让我们比较我们构建的不同实现的性能：输入长度 | 解释器 | JIT | 手写 | JIT加速比 | 手写加速比 9 | 45 ns | 3.8 ns | 3.8 ns | 11.7x | 11.9x 33 | 103 ns | 7.9 ns | 10.5 ns | 13.0x | 9.8x 129 | 597 ns | 30 ns | 32 ns | 19.7x | 18.6x 513 | 1,955 ns | 126 ns | 120 ns | 15.5x | 16.2x 2,049 | 8,301 ns | 470 ns | 393 ns | 17.7x | 21.1x 所以JIT和手写实现几乎并驾齐驱。有时JIT版本更快，有时手写版本更快。有一个流传的梗说AI没有帮助，因为“代码从来不是困难的部分”。我认为在某些领域这是真的，但在其他领域，编写代码绝对是困难的部分。JIT编译器就是一个很好的例子。对于许多软件来说，JIT编译器将大大有助于加速代码。JIT编译器的稀缺让我相信，历史上实现JIT编译器太困难，不值得。LLM降低了入门门槛，使得编写JIT编译器变得容易得多。这是pgrust背后的论点。数据库历来是最难构建的软件，并因此受到限制。现在，有了AI，我们可以对构建的软件类型更加雄心勃勃。感谢阅读，如果你想支持这个项目，支持pgrust的最佳方式是给我们在GitHub上点星。如果你想跟进：分享这个：在X上分享（在新窗口打开）X 在LinkedIn上分享（在新窗口打开）LinkedIn 在Facebook上分享（在新窗口打开）Facebook 通过电子邮件发送给朋友（在新窗口打开）电子邮件 打印（在新窗口打开）打印 留下回复 取消回复 您的电子邮件地址不会被公布。必填项已标记* 评论 * 姓名 * 电子邮件 * 网站 通过电子邮件通知我后续评论。通过电子邮件通知我新帖子。 Δ
5微秒内完成JIT编译：从零构建高性能正则引擎
1. 引言：JIT编译的挑战与机遇
历史上，快速JIT编译（Just-In-Time compilation）被视为一项黑魔法。要编写一个快速的JIT编译器，你需要精通汇编语言。一个典型的例子是：目前没有任何生产级数据库拥有自己的JIT编译器，它们要么使用LLVM，要么生成C/C++代码。这两种方案都存在编译时间过长的问题，限制了其适用性。如今，借助AI，通过直接针对汇编语言编写JIT编译器，实现快速编译变得前所未有的简单。这也是新数据库超越旧数据库的一个机会领域。在构建pgrust时，我最初认为实现JIT编译器会非常困难。最终，在AI的辅助下，我发现它比预期的要容易得多，这也成为了pgrust性能卓越的部分原因。pgrust的JIT编译器在约5微秒内完成代码编译，这使我们能够对每条SQL查询进行JIT编译，而不仅仅是其中的一部分。在这篇文章中，我将带你了解如何构建自己的快速JIT编译器。我们将以一个使用JIT编译的简单正则表达式引擎为例。
2. 为什么需要JIT编译
JIT编译是指在运行时（即“即时”）生成编译代码的实践。如果做得好，它可以带来巨大的性能提升，通常是2-5倍，有时甚至更多。JIT编译的主要应用场景是当你在运行时获得的信息会显著改变程序行为时。这在编程语言解释器中尤为常见，因为它们接收要执行的代码是在运行时。JIT编译器在编程语言之外的领域也很有用，例如数据解析。有时，直到运行时你才知道要解析的数据的模式（schema），JIT可以帮助解决这个问题。让我们先实现一个玩具正则表达式引擎。为了简单起见，我们只支持两个特性：字面量字符串和重复（即正则表达式中的*）。我们还将跳过解析器，直接将正则表达式表示为已解析的Rust结构。这意味着我们将能够支持诸如“apples”和“b(an)*”这样的字符串，但不支持交替（alternation）或后顾（lookbehind）等。在代码中，这非常简单。我们将有三种节点类型：字面量字符串节点、重复节点和连接节点（即两个节点的组合）。代码如下：enum Node { Literal( 'static str), Concatenation(Box<Node>, Box<Node>), Repetition(Box<Node>) } fn literal(text: 'static str) -> Node { Node::Literal(text) } fn concatenation(left: Node, right: Node) -> Node { Node::Concatenation(Box::new(left), Box::new(right)) } fn repetition(body: Node) -> Node { Node::Repetition(Box::new(body)) }
3. 解释器实现与性能瓶颈
为我们的正则表达式引擎编写解释器也很直接：fn match_node(node: Node, input: [u8], pos: usize, next: dyn Fn(usize) -> bool) -> bool { match node { Node::Literal(text) => { let literal = text.as_bytes(); input[pos..].starts_with(literal) && next(pos + literal.len()) } Node::Concatenation(left, right) => { match_node(left, input, pos, |left_end| { match_node(right, input, left_end, next) }) } Node::Repetition(body) => { match_node(body, input, pos, |body_end| { match_node(node, input, body_end, next) }) || next(pos) } } } fn interp_match(regex: Node, input: str) -> bool { let bytes = input.as_bytes(); match_node(regex, bytes, 0, |pos| pos == bytes.len()) } 这个正则表达式引擎非常简单，代码不到20行，但让我们看看它的性能表现。为了对比，我们将代码与针对该正则表达式专门实现的手写代码进行比较。我们的示例使用正则表达式“b(an)*”。手写代码看起来像这样：fn handwritten_b_an_star(input: str) -> bool { let bytes = input.as_bytes(); let mut pos = 0; if pos == bytes.len() || bytes[pos] != b'b' { return false; } pos += 1; while pos < bytes.len() { if bytes[pos] != b'a' { return false; } pos += 1; if pos == bytes.len() || bytes[pos] != b'n' { return false; } pos += 1; } true } （有方法可以优化这段代码并使其更快，但为了我们的目的，它作为一个很好的比较基准。）当我针对这两个实现进行基准测试时，手写版本比解释器快10-20倍。显然，还有很大的改进空间。现在，让我们看看如何使用JIT编译来获得一个性能与手写版本相当的正则表达式引擎。
4. JIT编译的步骤与设计决策
JIT编译代码有两个步骤。首先，为要运行的代码生成汇编。一旦有了代码，就将汇编代码打包成一个函数，你可以像调用程序中的任何其他代码一样调用它。为了生成汇编，我们将使用一种称为“复制并修补”（copy-and-patch）的方法的变体。其思想是，我们有一系列针对要JIT编译的不同操作的汇编模板。这些模板被称为“模板”（stencils）。当我们想要JIT编译一个操作时，我们取出相关的模板，并根据操作的具体情况做小的调整。这与填充真实的模板非常相似。通过将几个填充好的模板串在一起，我们可以在运行时构建一个与手写版本性能相似的程序。我们将采取以下路径：首先，我们查看为“b(an)”生成的ARM64代码。然后，我们将重复的指令序列转换为可重用的模板，编写一个发射器（emitter），从正则表达式AST填充和组合这些模板，最后将生成的指令复制到可执行内存中，以便Rust可以像普通函数一样调用它们。为了让你了解这是如何工作的，最容易的是从生成的代码开始，然后回溯到JIT编译器本身。再次强调，我们处理的是正则表达式“b(an)”。为了布局一些设计决策：我们将使用一个栈进行回溯。栈将跟踪如果正则表达式遇到死胡同我们应该去的状态。我们匹配的字符串将以空字节结尾。这意味着，如果我们到达字符串末尾，任何字符比较都会自动失败。这意味着我们不需要在任何点进行长度比较。对于程序的状态，我们将使用以下寄存器：x0 – 当前字符串位置和返回值；x1 – 用于回溯的栈顶；x2 – 用于回溯的栈底（这用于确定栈是否为空）；x9 – 用作临时变量。对于程序的输入，我们将传递：x0 – 指向字符串开头的指针；x1 – 指向我们将用作栈的位置的指针。
5. 生成的ARM64代码详解
现在我们已经处理了这些，让我们逐部分讲解生成的汇编。这特别针对macOS上的ARM64。首先，我们有初始化程序的序言（prologue）。它所做的只是通过将栈顶和栈底设置为传入的值来初始化栈：0: aa0103e2 mov x2, x1 接下来，我们有检查字符“b”的代码。如果它看到一个不是“b”的字符，我们跳转到一个处理回退逻辑的代码块。否则，我们推进字符串中的位置：; CHAR 'b' 4: 39400009 ldrb w9, [x0] ; 加载当前输入字节 8: 7101893f cmp w9, #0x62 ; 是'b'吗？ c: 54000281 b.ne 0x5c ; 否 - 回退块 10: 91000400 add x0, x0, #1 ; 是 - 推进输入 接下来，我们有重复部分“(an)*”。对于重复，我们需要进行回溯。如果我们在这里回溯，那意味着我们立即跳转到循环的末尾。这意味着我们需要将循环后的指令地址和我们在字符串中的位置都存储在栈上。14: d2800989 movz x9, #0x004c ; 构建恢复地址 18: f2a00009 movk x9, #0x0000, lsl #16 ; = 0x1_0000_004c 1c: f2c00029 movk x9, #0x0001, lsl #32 ; (循环出口) 20: f2e00009 movk x9, #0x0000, lsl #48 ; 24: a8810029 stp x9, x0, [x1], #16 ; 将(exit, pos)压入栈 有了这些，我们现在可以执行重复的主体。这将检查字符'a'和'n'，如果看到它们，就回到重复的顶部，但使用新的字符串位置。; CHAR 'a' 28: 39400009 ldrb w9, [x0] 2c: 7101853f cmp w9, #0x61 ; 'a'? 30: 54000161 b.ne 0x5c ; 否 - 回退块 34: 91000400 add x0, x0, #1 ; CHAR 'n' 38: 39400009 ldrb w9, [x0] 3c: 7101b93f cmp w9, #0x6e ; 'n'? 40: 540000e1 b.ne 0x5c ; 否 - 回退块 44: 91000400 add x0, x0, #1 ; JMP 48: 17fffff3 b 0x14 ; 回到循环顶部 现在我们过了循环。这是回溯时跳转的地方。一旦我们完成重复，我们就到了正则表达式的末尾。我们现在要做的就是检查是否在字符串的末尾。如果在字符串末尾，我们返回1表示成功。如果不是，那意味着正则表达式匹配失败，我们需要运行失败逻辑来进行回退。4c: 39400009 ldrb w9, [x0] 50: 35000069 cbnz w9, 0x5c ; 不在NUL - 回退块 54: d2800020 mov x0, #1 ; 成功 58: d65f03c0 ret 最后，我们有回退逻辑。这检查栈是否为空。如果为空，我们返回0。如果不为空，我们从栈中弹出回退地址和回退字符串位置，然后跳转到回退地址。5c: eb02003f cmp x1, x2 ; 还有帧吗？ 60: 54000060 b.eq 0x6c ; 否 - 放弃 64: a9ff0029 ldp x9, x0, [x1, #-16]! ; 弹出(resume, pos) 68: d61f0120 br x9 ; 跳转到那里 6c: d2800000 mov x0, #0 ; 不匹配 70: d65f03c0 ret
6. 构建模板（Stencils）
既然你已经看到了编译后的代码，你应该开始理解复制并修补编译器是如何工作的。我们有共同的指令集，它们之间只有微小的差异。对于这些功能块中的每一个，我们可以编写一个函数来生成相应的代码。每个函数将接受用于修改代码的值。例如，stencil_char的参数之一将是正则表达式中要比较的字符。我们将直接将该字符插入机器代码中。序言很简单，因为它只是一个代码块：const PROLOGUE_WORDS: usize = 1; fn stencil_prologue() -> [u32; PROLOGUE_WORDS] { [0xAA0103E2] // mov x2, x1 } 对于字符比较，我们需要插入我们正在比较的字符以及跳转到回退逻辑的位置：const CHAR_WORDS: usize = 4; fn stencil_char(byte: u8, stencil_pos: usize, fail_pos: usize) -> [u32; CHAR_WORDS] { [ 0x39400009, // ldrb w9, [x0] 0x7100013F | ((byte as u32) << 10), // cmp w9, #byte 0x54000001 | cond_branch_offset(stencil_pos + 2, fail_pos), // b.ne fail 0x91000400, // add x0, x0, #1 ] } 对于重复，我们有循环开始处压栈的代码和跳转到末尾的代码：const SPLIT_WORDS: usize = 5; fn stencil_split(resume_addr: u64) -> [u32; SPLIT_WORDS] { [ 0xD2800009 | addr_bits(resume_addr, 0), // movz x9, #addr[0..16] 0xF2A00009 | addr_bits(resume_addr, 1), // movk x9, #addr[16..32], lsl 16 0xF2C00009 | addr_bits(resume_addr, 2), // movk x9, #addr[32..48], lsl 32 0xF2E00009 | addr_bits(resume_addr, 3), // movk x9, #addr[48..64], lsl 48 0xA8810029, // stp x9, x0, [x1], #16 ] } const JMP_WORDS: usize = 1; fn stencil_jmp(stencil_pos: usize, target_pos: usize) -> [u32; JMP_WORDS] { [0x14000000 | branch_offset(stencil_pos, target_pos)] // b target } 然后我们有匹配和失败块，它们相当简洁：const MATCH_WORDS: usize = 4; fn stencil_match(stencil_pos: usize, fail_pos: usize) -> [u32; MATCH_WORDS] { [ 0x39400009, // ldrb w9, [x0] 0x35000009 | cond_branch_offset(stencil_pos + 1, fail_pos), // cbnz w9, fail 0xD2800020, // mov x0, #1 0xD65F03C0, // ret ] } const FAIL_WORDS: usize = 6; fn stencil_fail() -> [u32; FAIL_WORDS] { [ 0xEB02003F, // cmp x1, x2 0x54000060, // b.eq +3 (到下面的mov) 0xA9FF0029, // ldp x9, x0, [x1, #-16]! 0xD61F0120, // br x9 0xD2800000, // mov x0, #0 0xD65F03C0, // ret ] } 为了完整起见，这里是我们使用的辅助函数，它们帮助我们向指令中插入特定数据：// 计算条件分支（b.ne / cbnz）的分支偏移字段： // 从分支到目标的指令数，存储在位5..24。 fn cond_branch_offset(branch_pos: usize, target_pos: usize) -> u32 { let instr_count = target_pos as i64 - branch_pos as i64; // 可能为负 (((instr_count as u64) & 0x7FFFF) << 5) as u32 } // 计算无条件分支（b）的分支偏移字段： // 同样的想法，但存储在位0..26。 fn branch_offset(branch_pos: usize, target_pos: usize) -> u32 { let instr_count = target_pos as i64 - branch_pos as i64; // 可能为负 ((instr_count as u64) & 0x3FF_FFFF) as u32 } // 提取绝对地址的16位，为movz/movk立即数定位。 fn addr_bits(addr: u64, part: usize) -> u32 { (((addr >> (16 * part)) & 0xFFFF) as u32) << 5 }
7. 发射代码与加载机器码
现在驱动它的代码：// 计算一个节点编译成多少条指令。 fn node_words(node: Node) -> usize { match node { Node::Literal(text) => text.len() * CHAR_WORDS, Node::Concatenation(left, right) => node_words(left) + node_words(right), Node::Repetition(body) => SPLIT_WORDS + node_words(body) + JMP_WORDS, } } struct Emitter { code: Vec<u32>, fail: usize, // 共享失败块的单词偏移 base: u64, // code[0]的运行时地址，用于绝对地址空洞 } impl Emitter { // 返回下一条指令将被放置的偏移量。 fn pos(&self) -> usize { self.code.len() } // 将填充的模板追加到代码缓冲区。 fn emit(&mut self, stencil: &[u32]) { self.code.extend_from_slice(stencil); } // 为一个节点发射代码，递归到子节点。 fn emit_node(&mut self, node: Node) { match node { Node::Literal(text) => { for byte in text.as_bytes() { self.emit( &stencil_char(byte, self.pos(), self.fail)); } } Node::Concatenation(left, right) => { self.emit_node(left); self.emit_node(right); } Node::Repetition(body) => { let split_at = self.pos(); let exit = split_at + SPLIT_WORDS + node_words(body) + JMP_WORDS; self.emit( &stencil_split(self.base + exit as u64 * 4)); self.emit_node(body); self.emit( &stencil_jmp(self.pos(), split_at)); } } } } // 生成完整程序：序言、编译后的AST、MATCH、失败块。 fn generate_code(regex: Node, base: u64) -> Vec<u32> { let nwords = PROLOGUE_WORDS + node_words(regex) + MATCH_WORDS + FAIL_WORDS; let mut emitter = Emitter { code: Vec::with_capacity(nwords), fail: nwords - FAIL_WORDS, base, }; emitter.emit( &stencil_prologue()); emitter.emit_node(regex); let match_at = emitter.pos(); emitter.emit( &stencil_match(match_at, emitter.fail)); emitter.emit( &stencil_fail()); assert_eq!(emitter.pos(), nwords); emitter.code } 这就是困难的部分！就我个人而言，我发现AI在编写汇编时最有帮助。我的主要汇编经验是完成microcorruption CTF。我实际上从未自己写过汇编。我很难弄清楚需要的确切指令以及如何修改它们以获得我想要的输出。有了AI，我可以给我的编码代理一个关于JIT编译器如何工作的大致形状，它可以为我处理很多这些细节。为了完成我们的编译器，我们需要实际加载代码。为此，我们将使用mmap分配一块可读、可写且可执行的内存。然后我们将代码复制到该内存中，并将该内存块转换为一个函数，然后调用它：const BSTACK_MAX: usize = 4096; // 这些函数包含在mac系统库中 unsafe extern "C" { fn pthread_jit_write_protect_np(enabled: libc::c_int); fn sys_icache_invalidate(start: *mut libc::c_void, len: libc::size_t); } type MatchFn = unsafe extern "C" fn(input: *const u8, bstack: *mut u64) -> u64; struct Jit { buf: *mut u32, nbytes: usize, bstack: Vec<u64>, } impl Jit { fn compile(regex: Node) -> Jit { let nwords = PROLOGUE_WORDS + node_words(regex) + MATCH_WORDS + FAIL_WORDS; let nbytes = nwords * 4; unsafe { let buf = libc::mmap( std::ptr::null_mut(), nbytes, libc::PROT_READ | libc::PROT_WRITE | libc::PROT_EXEC, libc::MAP_PRIVATE | libc::MAP_ANON | libc::MAP_JIT, -1, 0, ) as *mut u32; assert!(buf as *mut libc::c_void != libc::MAP_FAILED, "mmap failed"); let code = generate_code(regex, buf as u64); pthread_jit_write_protect_np(0); // 使区域可写（Apple W^X） std::slice::from_raw_parts_mut(buf, code.len()).copy_from_slice( &code); pthread_jit_write_protect_np(1); // 恢复可执行 sys_icache_invalidate(buf as *mut libc::c_void, nbytes); Jit { buf, nbytes, bstack: vec![0; BSTACK_MAX * 2] } } } // 运行生成的代码。输入必须以NUL字节结尾。 fn is_match(&mut self, nul_terminated: &[u8]) -> bool { debug_assert_eq!(nul_terminated.last(), Some(&0)); unsafe { let matcher: MatchFn = std::mem::transmute(self.buf); matcher(nul_terminated.as_ptr(), self.bstack.as_mut_ptr()) != 0 } } } impl Drop for Jit { fn drop(&mut self) { unsafe { libc::munmap(self.buf as *mut libc::c_void, self.nbytes); } } }
8. 性能对比与结论
完成所有这些后，让我们比较我们构建的不同实现的性能：输入长度 | 解释器 | JIT | 手写 | JIT加速比 | 手写加速比 9 | 45 ns | 3.8 ns | 3.8 ns | 11.7x | 11.9x 33 | 103 ns | 7.9 ns | 10.5 ns | 13.0x | 9.8x 129 | 597 ns | 30 ns | 32 ns | 19.7x | 18.6x 513 | 1,955 ns | 126 ns | 120 ns | 15.5x | 16.2x 2,049 | 8,301 ns | 470 ns | 393 ns | 17.7x | 21.1x 所以JIT和手写实现几乎并驾齐驱。有时JIT版本更快，有时手写版本更快。有一个流传的梗说AI没有帮助，因为“代码从来不是困难的部分”。我认为在某些领域这是真的，但在其他领域，编写代码绝对是困难的部分。JIT编译器就是一个很好的例子。对于许多软件来说，JIT编译器将大大有助于加速代码。JIT编译器的稀缺让我相信，历史上实现JIT编译器太困难，不值得。LLM降低了入门门槛，使得编写JIT编译器变得容易得多。这是pgrust背后的论点。数据库历来是最难构建的软件，并因此受到限制。现在，有了AI，我们可以对构建的软件类型更加雄心勃勃。感谢阅读，如果你想支持这个项目，支持pgrust的最佳方式是给我们在GitHub上点星。如果你想跟进：分享这个：在X上分享（在新窗口打开）X 在LinkedIn上分享（在新窗口打开）LinkedIn 在Facebook上分享（在新窗口打开）Facebook 通过电子邮件发送给朋友（在新窗口打开）电子邮件 打印（在新窗口打开）打印 留下回复 取消回复 您的电子邮件地址不会被公布。必填项已标记* 评论 * 姓名 * 电子邮件 * 网站 通过电子邮件通知我后续评论。通过电子邮件通知我新帖子。 Δ