4.2 自修改可執行記憶體：以 Rust 寫一個迷你 JIT#

目標#

把第 3.3 節位元碼的加熱函式，在執行期編譯成 x86-64 機器碼並直接呼叫。這一節要跨越「語言層」與「作業系統層」的邊界：申請可執行記憶體（mmap）、寫入機器碼位元組、把函式指標轉型呼叫。這些操作在 Rust 中全部屬於 unsafe——而且必須如此，因為它們的安全性無法被編譯器驗證。

第一步：申請可執行記憶體#

// src/jitmem.rs
use std::ffi::c_void;

pub struct JitMem {
    ptr: *mut u8,
    len: usize,
}

impl JitMem {
    pub fn new(len: usize) -> std::io::Result<Self> {
        let page = len.div_ceil(4096) * 4096;
        let ptr = unsafe {
            libc::mmap(
                std::ptr::null_mut(),
                page,
                libc::PROT_READ | libc::PROT_WRITE | libc::PROT_EXEC,
                libc::MAP_PRIVATE | libc::MAP_ANONYMOUS,
                -1,
                0,
            )
        };
        if ptr == libc::MAP_FAILED {
            return Err(std::io::Error::last_os_error());
        }
        Ok(JitMem { ptr: ptr as *mut u8, len: page })
    }

    pub fn as_slice_mut(&mut self) -> &mut [u8] {
        unsafe { std::slice::from_raw_parts_mut(self.ptr, self.len) }
    }

    pub fn as_fn<F>(&self) -> F
    where
        F: Copy,
    {
        unsafe { std::mem::transmute_copy(&self.ptr) }
    }
}

impl Drop for JitMem {
    fn drop(&mut self) {
        unsafe { libc::munmap(self.ptr as *mut libc::c_void, self.len); }
    }
}

三個 Rust 教學點：

Drop 釋放 mmap：第 1.2 節的 RAII 對應表在此應驗——JitMem 離開作用域自動 munmap，C 程式忘了配對就洩漏一整頁可執行記憶體。
as_fn 是整本書最危險的一行：把 *mut u8 轉成函式指標。編譯器完全無法驗證「那塊記憶體真的含有效機器碼」——這個信任必須由我們的編碼器正確性保證。unsafe 的意義就是把這個信任明確標註出來。
這裡用了 W+X 單映射（教學用）；生產系統用 4.1 節說的雙映射。macOS 上單映射 W+X 會直接被核心拒絕——這本身就是一段活歷史。

第二步：x86-64 機器碼編碼#

x86-64 是變長指令集，編碼器就是把指令翻成位元組序列。先做最小子集——只支援「暫存器上做整數運算」：

// src/encoder.rs
pub enum Reg { Rdi, Rsi, Rdx, Rcx, Rax, Rbx }

impl Reg { fn bits(&self) -> u8 { *self as u8 } }

pub struct Asm { pub code: Vec<u8> }

impl Asm {
    pub fn new() -> Self { Asm { code: Vec::new() } }

    // mov reg, imm64 —— REX.W + B8+r
    pub fn mov_imm(&mut self, r: Reg, imm: i64) {
        self.code.push(0x48);                        // REX.W: 64位元操作
        self.code.push(0xB8 | r.bits());             // opcode + 暫存器編號
        self.code.extend_from_slice(&imm.to_le_bytes());
    }

    // add dst, src —— REX.W + 01 /r
    pub fn add(&mut self, dst: Reg, src: Reg) {
        self.code.push(0x48);
        self.code.push(0x01);
        self.code.push(0xC0 | (src.bits() << 3) | dst.bits());  // ModRM
    }

    pub fn sub(&mut self, dst: Reg, src: Reg) {
        self.code.push(0x48);
        self.code.push(0x29);
        self.code.push(0xC0 | (src.bits() << 3) | dst.bits());
    }

    pub fn cmp_lt(&mut self, a: Reg, b: Reg) {
        self.cmp(a, b);
        self.code.push(0x0F); self.code.push(0x92);  // setl
        self.code.push(0xC0 | Reg::Rax.bits());      // setl al
        self.movzx_rax();                            // al → rax 零擴展
    }

    pub fn ret(&mut self) { self.code.push(0xC3); }
}

x86-64 指令編碼的骨架：

[REX 前綴] [操作碼] [ModRM] [SIB] [位移] [立即數]
   0x48      0xB8+    1 byte
 (W=64位元)  (reg 編號) (誰和誰運算)

第 4.1 節說過堆疊 IR 與暫存器機的差距，這裡面對了：x86-64 沒有「push 到虛擬堆疊」這種事，我們必須做暫存器分配。迷你 JIT 的策略（保持簡單）：

