Compilación JIT de Expresiones Regulares en ARM64: Un Deep Dive en Copy-and-Patch con Rust

La compilación Just-In-Time (JIT) es una técnica fundamental en la optimización de rendimiento de sistemas donde la información crítica para la optimización solo está disponible en tiempo de ejecución. Históricamente, la implementación de compiladores JIT, especialmente aquellos que generan código máquina directamente, ha sido una tarea compleja, requiriendo un profundo conocimiento de la arquitectura del conjunto de instrucciones (ISA) y la gestión de memoria a bajo nivel. Esta complejidad ha limitado su adopción en dominios como las bases de datos, donde a menudo se recurre a backends de compilación como LLVM, que introducen latencias de compilación significativas.

El artículo argumenta que la asistencia de herramientas de IA ha reducido drásticamente la barrera de entrada para la escritura de compiladores JIT de alto rendimiento. Esto abre nuevas oportunidades para aplicar JIT en escenarios donde antes no era viable debido a la sobrecarga de desarrollo o la latencia de compilación. El ejemplo de un motor de expresiones regulares demuestra cómo un JIT puede alcanzar un rendimiento comparable al código escrito a mano, superando a los intérpretes en órdenes de magnitud, incluso para consultas individuales, al compilar en microsegundos.

Arquitectura del Sistema

El sistema se compone de un parser de expresiones regulares (simplificado a un AST pre-construido), un emisor de código JIT y un cargador de código máquina. El AST de la expresión regular se representa mediante un enum Node en Rust, con variantes para literales, concatenaciones y repeticiones. El corazón del JIT es la técnica de copy-and-patch, donde se definen 'stencils' (plantillas) de código ensamblador ARM64 para operaciones comunes (comparación de caracteres, saltos, manejo de repeticiones).

El Emitter recorre el AST y, para cada nodo, selecciona el stencil apropiado, 'parchea' los valores específicos (como el carácter a comparar o la dirección de salto) directamente en las instrucciones de máquina, y las concatena en un buffer de u32. Se utiliza una pila (backtracking stack) para manejar las repeticiones y los fallos de coincidencia, almacenando direcciones de retorno y posiciones en la cadena. La gestión de memoria se realiza con mmap para asignar una región de memoria con permisos de lectura, escritura y ejecución (PROT_READ | PROT_WRITE | PROT_EXEC). Tras copiar el código generado, se utilizan funciones específicas de macOS (pthread_jit_write_protect_np, sys_icache_invalidate) para asegurar la coherencia de la caché de instrucciones y la protección de escritura, antes de transmutar el puntero a la memoria en una función invocable de Rust.

enum Node {
    Literal(&'static str),
    Concatenation(Box<Node>, Box<Node>),
    Repetition(Box<Node>),
}

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

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;
    // ... (copy code, protect, invalidate cache)
    let matcher: MatchFn = std::mem::transmute(self.buf);
    matcher(nul_terminated.as_ptr(), self.bstack.as_mut_ptr()) != 0
}

Fundamentos Teóricos

La idea de la compilación JIT y la optimización de código en tiempo de ejecución tiene raíces profundas en la investigación de lenguajes de programación y sistemas operativos. Conceptos como la generación de código dinámico y la auto-modificación de código se exploraron en los primeros días de la computación. La técnica de copy-and-patch es una forma simplificada de generación de código, a menudo utilizada en compiladores JIT ligeros o en sistemas embebidos, donde la sobrecarga de un compilador completo es inaceptable. El manejo de expresiones regulares mediante autómatas finitos (deterministas o no deterministas) es un pilar de la teoría de la computación, formalizado por Kleene en los años 50. La implementación de un motor de regex con backtracking implícito en el JIT se relaciona con la simulación de un NFA (Non-deterministic Finite Automaton) con una pila explícita para manejar las transiciones epsilon y las alternativas, un concepto bien establecido en la literatura de compiladores y diseño de lenguajes.