JIT Compiling Code in 5μs

在 5 微秒内进行 JIT 代码编译

Historically, fast JIT compilation was a black art. To write a fast JIT compiler, you would need to know how to write assembly. Case in point: there is no production-ready database today that has its own JIT compiler. They all either use LLVM or generate C/C++ code. Both of these options suffer from high compile times, which limits their applicability. 从历史上看，快速 JIT（即时）编译是一门“黑魔法”。要编写一个快速的 JIT 编译器，你需要精通汇编语言。事实证明：目前没有任何生产就绪的数据库拥有自己的 JIT 编译器。它们要么使用 LLVM，要么生成 C/C++ 代码。这两种方案都存在编译时间过长的问题，从而限制了它们的适用性。

Now, with the use of AI, it’s easier than ever to write a JIT compiler with fast compile times by directly targeting assembly. This is also one area of opportunity for new databases to improve on old ones. When building pgrust, I initially thought it would be really hard to implement a JIT compiler. In the end, I found it much easier than I expected due to AI assistance and it ends up being part of the reason why pgrust is so fast. 现在，借助人工智能，通过直接针对汇编语言来编写具有快速编译时间的 JIT 编译器变得前所未有的简单。这也是新一代数据库超越旧数据库的一个机会。在构建 pgrust 时，我最初认为实现一个 JIT 编译器会非常困难。最终，由于 AI 的辅助，我发现这比预期的要容易得多，这也成为了 pgrust 速度如此之快的原因之一。

The pgrust JIT compiler compiles code in around 5μs, which enables us to JIT compile every SQL query, not just a subset of them. In this post, I’ll walk you through how you can build your own fast JIT compiler. We’ll build a simple regular expression engine that uses JIT compilation as an example. pgrust 的 JIT 编译器编译代码仅需约 5 微秒，这使我们能够对每一个 SQL 查询进行 JIT 编译，而不仅仅是其中的一部分。在这篇文章中，我将带你了解如何构建自己的快速 JIT 编译器。我们将以构建一个使用 JIT 编译的简单正则表达式引擎为例。

Why JIT Compilation

为什么需要 JIT 编译

JIT compilation is the practice of generating compiled code at runtime or “Just In Time”. When done right, it can result in big performance wins, often on the order of 2-5x and sometimes even more. The main use case for JIT compilation is when there’s information you gain at runtime that drastically alters the behavior of your program. This is particularly common with programming language interpreters; they receive the code to execute at runtime. JIT 编译是在运行时或“即时”生成编译代码的做法。如果运用得当，它可以带来巨大的性能提升，通常在 2-5 倍甚至更多。JIT 编译的主要应用场景是：当你在运行时获得的信息会极大地改变程序行为时。这在编程语言解释器中尤为常见；它们在运行时接收需要执行的代码。

JIT compilers are also useful in domains beyond programming languages, such as parsing data. Sometimes you don’t know the schema of the data you’re parsing until runtime, and a JIT can help with that. JIT 编译器在编程语言之外的领域也很有用，例如数据解析。有时，直到运行时你才知道要解析的数据模式（schema），而 JIT 可以为此提供帮助。

To kick things off, let’s implement a toy regular expression engine. To keep things simple, we’ll support only two features: literal strings and repetition (i.e. the regex *). We’ll also skip the parser and represent the regular expression as already parsed Rust structures. This means we’ll be able to support strings such as: apples b(an)* but no alternation or lookbehind or anything like that. 首先，让我们实现一个玩具正则表达式引擎。为了保持简单，我们只支持两个功能：字面量字符串和重复（即正则表达式中的 *）。我们还将跳过解析器，直接将正则表达式表示为已解析的 Rust 结构。这意味着我们将能够支持诸如 apples 或 b(an)* 之类的字符串，但不支持交替（alternation）、后向查找（lookbehind）等复杂功能。

