Achieving JIT Compilation in $5\mu s$

Historically, the ability to create a high-performance Just-In-Time (JIT) compiler was considered a common skill "black art." To succeed, a developer needed a deep, intrinsic mastery of assembly language.

Because of this steep learning curve, you will notice a glaring gap in the current market: no production-ready database currently features its own integrated JIT compiler. Existing alternatives typically suffer from excessive compilation overhead, which severely restricts where they can be practically applied.

The AI Paradigm Shift

The landscape has changed. With the advent of AI, targeting assembly directly has become significantly more accessible. This creates a massive opportunity for modern databases to leapfrog legacy systems.

While building pgrust, I initially feared that implementing a JIT compiler would be an insurmountable hurdle. However, thanks to AI assistance, the process was surprisingly smooth. This implementation is a core reason why pgrust achieves such high speeds.

The Result: The pgrust JIT compiler operates in approximately $5\mu s$. This efficiency allows us to JIT compile every single SQL query, rather than just a selective few.

---

Building a Fast JIT: A Practical Example

To demonstrate how to build a high-speed JIT compiler, we will create a simplified regular expression engine.

1. Scope and Constraints

To keep the demonstration focused, we will limit the engine to:

• Literal strings
• Repetition (the * operator)

We will bypass the parsing stage and assume the regex is already represented as Rust structures. For instance, we can handle apples b(an)*, but we will not support lookbehinds or alternation.

2. The AST Structure

We define three primary node types: Literal, Concatenation (joining two nodes), and Repetition.

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. The Interpreter Approach

A standard interpreter for this logic is straightforward to write:

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

---

Performance Analysis

While the interpreter is concise (under 20 lines), it is slow. To see the difference, let's compare it to a handwritten function specifically optimized for the regex 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 
}

Comparison Table

---

The JIT Strategy: Copy-and-Patch

To achieve handwritten speeds with general-purpose flexibility, we use a technique called copy-and-patch.

The Concept: Instead of generating assembly from scratch, we use pre-defined assembly stencils for each operation. We then "patch" these stencils with specific runtime values and chain them together.

The JIT Pipeline

Implementation Roadmap

• Analyze ARM64 machine code for b(an)*.
• Convert instruction sequences into reusable stencils.
• Build an emitter to combine stencils based on the AST.
• Transfer the final binary blob into executable memory for Rust to invoke.

Design Note: To handle the complexities of regex matching, we will utilize a stack to manage backtracking states.