Cranelift in 461 lines: AST to function pointer

The implementation manual for our stub: a toy language compiled to callable machine code, and the entire cranelift JIT recipe fits in one file. This chapter builds the recipe step by step — what a JIT library actually has to hand you, the object ladder, the declare/define/finalize ceremony, the per-node translation table, and the lifetime contract that makes the final transmute sound — then maps each step into jit.rs. Read it before touching experiments/src/jit.rs, because every ceremony the stub needs appears here first.

The problem in one sentence

M19 needs to turn an Expr tree into a fn(*const f64) -> f64 it can call millions of times — which means generating machine code into executable memory at runtime, in ~tens of microseconds, without ever letting the function pointer outlive the memory it points into.

The concepts, step by step

Step 1 — what a JIT library does: IR in, function pointer out

Cranelift is a code generator: you hand it a function written in CLIF (Cranelift’s intermediate representation — typed instructions like fadd/load organized in basic blocks), and it gives back native machine code placed in executable memory, plus a raw pointer you can call. CLIF is in SSA form (static single assignment — every value is defined exactly once; re-assignment becomes new values, and control-flow merges pass values as block parameters). You never write SSA by hand: a helper called FunctionBuilder maintains it while you emit instructions one at a time. So the whole job of our stub is a recursive walk: Expr node in, CLIF instruction out, then one call to compile.

Step 2 — the object ladder (compare wgpu’s, topic 18)