In code this is pretty simple. We’ll have 3 types of Nodes: a literal string node, a repetition node, and a concatenation node, which is the combination of two nodes. This ends up looking like this: 在代码实现上这非常简单。我们将有 3 种类型的节点：字面量字符串节点、重复节点，以及连接两个节点的连接节点。最终代码如下所示：

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

Writing an interpreter for our regular expression engine is also straightforward: 为我们的正则表达式引擎编写解释器也非常直观：

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())
}

Now this regular expression engine is pretty simple. It’s under 20 lines of code, but let’s see how it does in terms of performance. For comparison, we’ll compare the code against handwritten code implemented specifically for the regex. For our example we’ll use the regex b(an)*. The handwritten code ends up looking like: 这个正则表达式引擎非常简单，代码不到 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
}

(There are ways you could optimize this code and make it much faster, but for our purposes it serves as a good comparison) When I benchmark a couple of examples against these two, I get that the handwritten version is 10-20x faster than the interpreter. Clearly a lot of room for improvement. Now let’s take a look at how we can use JIT compilation to get a general regular expression engine that performs as well as the handwritten version. （虽然有方法可以优化这段代码使其更快，但就我们的目的而言，它是一个很好的对比基准。）当我针对这两个版本进行基准测试时，发现手写版本比解释器版本快 10-20 倍。显然，还有很大的改进空间。现在让我们看看如何使用 JIT 编译来获得一个性能与手写版本相当的通用正则表达式引擎。

How to JIT Compile

如何进行 JIT 编译

There are two steps to JIT compile code. First you generate the assembly for the code you want to run. Once you have the code, you then package the assembly code into a function that you can call like any other code into your program. JIT 编译代码分为两个步骤。首先，为你想要运行的代码生成汇编指令。一旦有了代码，就将这些汇编代码封装成一个函数，使其可以像程序中的任何其他代码一样被调用。

To generate the assembly, we will use a variant of an approach called copy-and-patch. The idea is that we have a series of templates in assembly for the different operations we want to JIT compile. These templates are called “stencils”. When we want to JIT compile an operation, we take the associated stencil and make small tweaks based on the specifics of the operation. Very similar to filling in a real stencil. By stringing together several of these filled stencils, we can construct a program at runtime that has similar performance to the handwritten version. 为了生成汇编代码，我们将使用一种称为“复制与修补”（copy-and-patch）方法的变体。其核心思想是：针对我们想要 JIT 编译的不同操作，准备一系列汇编模板。这些模板被称为“模具”（stencils）。当我们想要 JIT 编译某个操作时，就取出对应的模具，并根据操作的具体细节进行微调。这非常类似于填充真实的模具。通过将多个填充好的模具串联起来，我们可以在运行时构建出一个性能与手写版本相当的程序。

Here’s the path we’ll take: first we’ll look at the ARM64 code we want to generate for b(an)*. Then we’ll turn repeated instruction sequences into reusable stencils, write an emitter that fills and combines those stencils from the regex AST, and finally copy the generated instructions into executable memory so Rust can call them like a normal function. To walk you through how this works, it’s easiest to start with the generated code and work backwards to the JIT compiler itself. Again, we’re working with the regex “b(an)*”. To lay out some design decisions: We’ll use a stack for backtracking. The stack will keep track of the state we should go to if we hit a dead end in the… 我们将采取以下路径：首先，查看我们想要为 b(an)* 生成的 ARM64 代码。然后，将重复的指令序列转化为可重用的模具，编写一个发射器（emitter）来根据正则表达式的抽象语法树（AST）填充并组合这些模具，最后将生成的指令复制到可执行内存中，以便 Rust 可以像调用普通函数一样调用它们。为了向你解释其工作原理，最简单的方法是从生成的代码开始，反向推导到 JIT 编译器本身。再次强调，我们处理的是正则表达式 b(an)*。在设计决策方面：我们将使用一个栈来进行回溯。该栈将记录如果我们陷入死胡同时应该返回的状态……