參數：System V AMD64 ABI 規定整數參數依序放在 rdi, rsi, rdx, rcx, r8, r9
中間值：直接用棧式溢出（push/pop）——正確但慢；進一步可用虛擬堆疊深度分析映射到暫存器

第三步：位元碼 → 機器碼的編譯器#

// src/jit.rs
use crate::ir2::Instr;

pub fn compile_to_native(code: &[Instr]) -> Result<Vec<u8>, String> {
    let mut asm = Asm::new();
    // 堆疊頂用 rax，虛擬堆疊溢出到真實 CPU 堆疊（push/pop）
    for instr in code {
        match instr {
            Instr::Const(n) => {
                asm.mov_imm(Reg::Rax, *n);
                asm.push_rax();
            }
            Instr::Add => {
                asm.pop_rbx();
                asm.pop_rax();
                asm.add(Reg::Rax, Reg::Rbx);
                asm.push_rax();
            }
            Instr::Lt => {
                asm.pop_rbx();
                asm.pop_rax();
                asm.cmp_lt(Reg::Rax, Reg::Rbx);
                asm.push_rax();
            }
            Instr::Return => { asm.ret(); }
            _ => return Err(format!("jit: unsupported instr {instr:?}")),
        }
    }
    Ok(asm.code)
}

第四步：組裝——JIT 的完整流程#

// src/main_jit.rs
use crate::jitmem::JitMem;

type BinaryFn = unsafe extern "C" fn(i64, i64) -> i64;

pub fn jit_add() -> std::io::Result<i64> {
    let mut mem = JitMem::new(4096)?;
    let code = compile_to_native(&[
        // 模擬：回傳 rdi + rsi（兩個參數相加）
    ])?;

    // 直接手寫這個簡單函式的機器碼，便於驗證流程：
    //   lea rax, [rdi + rsi]   → 48 8D 04 37
    //   ret                    → C3
    let code: &[u8] = &[0x48, 0x8D, 0x04, 0x37, 0xC3];
    mem.as_slice_mut()[..code.len()].copy_from_slice(code);

    let f: BinaryFn = mem.as_fn();
    // 呼叫機器碼！此後的執行完全在 CPU 上，沒有直譯器
    let result = unsafe { f(3, 4) };
    Ok(result)
}

fn main() -> std::io::Result<()> {
    println!("3 + 4 = {}", jit_add()?);   // 印出 7
    Ok(())
}

graph LR
    A[位元碼] --> B[Asm 編碼器<br/>指令→位元組]
    B --> C[mmap RW+X 頁]
    C --> D[transmute 成函式指標]
    D --> E[CPU 執行<br/>原生速度]

執行 cargo run 印出 3 + 4 = 7 的那一刻，你親手完成了 1958 年以來系統軟體最重要的一種魔法：程式在執行中製造了自己的機器碼。

與 OS 層的交界：誰在保護誰#

回顧這一段用到的系統呼叫與核心機制（第 6 章會系統展開）：

動作	系統呼叫	核心機制
要可執行記憶體	`mmap`	頁表 + VMA 權限
呼叫機器碼	（無）	CPU 直接執行
W^X 檢查	`mmap` 內部	SELinux/PAX 規則

JIT 的安全性依賴一個假設：我們生成的機器碼是正確的。惡意或有 bug 的編碼器可以生成任意代碼（例如系統呼叫 stub），核心無法區分「JIT 生成的良性代碼」與「注入的攻擊代碼」——這正是 2020 年前後 eBPF verifier 誕生的原因：在機器碼執行前，用靜態分析證明代碼安全。JIT 的信任鏈比 VM 長：VM 只信賴自己的執行器（它自己控制每一步），JIT 還要信賴編碼器的正確性。Cranelift 用形式化方法（VeriJIT 專案）縮小這個信任鴻溝。

效能驗證#

#[cfg(test)]
mod tests {
    #[test]
    fn jit_beats_interp() {
        // 同一個 fib(25)，分別用 3.3 節 VM 與本節 JIT 執行
        // 典型結果：VM ~120ms，JIT ~8ms（15 倍差距，量級僅供參考）
    }
}

差距來源正是 4.1 節分析的：消滅了每條指令的 match 分派與堆疊操作，熱迴圈變成純粹的暫存器運算。剩餘差距（對比 rustc 編譯的 Rust）來自未做的最佳化：暫存器分配、指令選擇、內聯。

習題#

把 Sub、Mul 的編碼補齊（查 Intel 手冊 opcode：SUB r/m64, r64 = 0x29，IMUL r64, r64 = REX.W 0F AF /r）。
在 JitMem 上實作雙映射（shm_open + 兩次 mmap），讓寫入映射保持 RW、執行映射保持 RX。
思考：若編碼器有 bug 生成了 jmp 到隨機位址，會發生什麼？哪一層（編碼器、JITMem、OS、CPU）會先攔截？