Like every runtime-code system, cranelift splits expensive long-lived containers from cheap per-function scratch:

 JITBuilder ──► JITModule            (owns memory for code+data)
                  ├─ ctx: codegen::Context     (one function's CLIF)
                  ├─ builder_context: FunctionBuilderContext (reused scratch)
                  └─ declare/define/finalize API

 FunctionBuilder(&mut ctx.func)      (SSA construction helper —
                                      you emit ops, IT handles
                                      block params/phi nodes)

Same shape as topic 18’s Instance→Device→Pipeline: expensive long-lived containers (JITModule owns the executable memory), cheap per-function contexts (reused between compiles), and an explicit “finalize” moment after which you hold a raw pointer. Why it matters: the ladder tells you what to hoist — create the module once, reuse the contexts per expression (question 3 measures the difference).

Step 3 — the compile ladder: declare, define, finalize (memorize this)

Compilation is a fixed seven-rung sequence — the split between define (generate code) and finalize (patch relocations — addresses of other functions/data unknown until everything is placed) is the part that surprises:

 1. translate AST → CLIF            (FunctionTranslator walk)
 2. module.declare_function(name, Linkage::Export, &sig)  → id
 3. module.define_function(id, &mut ctx)   ← compilation happens
 4. module.clear_context(&mut ctx)         ← reuse scratch
 5. module.finalize_definitions()          ← relocations patched
 6. module.get_finalized_function(id)      → *const u8   (:90)
 7. unsafe { mem::transmute::<_, fn(f64...)->f64>(ptr) }

The same ladder as our stub will run it:

#![allow(unused)]
fn main() {
// CLIF in, callable pointer out — the whole recipe
fn compile(&mut self, expr: &Expr) -> fn(*const f64) -> f64 {
    let mut b = FunctionBuilder::new(&mut self.ctx.func, &mut self.b_ctx);
    let block = b.create_block();
    b.append_block_params_for_function_params(block);
    b.switch_to_block(block);
    b.seal_block(block);                          // one block: seal immediately
    let row_ptr = b.block_params(block)[0];
    let v = translate(&mut b, expr, row_ptr);     // the §Step-4 table, recursively
    b.ins().return_(&[v]);
    b.finalize();
    let id = self.module.declare_function("f", Linkage::Export, &sig)?;
    self.module.define_function(id, &mut self.ctx)?;  // ← compilation happens
    self.module.clear_context(&mut self.ctx);
    self.module.finalize_definitions()?;              // ← relocations patched
    unsafe { mem::transmute(self.module.get_finalized_function(id)) }
}   // sound only while the JITModule lives — CompiledExpr must own it
}

The SSA ceremony (create_block, append_block_params..., switch_to_block, seal_block — sealing tells the builder no more predecessors will arrive, so it can resolve block params) collapses to four lines because a pure expression needs exactly one block.

Step 4 — translating an expression: one CLIF op per Expr node

The demo’s translator (jit.rs:189+) is statement-oriented; our Expr is pure — simpler. The entire translation is this table, applied by recursion:

 Col(i)   → load: builder.ins().load(F64, MemFlags::trusted(),
                                     row_ptr, (i*8) as i32)
 Const(c) → builder.ins().f64const(c)
 Add(a,b) → builder.ins().fadd(va, vb)
 Mul(a,b) → builder.ins().fmul(va, vb)
 Lt(a,b)  → cmp = builder.ins().fcmp(FloatCC::LessThan, va, vb)
            → select(cmp, one, zero)  (we keep f64 1.0/0.0)
 And(a,b) → both sides as f64 0/1 → fmin or fmul (branch-free —
            topic 17's predication instinct, now in codegen)

Signature: fn(*const f64) -> f64 — one pointer param (AbiParam::new(types::I64) or a real pointer type via module.target_config().pointer_type()), one F64 return. Note the comparisons stay branch-free (fcmp + select, values not jumps) — generated straight-line code with no control flow is exactly what Step 6’s “quality gap vanishes” claim relies on.

Step 5 — the lifetime contract: the pointer is borrowed, not owned

get_finalized_function returns a raw pointer into memory the JITModule owns; transmute erases that relationship, and Rust can no longer save you. The pointer is valid exactly as long as the JITModule lives — so our CompiledExpr must own the module (CompiledExpr { module, func }; drop order = use-after-free otherwise). postgres solves the same lifetime with per-context resource trackers (llvmjit.c:288); the obligation is universal to JITs, only the spelling differs. The other half of the unsafe contract is the signature: the transmuted type must match the CLIF signature and ABI exactly (question 4 spells out every precondition).

Step 6 — the design point: cranelift vs LLVM, and the gotcha list

                 cranelift            LLVM -O3
 compile speed   ~10-100× faster      baseline
 code quality    ~ -O0..-O1           best
 passes          e-graph based        ~100 passes
                 mid-end (aegraph)
 written in      Rust (no FFI)        C++ (bindgen pain)
 designed for    wasmtime JIT         everything

Cranelift ≈ Umbra’s Flying Start as a design point (fast, single-tier, good-enough). For straight-line f64 arithmetic the quality gap vs LLVM nearly vanishes — no loops to optimize, and OUR loop (over rows) stays in Rust and gets rustc -O.

Gotchas for the stub:

• Version lock: cranelift crates move together — Cargo.toml pins matching versions of cranelift-{jit,module,frontend,codegen,native}.
• cranelift_native::builder() detects the host ISA; enable is_pic false default is fine for JIT.
• MemFlags::trusted() = aligned + notrap: we promise row_ptr is valid — the unsafe contract lives at the eval() call site.
• Floats: use fcmp+select, NOT bint/bitcast tricks — CLIF’s bool handling changed across versions; select on f64 is stable.
• The module must not be dropped: CompiledExpr { module, func } with func called through a stored raw pointer.

Where each step lives in the code

Read jit.rs top to bottom once, then re-read compile() (:55-92) against Step 3’s seven rungs until each line maps to a rung; the FunctionTranslator walk (:189+) is Step 4 with statements added that our pure Expr doesn’t need.

Questions for notes.md

1. Why does define_function (:78) not yet give you a callable — what do relocations still need (addresses of other functions/ data), and which of our Expr nodes would introduce one (none — pure arithmetic; a pow() call would)?
2. FunctionBuilder “handles SSA construction” — what does that mean concretely for a var assigned in two branches (block params instead of phi nodes — how do they differ)?
3. Time compile() in jit_bench across expr depths 2..12. Is it linear in node count? Where does the constant term come from (ISA setup? module init? — hoist GLOBAL vs per-expr state and measure both ways)?
4. The demo transmutes to fn(f64) -> f64. Spell out every precondition that makes our fn(*const f64) -> f64 transmute sound (ABI = System V default? signature match? module alive? W^X handled by JITModule?).
5. M19: eval.rs values aren’t all f64 (nodes, strings, nulls). Which subset of Cypher expressions compiles to this f64 scheme directly, and what’s the fallback boundary (per-node fallback vs whole-expression bailout — pick one and defend it)?

References

Code

• cranelift-jit-demo — src/jit.rs — read it top to bottom; src/frontend.rs (the toy parser) can be skipped, we already have Expr