Chapter 22: JIT and the REPL

New files

src/jit.rs

src/repl.rs

src/bin/fun.rs

Changed files

src/llvm/mod.rs

src/llvm/control.rs

rt/src/heap.rs

src/session.rs

src/rename.rs

build.rs

src/core.rs

src/core_type.rs

src/elaborate.rs

src/erase.rs

src/lib.rs

src/main.rs

Motivation

An executable can run after its compiler has exited. With a JIT, we keep the compiler around and run the generated code inside its process. This gives us something rather fun to build next, a REPL that compiles the definitions we type.

LLVM already knows how to produce machine instructions. The tricky part is keeping them alive for as long as the program needs them. Enter saved = make 40, then call saved 2 on the next line. That closure might hold both a captured integer and an address in the previous input’s code. We must retain both.

We’ll keep the compiler pipeline and calling convention from Chapter 21. Around them we’ll build a session that remembers accepted definitions and leaves them intact when a new input fails. A typo on the next line shouldn’t destroy the work we’ve already done.

The Language, Now

There is no new Fun syntax. The REPL accepts expressions and declarations. Explicit braces put several declarations into one input. A file loaded with :load can use ordinary multiline layout.

make x = \y -> x + y

saved = make 40

Entering saved 2 prints 42. A pure expression is wrapped in the Prelude’s print, so its type must have a Show instance. An IO expression executes its action and discards the result. Inspect a function with :t rather than trying to print it.

Command	Meaning
`:t EXPR`	Infer the principal qualified type. `:type` is an alias.
`:k TYPE`	Infer the kind of a constructor or type application.
`:core EXPR`, `:anf EXPR`, `:llvm EXPR`	Inspect an intermediate program without executing it.
`:load FILE`	Replace definitions after successful checking and initialization.
`:reset`	Release retained code, roots, and definitions.
`:engine NAME`	Choose `jit`, `core`, `anf`, or `flat`. Preserve definitions.
`:quit`	Leave the session. `:q` is an alias.

Loading a file does not execute its main. Enter main explicitly to run it. Duplicate user definitions are errors, as in a source file. Reset or load a replacement to change them. A compiler error leaves accepted definitions intact. Put Prelude overrides in the first declaration block or a loaded file. Later inputs cannot change the names visible to earlier definitions. A fixity must accompany its operator definition in the same input.

Theory

A closure is a pair of code and environment, even when our representation stores several captured fields (Marlow and Peyton Jones 2004). Reachability of the environment is only half the lifetime problem. The code address must remain executable too. The Boehm collector understands heap reachability, but does not own LLVM’s machine code or know when an execution engine can be discarded (Boehm and Weiser 1988).

We give a session ownership of all its execution engines. Each engine owns one input’s module. A separate table holds global words in registered root slots. Earlier modules can therefore remain immutable, and later modules import their values through external global declarations. Reset destroys the engines before unregistering the slots. The enclosing LLVM Context outlives both.

Typing a bad definition shouldn’t destroy the good ones we’ve already entered. We therefore prepare a candidate separately and retain it only after checking and initialization succeed. Inspection commands stop before initialization. Asking for the type of error "x" is safe. Running it is another matter. Keeping those paths separate lets us explore programs without accidentally executing them (Pierce 2002).

Implementation

A returning entry

A native executable starts with main, initializes globals, executes the IO action, and exits. A REPL cannot exit after every successful input. JIT lowering instead emits a C-callable void fun_entry(void): initialize new globals, apply the query action to unit, then return. Fast entries still use tailcc. Closure calls still use the same C eval/apply path. Calls across that ABI boundary are explicitly notail. A wide-call regression caught LLVM inferring a tail transfer from a tailcc function into C apply, which corrupted the callee-pop return frame on AArch64. Verification did not catch this target-level failure. Executing the same program on both backends did. The fix also applies to Chapter 21.

Backend::for_jit selects this entry and external global storage. It reuses Chapter 21’s worklist emission, verifier, and pass manager. We run O1 for the interactive path. Native build retains its O0–O3 choice.

#![allow(unused)]
fn main() {
/// A returning entry and externally rooted globals replace the native process startup.
pub fn for_jit(
    context: &'ctx Context,
    program: &'program flat::Program,
    inherited: BTreeSet<Id>,
) -> Result<Self> {
    Self::create(context, program, 1, true, inherited)
}
}

Resolve symbols before execution

The system linker resolves a native object’s runtime calls. MCJIT needs the same addresses inside the compiler. runtime_symbols takes the addresses of the linked runtime’s C exports, keeping them in the binary. map_symbols checks every external function and installs its address with add_global_mapping. Earlier fast and slow entries are mapped to their retained engines in the same step. LLVM intrinsics remain LLVM’s responsibility.

#![allow(unused)]
fn main() {
use crate::{flat, llvm::Backend};
use fun_rt::heap::GlobalRoot;
use inkwell::{
    OptimizationLevel, context::Context, execution_engine::ExecutionEngine,
};
use std::collections::BTreeMap;
type GlobalKey = (String, usize);
/// Keep every engine alive until no REPL binding can reference its closures.
/// Fields drop in declaration order: release code before unregistering the slots.
pub struct Jit<'ctx> {
    engines: Vec<ExecutionEngine<'ctx>>,
    roots: BTreeMap<GlobalKey, GlobalRoot>,
    queries: Vec<GlobalRoot>,
    code: BTreeMap<GlobalKey, (usize, usize, usize)>,
    types: BTreeMap<GlobalKey, Vec<String>>,
    context: &'ctx Context,
}
impl<'ctx> Jit<'ctx> {
    pub fn new(context: &'ctx Context) -> Self {
        fun_rt::rt_init();
        Self {
            engines: vec![],
            roots: BTreeMap::new(),
            queries: vec![],
            code: BTreeMap::new(),
            types: BTreeMap::new(),
            context,
        }
    }
    /// Compile a checked program, import accepted globals, then initialize only new bindings.
    /// Runtime errors retain the native runtime's process-exit semantics.
    pub fn run(&mut self, program: &flat::Program) -> Result<(), String> {
        let mut occurrences = BTreeMap::new();
        let mut keys = BTreeMap::new();
        for global in &program.globals {
            let count =
                occurrences.entry(global.name.text.clone()).or_insert(0);
            keys.insert(
                global.name.clone(),
                (global.name.text.clone(), *count),
            );
            *count += 1;
        }
        let mut checked_types = BTreeMap::new();
        for (id, shape) in &program.metadata.global_types {
            let Some(key) = keys.get(id) else { continue };
            if let Some(previous) = self.types.get(key)
                && previous != shape
            {
                return Err(format!(
                    "input changes the accepted type of {}; use :reset or :load",
                    key.0
                ));
            }
            checked_types.insert(key.clone(), shape.clone());
        }
        let inherited = keys
            .iter()
            .filter(|(_, key)| self.roots.contains_key(*key))
            .map(|(id, _)| id.clone())
            .collect();
        let backend = Backend::for_jit(self.context, program, inherited)
            .map_err(|e| e.to_string())?;
        backend.optimize(1).map_err(|e| e.to_string())?;
        let engine = backend
            .module
            .create_jit_execution_engine(OptimizationLevel::Less)
            .map_err(|e| format!("JIT setup failed: {e}"))?;
        let mut imports = vec![];
        for global in &program.globals {
            if let Some(f) = global.function
                && let Some((fast, slow, arity)) =
                    self.code.get(&keys[&global.name])
            {
                if *arity != program.funcs[f.0 as usize].params.len() {
                    return Err(format!(
                        "input changes the accepted arity of {}; use :reset or :load",
                        global.name.text
                    ));
                }
                imports.push((format!("fun.f{}", f.0), *fast));
                imports.push((format!("fun.f{}.slow", f.0), *slow));
            }
        }
        map_symbols(&backend.module, &engine, &imports)?;
        let mut pending = BTreeMap::new();
        for global in &program.globals {
            let name = &keys[&global.name];
            let slots = if let Some(slots) = self.roots.get(name) {
                slots
            } else {
                pending
                    .entry(name.clone())
                    .or_insert_with(|| GlobalRoot::new(0))
            };
            let value = backend
                .module
                .get_global(&format!("fun.g{}", global.name.name.unique));
            let Some(value) = value else { continue }; // O1 removed an unused declaration.
            engine.add_global_mapping(&value, slots.as_ptr() as usize);
        }
        let mut pending_code = BTreeMap::new();
        for global in &program.globals {
            let key = &keys[&global.name];
            if let Some(f) = global.function
                && !self.code.contains_key(key)
            {
                let fast = engine
                    .get_function_address(&format!("fun.f{}", f.0))
                    .map_err(|e| e.to_string())?;
                let slow = engine
                    .get_function_address(&format!("fun.f{}.slow", f.0))
                    .map_err(|e| e.to_string())?;
                pending_code.insert(
                    key.clone(),
                    (fast, slow, program.funcs[f.0 as usize].params.len()),
                );
            }
        }
        call_entry(&engine)?;
        // A transient entry gets a fresh slot on every input; retain even that slot
        // while its module exists, so no relocation can ever address freed storage.
        if let Some(query) =
            pending.remove(&keys[&program.metadata.main.name])
        {
            self.queries.push(query);
        }
        let query = &keys[&program.metadata.main.name];
        pending_code.remove(query);
        checked_types.remove(query);
        self.code.append(&mut pending_code);
        self.types.append(&mut checked_types);
        self.roots.append(&mut pending);
        self.engines.push(engine);
        Ok(())
    }
    /// Introspection for tests: accepted globals survive later inputs without reinitialization.
    pub fn global_word(&self, name: &str) -> Option<u64> {
        self.roots
            .iter()
            .rfind(|((text, _), _)| text == name)
            .map(|(_, root)| root.get())
    }
}
/// One explicit symbol table links the JIT to the very same ABI as native objects.
pub fn runtime_symbols() -> Vec<(&'static str, usize)> {
    macro_rules! entries {
        ($module:path => [$($name:ident),* $(,)?]) => {{
            use $module as runtime;
            vec![$((stringify!($name), runtime::$name as *const () as usize)),*]
        }};
    }
    let mut symbols = entries!(fun_rt => [
        rt_init,
        rt_alloc,
        rt_alloc_atomic,
        rt_apply,
        rt_apply1,
        rt_apply2,
        rt_apply3,
        rt_apply4,
        rt_closure_set,
        rt_make_pap,
        rt_div_zero,
        rt_error,
        rt_exit,
        rt_pattern_fail,
    ]);
    symbols.extend(entries!(fun_rt::abi => [
        rt_prim_intAdd,
        rt_prim_intSub,
        rt_prim_intMul,
        rt_prim_intNeg,
        rt_prim_intQuot,
        rt_prim_intRem,
        rt_prim_intDiv,
        rt_prim_intMod,
        rt_prim_intEq,
        rt_prim_intLt,
        rt_prim_intLe,
        rt_prim_dblAdd,
        rt_prim_dblSub,
        rt_prim_dblMul,
        rt_prim_dblDiv,
        rt_prim_dblNeg,
        rt_prim_dblSqrt,
        rt_prim_dblExp,
        rt_prim_dblLog,
        rt_prim_dblSin,
        rt_prim_dblCos,
        rt_prim_dblPow,
        rt_prim_dblEq,
        rt_prim_dblLt,
        rt_prim_dblLe,
        rt_prim_intToDbl,
        rt_prim_dblFloor,
        rt_prim_dblRound,
        rt_prim_dblTruncate,
        rt_prim_chrOrd,
        rt_prim_chrChr,
        rt_prim_chrEq,
        rt_prim_chrLt,
        rt_prim_strAppend,
        rt_prim_strLength,
        rt_prim_strEq,
        rt_prim_strCompare,
        rt_prim_strUnpack,
        rt_prim_strPack,
        rt_prim_showInt,
        rt_prim_showDouble,
        rt_prim_showChar,
        rt_prim_showString,
        rt_prim_readInt,
        rt_prim_error,
        rt_prim_putStr,
        rt_prim_getLine,
        rt_prim_exit,
    ]));
    symbols
}
/// Enter a finalized module without exposing a callable to the rest of the compiler.
fn call_entry(engine: &ExecutionEngine<'_>) -> Result<(), String> {
    // SAFETY: the verified backend emits exactly void fun_entry(void), C ABI.
    // All imports are explicitly mapped; code, constants, and root slots remain live.
    // This engine is never modified after lookup, and no pointer escapes this call.
    unsafe {
        let entry = engine
            .get_function::<unsafe extern "C" fn()>("fun_entry")
            .map_err(|e| e.to_string())?;
        entry.call();
    }
    Ok(())
}
/// Resolve declarations before MCJIT finalizes the module. Intrinsics belong to LLVM.
fn map_symbols(
    module: &inkwell::module::Module<'_>,
    engine: &ExecutionEngine<'_>,
    imports: &[(String, usize)],
) -> Result<(), String> {
    let symbols = runtime_symbols();
    for function in module
        .get_functions()
        .filter(|f| f.count_basic_blocks() == 0)
    {
        let name =
            function.get_name().to_str().map_err(|e| e.to_string())?;
        if name.starts_with("llvm.") {
            continue;
        }
        let address = symbols
            .iter()
            .find(|(symbol, _)| *symbol == name)
            .map(|(_, address)| *address)
            .or_else(|| {
                imports
                    .iter()
                    .find(|(symbol, _)| symbol == name)
                    .map(|(_, address)| *address)
            })
            .ok_or_else(|| {
                format!("unmapped JIT runtime symbol: {name}")
            })?;
        engine.add_global_mapping(&function, address);
    }
    Ok(())
}
}

An unrecognized import is an error before execution. We do not depend on the host executable exporting symbols by name. String operations, formatting, application, errors, and IO all reach the Chapter 22 runtime rlib, whose ABI matches the static library used by native executables.

Keep code and roots together

The session’s ownership is explicit:

#![allow(unused)]
fn main() {
use crate::{flat, llvm::Backend};
use fun_rt::heap::GlobalRoot;
use inkwell::{
    OptimizationLevel, context::Context, execution_engine::ExecutionEngine,
};
use std::collections::BTreeMap;
type GlobalKey = (String, usize);
impl<'ctx> Jit<'ctx> {
    pub fn new(context: &'ctx Context) -> Self {
        fun_rt::rt_init();
        Self {
            engines: vec![],
            roots: BTreeMap::new(),
            queries: vec![],
            code: BTreeMap::new(),
            types: BTreeMap::new(),
            context,
        }
    }
    /// Compile a checked program, import accepted globals, then initialize only new bindings.
    /// Runtime errors retain the native runtime's process-exit semantics.
    pub fn run(&mut self, program: &flat::Program) -> Result<(), String> {
        let mut occurrences = BTreeMap::new();
        let mut keys = BTreeMap::new();
        for global in &program.globals {
            let count =
                occurrences.entry(global.name.text.clone()).or_insert(0);
            keys.insert(
                global.name.clone(),
                (global.name.text.clone(), *count),
            );
            *count += 1;
        }
        let mut checked_types = BTreeMap::new();
        for (id, shape) in &program.metadata.global_types {
            let Some(key) = keys.get(id) else { continue };
            if let Some(previous) = self.types.get(key)
                && previous != shape
            {
                return Err(format!(
                    "input changes the accepted type of {}; use :reset or :load",
                    key.0
                ));
            }
            checked_types.insert(key.clone(), shape.clone());
        }
        let inherited = keys
            .iter()
            .filter(|(_, key)| self.roots.contains_key(*key))
            .map(|(id, _)| id.clone())
            .collect();
        let backend = Backend::for_jit(self.context, program, inherited)
            .map_err(|e| e.to_string())?;
        backend.optimize(1).map_err(|e| e.to_string())?;
        let engine = backend
            .module
            .create_jit_execution_engine(OptimizationLevel::Less)
            .map_err(|e| format!("JIT setup failed: {e}"))?;
        let mut imports = vec![];
        for global in &program.globals {
            if let Some(f) = global.function
                && let Some((fast, slow, arity)) =
                    self.code.get(&keys[&global.name])
            {
                if *arity != program.funcs[f.0 as usize].params.len() {
                    return Err(format!(
                        "input changes the accepted arity of {}; use :reset or :load",
                        global.name.text
                    ));
                }
                imports.push((format!("fun.f{}", f.0), *fast));
                imports.push((format!("fun.f{}.slow", f.0), *slow));
            }
        }
        map_symbols(&backend.module, &engine, &imports)?;
        let mut pending = BTreeMap::new();
        for global in &program.globals {
            let name = &keys[&global.name];
            let slots = if let Some(slots) = self.roots.get(name) {
                slots
            } else {
                pending
                    .entry(name.clone())
                    .or_insert_with(|| GlobalRoot::new(0))
            };
            let value = backend
                .module
                .get_global(&format!("fun.g{}", global.name.name.unique));
            let Some(value) = value else { continue }; // O1 removed an unused declaration.
            engine.add_global_mapping(&value, slots.as_ptr() as usize);
        }
        let mut pending_code = BTreeMap::new();
        for global in &program.globals {
            let key = &keys[&global.name];
            if let Some(f) = global.function
                && !self.code.contains_key(key)
            {
                let fast = engine
                    .get_function_address(&format!("fun.f{}", f.0))
                    .map_err(|e| e.to_string())?;
                let slow = engine
                    .get_function_address(&format!("fun.f{}.slow", f.0))
                    .map_err(|e| e.to_string())?;
                pending_code.insert(
                    key.clone(),
                    (fast, slow, program.funcs[f.0 as usize].params.len()),
                );
            }
        }
        call_entry(&engine)?;
        // A transient entry gets a fresh slot on every input; retain even that slot
        // while its module exists, so no relocation can ever address freed storage.
        if let Some(query) =
            pending.remove(&keys[&program.metadata.main.name])
        {
            self.queries.push(query);
        }
        let query = &keys[&program.metadata.main.name];
        pending_code.remove(query);
        checked_types.remove(query);
        self.code.append(&mut pending_code);
        self.types.append(&mut checked_types);
        self.roots.append(&mut pending);
        self.engines.push(engine);
        Ok(())
    }
    /// Introspection for tests: accepted globals survive later inputs without reinitialization.
    pub fn global_word(&self, name: &str) -> Option<u64> {
        self.roots
            .iter()
            .rfind(|((text, _), _)| text == name)
            .map(|(_, root)| root.get())
    }
}
/// One explicit symbol table links the JIT to the very same ABI as native objects.
pub fn runtime_symbols() -> Vec<(&'static str, usize)> {
    macro_rules! entries {
        ($module:path => [$($name:ident),* $(,)?]) => {{
            use $module as runtime;
            vec![$((stringify!($name), runtime::$name as *const () as usize)),*]
        }};
    }
    let mut symbols = entries!(fun_rt => [
        rt_init,
        rt_alloc,
        rt_alloc_atomic,
        rt_apply,
        rt_apply1,
        rt_apply2,
        rt_apply3,
        rt_apply4,
        rt_closure_set,
        rt_make_pap,
        rt_div_zero,
        rt_error,
        rt_exit,
        rt_pattern_fail,
    ]);
    symbols.extend(entries!(fun_rt::abi => [
        rt_prim_intAdd,
        rt_prim_intSub,
        rt_prim_intMul,
        rt_prim_intNeg,
        rt_prim_intQuot,
        rt_prim_intRem,
        rt_prim_intDiv,
        rt_prim_intMod,
        rt_prim_intEq,
        rt_prim_intLt,
        rt_prim_intLe,
        rt_prim_dblAdd,
        rt_prim_dblSub,
        rt_prim_dblMul,
        rt_prim_dblDiv,
        rt_prim_dblNeg,
        rt_prim_dblSqrt,
        rt_prim_dblExp,
        rt_prim_dblLog,
        rt_prim_dblSin,
        rt_prim_dblCos,
        rt_prim_dblPow,
        rt_prim_dblEq,
        rt_prim_dblLt,
        rt_prim_dblLe,
        rt_prim_intToDbl,
        rt_prim_dblFloor,
        rt_prim_dblRound,
        rt_prim_dblTruncate,
        rt_prim_chrOrd,
        rt_prim_chrChr,
        rt_prim_chrEq,
        rt_prim_chrLt,
        rt_prim_strAppend,
        rt_prim_strLength,
        rt_prim_strEq,
        rt_prim_strCompare,
        rt_prim_strUnpack,
        rt_prim_strPack,
        rt_prim_showInt,
        rt_prim_showDouble,
        rt_prim_showChar,
        rt_prim_showString,
        rt_prim_readInt,
        rt_prim_error,
        rt_prim_putStr,
        rt_prim_getLine,
        rt_prim_exit,
    ]));
    symbols
}
/// Resolve declarations before MCJIT finalizes the module. Intrinsics belong to LLVM.
fn map_symbols(
    module: &inkwell::module::Module<'_>,
    engine: &ExecutionEngine<'_>,
    imports: &[(String, usize)],
) -> Result<(), String> {
    let symbols = runtime_symbols();
    for function in module
        .get_functions()
        .filter(|f| f.count_basic_blocks() == 0)
    {
        let name =
            function.get_name().to_str().map_err(|e| e.to_string())?;
        if name.starts_with("llvm.") {
            continue;
        }
        let address = symbols
            .iter()
            .find(|(symbol, _)| *symbol == name)
            .map(|(_, address)| *address)
            .or_else(|| {
                imports
                    .iter()
                    .find(|(symbol, _)| symbol == name)
                    .map(|(_, address)| *address)
            })
            .ok_or_else(|| {
                format!("unmapped JIT runtime symbol: {name}")
            })?;
        engine.add_global_mapping(&function, address);
    }
    Ok(())
}
/// Enter a finalized module without exposing a callable to the rest of the compiler.
fn call_entry(engine: &ExecutionEngine<'_>) -> Result<(), String> {
    // SAFETY: the verified backend emits exactly void fun_entry(void), C ABI.
    // All imports are explicitly mapped; code, constants, and root slots remain live.
    // This engine is never modified after lookup, and no pointer escapes this call.
    unsafe {
        let entry = engine
            .get_function::<unsafe extern "C" fn()>("fun_entry")
            .map_err(|e| e.to_string())?;
        entry.call();
    }
    Ok(())
}
/// Keep every engine alive until no REPL binding can reference its closures.
/// Fields drop in declaration order: release code before unregistering the slots.
pub struct Jit<'ctx> {
    engines: Vec<ExecutionEngine<'ctx>>,
    roots: BTreeMap<GlobalKey, GlobalRoot>,
    queries: Vec<GlobalRoot>,
    code: BTreeMap<GlobalKey, (usize, usize, usize)>,
    types: BTreeMap<GlobalKey, Vec<String>>,
    context: &'ctx Context,
}
}

A global slot uses the runtime’s GlobalRoot owner. Its boxed storage stays at a stable address and is registered with Boehm. UnsafeCell permits generated code to write the slot without conflicting with Rust’s reference rules. Mapping an LLVM global to that address makes both generated loads and the collector see the same word. A Rust container alone would not provide this guarantee. Boehm does not scan Rust’s ordinary heap transitively.

The source group order distinguishes Prelude globals from user globals that shadow their spelling. Session-local renamer uniques identify LLVM declarations within a compilation, but adding declarations may change their numbers. We need something else for persistent keys. Accepted slots are reused and their initializers are skipped. Before importing a value, an alpha-normalized type fingerprint checks that rechecking has not changed its type, and the saved syntactic arity checks the fast entry’s argument count. Erasure preserves this flat metadata. The fingerprint traversal uses an explicit worklist. Query slots are fresh and retained with their engines, so no surviving relocation points into freed storage.

Jit::run finishes all mappings before retrieving a typed callable. Its unsafe call relies on a few facts we check beforehand. The backend emits exactly the C signature requested from Inkwell. All runtime imports are resolved, and the engine and roots remain alive throughout the call. The callable is used immediately. We never add code to that engine afterward. Inkwell documents that later additions can invalidate a previously retrieved function pointer.

#![allow(unused)]
fn main() {
use crate::{flat, llvm::Backend};
use fun_rt::heap::GlobalRoot;
use inkwell::{
    OptimizationLevel, context::Context, execution_engine::ExecutionEngine,
};
use std::collections::BTreeMap;
type GlobalKey = (String, usize);
/// Keep every engine alive until no REPL binding can reference its closures.
/// Fields drop in declaration order: release code before unregistering the slots.
pub struct Jit<'ctx> {
    engines: Vec<ExecutionEngine<'ctx>>,
    roots: BTreeMap<GlobalKey, GlobalRoot>,
    queries: Vec<GlobalRoot>,
    code: BTreeMap<GlobalKey, (usize, usize, usize)>,
    types: BTreeMap<GlobalKey, Vec<String>>,
    context: &'ctx Context,
}
impl<'ctx> Jit<'ctx> {
    pub fn new(context: &'ctx Context) -> Self {
        fun_rt::rt_init();
        Self {
            engines: vec![],
            roots: BTreeMap::new(),
            queries: vec![],
            code: BTreeMap::new(),
            types: BTreeMap::new(),
            context,
        }
    }
    /// Compile a checked program, import accepted globals, then initialize only new bindings.
    /// Runtime errors retain the native runtime's process-exit semantics.
    pub fn run(&mut self, program: &flat::Program) -> Result<(), String> {
        let mut occurrences = BTreeMap::new();
        let mut keys = BTreeMap::new();
        for global in &program.globals {
            let count =
                occurrences.entry(global.name.text.clone()).or_insert(0);
            keys.insert(
                global.name.clone(),
                (global.name.text.clone(), *count),
            );
            *count += 1;
        }
        let mut checked_types = BTreeMap::new();
        for (id, shape) in &program.metadata.global_types {
            let Some(key) = keys.get(id) else { continue };
            if let Some(previous) = self.types.get(key)
                && previous != shape
            {
                return Err(format!(
                    "input changes the accepted type of {}; use :reset or :load",
                    key.0
                ));
            }
            checked_types.insert(key.clone(), shape.clone());
        }
        let inherited = keys
            .iter()
            .filter(|(_, key)| self.roots.contains_key(*key))
            .map(|(id, _)| id.clone())
            .collect();
        let backend = Backend::for_jit(self.context, program, inherited)
            .map_err(|e| e.to_string())?;
        backend.optimize(1).map_err(|e| e.to_string())?;
        let engine = backend
            .module
            .create_jit_execution_engine(OptimizationLevel::Less)
            .map_err(|e| format!("JIT setup failed: {e}"))?;
        let mut imports = vec![];
        for global in &program.globals {
            if let Some(f) = global.function
                && let Some((fast, slow, arity)) =
                    self.code.get(&keys[&global.name])
            {
                if *arity != program.funcs[f.0 as usize].params.len() {
                    return Err(format!(
                        "input changes the accepted arity of {}; use :reset or :load",
                        global.name.text
                    ));
                }
                imports.push((format!("fun.f{}", f.0), *fast));
                imports.push((format!("fun.f{}.slow", f.0), *slow));
            }
        }
        map_symbols(&backend.module, &engine, &imports)?;
        let mut pending = BTreeMap::new();
        for global in &program.globals {
            let name = &keys[&global.name];
            let slots = if let Some(slots) = self.roots.get(name) {
                slots
            } else {
                pending
                    .entry(name.clone())
                    .or_insert_with(|| GlobalRoot::new(0))
            };
            let value = backend
                .module
                .get_global(&format!("fun.g{}", global.name.name.unique));
            let Some(value) = value else { continue }; // O1 removed an unused declaration.
            engine.add_global_mapping(&value, slots.as_ptr() as usize);
        }
        let mut pending_code = BTreeMap::new();
        for global in &program.globals {
            let key = &keys[&global.name];
            if let Some(f) = global.function
                && !self.code.contains_key(key)
            {
                let fast = engine
                    .get_function_address(&format!("fun.f{}", f.0))
                    .map_err(|e| e.to_string())?;
                let slow = engine
                    .get_function_address(&format!("fun.f{}.slow", f.0))
                    .map_err(|e| e.to_string())?;
                pending_code.insert(
                    key.clone(),
                    (fast, slow, program.funcs[f.0 as usize].params.len()),
                );
            }
        }
        call_entry(&engine)?;
        // A transient entry gets a fresh slot on every input; retain even that slot
        // while its module exists, so no relocation can ever address freed storage.
        if let Some(query) =
            pending.remove(&keys[&program.metadata.main.name])
        {
            self.queries.push(query);
        }
        let query = &keys[&program.metadata.main.name];
        pending_code.remove(query);
        checked_types.remove(query);
        self.code.append(&mut pending_code);
        self.types.append(&mut checked_types);
        self.roots.append(&mut pending);
        self.engines.push(engine);
        Ok(())
    }
    /// Introspection for tests: accepted globals survive later inputs without reinitialization.
    pub fn global_word(&self, name: &str) -> Option<u64> {
        self.roots
            .iter()
            .rfind(|((text, _), _)| text == name)
            .map(|(_, root)| root.get())
    }
}
/// One explicit symbol table links the JIT to the very same ABI as native objects.
pub fn runtime_symbols() -> Vec<(&'static str, usize)> {
    macro_rules! entries {
        ($module:path => [$($name:ident),* $(,)?]) => {{
            use $module as runtime;
            vec![$((stringify!($name), runtime::$name as *const () as usize)),*]
        }};
    }
    let mut symbols = entries!(fun_rt => [
        rt_init,
        rt_alloc,
        rt_alloc_atomic,
        rt_apply,
        rt_apply1,
        rt_apply2,
        rt_apply3,
        rt_apply4,
        rt_closure_set,
        rt_make_pap,
        rt_div_zero,
        rt_error,
        rt_exit,
        rt_pattern_fail,
    ]);
    symbols.extend(entries!(fun_rt::abi => [
        rt_prim_intAdd,
        rt_prim_intSub,
        rt_prim_intMul,
        rt_prim_intNeg,
        rt_prim_intQuot,
        rt_prim_intRem,
        rt_prim_intDiv,
        rt_prim_intMod,
        rt_prim_intEq,
        rt_prim_intLt,
        rt_prim_intLe,
        rt_prim_dblAdd,
        rt_prim_dblSub,
        rt_prim_dblMul,
        rt_prim_dblDiv,
        rt_prim_dblNeg,
        rt_prim_dblSqrt,
        rt_prim_dblExp,
        rt_prim_dblLog,
        rt_prim_dblSin,
        rt_prim_dblCos,
        rt_prim_dblPow,
        rt_prim_dblEq,
        rt_prim_dblLt,
        rt_prim_dblLe,
        rt_prim_intToDbl,
        rt_prim_dblFloor,
        rt_prim_dblRound,
        rt_prim_dblTruncate,
        rt_prim_chrOrd,
        rt_prim_chrChr,
        rt_prim_chrEq,
        rt_prim_chrLt,
        rt_prim_strAppend,
        rt_prim_strLength,
        rt_prim_strEq,
        rt_prim_strCompare,
        rt_prim_strUnpack,
        rt_prim_strPack,
        rt_prim_showInt,
        rt_prim_showDouble,
        rt_prim_showChar,
        rt_prim_showString,
        rt_prim_readInt,
        rt_prim_error,
        rt_prim_putStr,
        rt_prim_getLine,
        rt_prim_exit,
    ]));
    symbols
}
/// Resolve declarations before MCJIT finalizes the module. Intrinsics belong to LLVM.
fn map_symbols(
    module: &inkwell::module::Module<'_>,
    engine: &ExecutionEngine<'_>,
    imports: &[(String, usize)],
) -> Result<(), String> {
    let symbols = runtime_symbols();
    for function in module
        .get_functions()
        .filter(|f| f.count_basic_blocks() == 0)
    {
        let name =
            function.get_name().to_str().map_err(|e| e.to_string())?;
        if name.starts_with("llvm.") {
            continue;
        }
        let address = symbols
            .iter()
            .find(|(symbol, _)| *symbol == name)
            .map(|(_, address)| *address)
            .or_else(|| {
                imports
                    .iter()
                    .find(|(symbol, _)| symbol == name)
                    .map(|(_, address)| *address)
            })
            .ok_or_else(|| {
                format!("unmapped JIT runtime symbol: {name}")
            })?;
        engine.add_global_mapping(&function, address);
    }
    Ok(())
}
/// Enter a finalized module without exposing a callable to the rest of the compiler.
fn call_entry(engine: &ExecutionEngine<'_>) -> Result<(), String> {
    // SAFETY: the verified backend emits exactly void fun_entry(void), C ABI.
    // All imports are explicitly mapped; code, constants, and root slots remain live.
    // This engine is never modified after lookup, and no pointer escapes this call.
    unsafe {
        let entry = engine
            .get_function::<unsafe extern "C" fn()>("fun_entry")
            .map_err(|e| e.to_string())?;
        entry.call();
    }
    Ok(())
}
}

Commit definitions after success

Repl::line dispatches commands in a loop. Declarations are added to a candidate source list, checked with a private $repl entry, and initialized before the list replaces the accepted one. The private name cannot be written in Fun and does not collide with a loaded main. A load similarly constructs a replacement session before swapping it in.

Expression execution uses ordinary language operations:

#![allow(unused)]
fn main() {
use crate::{
    ast, core,
    jit::Jit,
    session::{Options, Session},
    typecheck::TypedProgram,
};
use inkwell::context::Context;
use rustyline::{DefaultEditor, error::ReadlineError};
use std::io::{self, IsTerminal, Write};
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
/// The JIT is the interactive default; interpreters remain useful for comparison.
pub enum Engine {
    Jit,
    Core,
    Anf,
    Flat,
}
#[derive(Clone)]
struct Input {
    name: String,
    source: String,
}
/// Source definitions and executable definitions commit together after a successful input.
pub struct Repl<'ctx> {
    inputs: Vec<Input>,
    jit: Jit<'ctx>,
    context: &'ctx Context,
    engine: Engine,
    lint: bool,
    fuel: usize,
}
impl<'ctx> Repl<'ctx> {
    pub fn new(context: &'ctx Context, lint: bool, fuel: usize) -> Self {
        Self {
            inputs: vec![],
            jit: Jit::new(context),
            context,
            engine: Engine::Jit,
            lint,
            fuel,
        }
    }
    fn reset(&mut self) {
        self.jit = Jit::new(self.context);
        self.inputs.clear();
    }
    /// Process one input. False requests a clean exit; diagnostics leave the session usable.
    pub fn line(&mut self, line: &str) -> Result<bool, String> {
        let line = line.trim();
        if line.is_empty() {
            return Ok(true);
        }
        let (command, rest) =
            line.split_once(char::is_whitespace).unwrap_or((line, ""));
        let rest = rest.trim();
        match command {
            ":q" | ":quit" => return Ok(false),
            ":reset" => {
                self.reset();
                println!("Reset.");
            }
            ":engine" => {
                if rest.is_empty() {
                    println!(
                        "{}",
                        format!("{:?}", self.engine).to_lowercase()
                    );
                } else {
                    let engine = match rest {
                        "jit" => Engine::Jit,
                        "core" => Engine::Core,
                        "anf" => Engine::Anf,
                        "flat" => Engine::Flat,
                        _ => {
                            return Err(
                                "expected :engine jit|core|anf|flat".into(),
                            );
                        }
                    };
                    // Switching engines rebuilds pure globals; source definitions survive.
                    self.jit = Jit::new(self.context);
                    self.engine = engine;
                    println!("Engine: {rest}");
                }
            }
            ":load" => {
                if rest.is_empty() {
                    return Err("expected :load FILE".into());
                }
                let input = Input {
                    name: rest.into(),
                    source: std::fs::read_to_string(rest)
                        .map_err(|e| e.to_string())?,
                };
                // Keep the old session until both checking and JIT preparation succeed.
                let mut next =
                    Self::new(self.context, self.lint, self.fuel);
                next.engine = self.engine;
                next.evaluate(std::slice::from_ref(&input), "pure ()")?;
                next.inputs.push(input);
                *self = next;
                println!("Loaded {rest}");
            }
            ":t" | ":type" => {
                let (_, typed) = prepare(&self.inputs, rest, false)?;
                println!("{rest} : {}", typed.main_description());
            }
            ":k" => {
                let (mut sess, typed) = prepare(&self.inputs, "()", false)?;
                let file = sess
                    .sources
                    .add("<kind>", format!("type ReplKind = {rest}"));
                let syntax = crate::parse::parse(&mut sess, file)
                    .map_err(|_| diagnostics(&sess))?;
                let alias =
                    syntax.aliases.first().ok_or("expected :k TYPE")?;
                let mut store = crate::kinds::Store::default();
                let expanded = crate::kinds::expand(
                    &alias.body,
                    &typed.surface.data.aliases,
                    alias.span,
                )
                .map_err(|e| e.message)?;
                let kind = store
                    .infer(
                        &expanded,
                        &typed.surface.data.kinds,
                        &mut std::collections::BTreeMap::new(),
                        false,
                        alias.span,
                    )
                    .map_err(|e| e.message)?;
                println!("{rest} : {}", store.finish(&kind));
            }
            ":core" | ":anf" | ":llvm" => self.dump(rest, command)?,
            ":help" => println!(
                "EXPR; declarations; :t EXPR; :k TYPE; :core/:anf/:llvm EXPR; :load FILE; :reset; :engine jit|core|anf|flat; :quit"
            ),
            _ if command.starts_with(':') => {
                return Err(format!(
                    "unknown command: {command}; use :help"
                ));
            }
            _ => {
                // Trial parsing distinguishes declarations from expressions without publishing diagnostics.
                let mut trial = Session::default();
                let file = trial.sources.add("<input>", line.into());
                if crate::parse::parse(&mut trial, file).is_ok() {
                    let mut candidate = self.inputs.clone();
                    candidate.push(Input {
                        name: format!("<input {}>", candidate.len() + 1),
                        source: line.into(),
                    });
                    self.evaluate(&candidate, "pure ()")?;
                    self.inputs = candidate;
                } else {
                    let inputs = self.inputs.clone();
                    let (_, typed) = prepare(&inputs, line, false)?;
                    let query = action_query(line, &typed.main_type);
                    self.evaluate(&inputs, &query)?;
                }
            }
        }
        Ok(true)
    }
    fn evaluate(
        &mut self,
        inputs: &[Input],
        query: &str,
    ) -> Result<(), String> {
        let (mut sess, typed) = prepare(inputs, query, true)?;
        let core = lower_core(&mut sess, typed, self.lint)?;
        if self.engine == Engine::Core {
            return crate::eval::run_io(&core, self.fuel)
                .map_err(|e| e.to_string());
        }
        let anf = lower_anf(&mut sess, core, self.lint)?;
        if self.engine == Engine::Anf {
            return crate::anf_machine::run(&anf, self.fuel, true)
                .map(|_| ())
                .map_err(|e| e.to_string());
        }
        let flat = crate::closure::convert(&mut sess, anf);
        crate::flat_lint::lint(&flat)?;
        if self.engine == Engine::Jit {
            match self.jit.run(&flat) {
                Ok(()) => return Ok(()),
                Err(e) if e.starts_with("JIT setup failed:") => {
                    eprintln!("fun: {e}; using flat interpreter");
                    self.engine = Engine::Flat;
                }
                Err(e) => return Err(e),
            }
        }
        crate::flat_machine::run(&flat, self.fuel, true)
            .map(|_| ())
            .map_err(|e| e.to_string())
    }
    fn dump(&self, expression: &str, pass: &str) -> Result<(), String> {
        let (mut sess, mut typed) =
            prepare(&self.inputs, expression, false)?;
        if pass == ":llvm" {
            let query = action_query(expression, &typed.main_type);
            (sess, typed) = prepare(&self.inputs, &query, true)?;
        }
        let core = lower_core(&mut sess, typed, self.lint)?;
        if pass == ":core" {
            for group in &core.groups {
                println!("{group}");
            }
            return Ok(());
        }
        let anf = lower_anf(&mut sess, core, self.lint)?;
        if pass == ":anf" {
            for (name, body) in &anf.globals {
                println!("{name} = {body}");
            }
            return Ok(());
        }
        let flat = crate::closure::convert(&mut sess, anf);
        let backend = crate::llvm::Backend::for_jit(
            self.context,
            &flat,
            Default::default(),
        )
        .map_err(|e| e.to_string())?;
        print!("{}", backend.ir());
        Ok(())
    }
}
fn diagnostics(sess: &Session) -> String {
    let mut bytes = vec![];
    for diagnostic in &sess.diagnostics {
        let _ = diagnostic.render(&sess.sources, false, &mut bytes);
    }
    String::from_utf8_lossy(&bytes).trim_end().into()
}
fn append(to: &mut ast::Program, mut from: ast::Program) {
    to.bindings.append(&mut from.bindings);
    to.fixities.append(&mut from.fixities);
    to.data.append(&mut from.data);
    to.aliases.append(&mut from.aliases);
    to.classes.append(&mut from.classes);
    to.instances.append(&mut from.instances);
}
fn prepare(
    inputs: &[Input],
    query: &str,
    io: bool,
) -> Result<(Session, TypedProgram), String> {
    let mut sess = Session {
        options: Options {
            require_io: io,
            ..Options::default()
        },
        entry_name: Some("$repl".into()),
        ..Session::default()
    };
    let mut program = None;
    let mut later_files = vec![];
    for (index, input) in inputs.iter().enumerate() {
        let file =
            sess.sources.add(input.name.clone(), input.source.clone());
        if index > 0 {
            later_files.push(file);
        }
        let parsed = crate::parse::parse(&mut sess, file)
            .map_err(|_| diagnostics(&sess))?;
        if index > 0
            && parsed.fixities.iter().any(|f| {
                f.operators.iter().any(|op| {
                    !parsed.bindings.iter().any(|b| b.name == *op)
                })
            })
        {
            return Err("a REPL fixity must accompany its operator definition in the same input".into());
        }
        if let Some(program) = &mut program {
            append(program, parsed);
        } else {
            program = Some(parsed);
        }
    }
    let file = sess.sources.add("<query>", format!("main = {query}"));
    let mut entry = crate::parse::parse(&mut sess, file)
        .map_err(|_| diagnostics(&sess))?;
    entry.bindings[0].name = "$repl".into();
    let program = if let Some(mut program) = program {
        append(&mut program, entry);
        program
    } else {
        entry
    };
    let typed = crate::typecheck::check(&mut sess, program)
        .map_err(|_| diagnostics(&sess))?;
    // Retained code and replaying interpreters must agree on name resolution.
    // Put Prelude overrides in the initial declaration block or a loaded file.
    let prelude_names = typed
        .surface
        .groups
        .iter()
        .flat_map(|g| &g.binds)
        .filter(|b| Some(b.span.file) == sess.prelude_file)
        .map(|b| b.id.text.as_str())
        .collect::<std::collections::BTreeSet<_>>();
    for binding in typed.surface.groups.iter().flat_map(|g| &g.binds) {
        if later_files.contains(&binding.span.file)
            && prelude_names.contains(binding.id.text.as_str())
        {
            return Err(format!(
                "{} shadows the Prelude after earlier definitions; use :reset or :load",
                binding.id.text
            ));
        }
    }
    Ok((sess, typed))
}
fn lower_core(
    sess: &mut Session,
    typed: TypedProgram,
    lint: bool,
) -> Result<core::Program, String> {
    let core = crate::elaborate::elaborate(sess, typed)
        .map_err(|_| diagnostics(sess))?;
    if lint {
        core::lint(&core)?;
    }
    let core = crate::match_compile::compile(sess, core)
        .map_err(|_| diagnostics(sess))?;
    if lint {
        core::lint(&core)?;
    }
    Ok(core)
}
fn lower_anf(
    sess: &mut Session,
    core: core::Program,
    lint: bool,
) -> Result<crate::anf::Program, String> {
    let mut anf = crate::anf::convert(sess, crate::erase::erase(core));
    crate::anf_simplify::simplify(sess, &mut anf);
    if lint {
        crate::anf_lint::lint(&anf)?;
    }
    Ok(anf)
}
/// Terminal input uses rustyline; redirected input has deterministic, prompt-free output.
pub fn run(
    lint: bool,
    fuel: usize,
) -> Result<(), Box<dyn std::error::Error>> {
    let context = Context::create();
    let mut repl = Repl::new(&context, lint, fuel);
    let interactive = io::stdin().is_terminal();
    let mut editor = if interactive {
        Some(DefaultEditor::new()?)
    } else {
        None
    };
    if interactive {
        println!("Fun: JIT ready on first input. Type :help for commands.");
    }
    loop {
        let line = if let Some(editor) = &mut editor {
            match editor.readline("fun> ") {
                Ok(line) => {
                    editor.add_history_entry(&line)?;
                    line
                }
                Err(ReadlineError::Eof | ReadlineError::Interrupted) => {
                    break;
                }
                Err(e) => return Err(e.into()),
            }
        } else {
            let mut line = String::new();
            if io::stdin().read_line(&mut line)? == 0 {
                break;
            }
            line
        };
        match repl.line(&line) {
            Ok(false) => break,
            Ok(true) => {}
            Err(e) => eprintln!("fun: {e}"),
        }
        io::stdout().flush()?;
    }
    Ok(())
}
/// Both display and IO sequencing are ordinary, type-checked Prelude operations.
fn action_query(expression: &str, ty: &crate::types::Type) -> String {
    if matches!(ty, crate::types::Type::Con(id, _) if id.name == "IO") {
        format!("({expression}) >> pure ()")
    } else {
        format!("print ({expression})")
    }
}
}

The resulting query passes through inference and elaboration again. The wrapper does not bypass class solving or IO abstraction. Type, kind, and IR commands stop before execution. All new traversal and command control uses loops. No larger Rust stack or stack-growth library is introduced. Recursive frontend passes inherited from earlier chapters retain their existing limits.

We recheck accepted source and lower a candidate module on each input. Existing global values and fast/slow function entries are imported, preserving both initialization and code identity. Unused regenerated local bodies are removed by optimization, but compilation work and retained machine code grow with the session. This makes ownership easy to inspect. Incremental checking and code reclamation would require additional dependency and lifetime machinery.

Note. Runtime errors and exitWith retain native process-exit semantics, including in the JIT REPL. We never unwind or jump across generated frames. Compiler errors are recoverable. If creation of the execution engine fails, the REPL selects Flat. Verifier failures are reported rather than hidden by fallback. Explicit run --engine=jit reports setup failure.

Running It

Build with the LLVM 22 and Boehm setup from Chapter 21:

$ python3 scripts/build_runtime.py 22
$ cargo build -p ch22_jit
$ target/debug/fun repl --lint
Fun: JIT ready on first input. Type :help for commands.
fun> make x = \y -> x + y
fun> saved = make 40
fun> saved 2
42
fun> :t saved
saved : Int -> Int
fun> :k Maybe
Maybe : * -> *
fun> :reset
Reset.
fun> :quit
$ target/debug/fun run crates/ch22_jit/tests/programs/native/acceptance.fun --engine=jit --lint
5.0
3.0
6.0
[3.0,10.0]
True
50000005000000
$ target/debug/fun dump --pass=llvm crates/ch22_jit/tests/programs/jit/returned-function.fun > build/jit-example.ll

fun and fun-ch22 share one driver. The CLI keeps Flat as its default run engine. The REPL defaults to JIT. --fuel and the test-only --print-value remain interpreter facilities. Redirected REPL input omits prompts and the banner so transcripts can be compared byte for byte.

The JIT should give the same answers as the interpreters and standalone executables. It also needs to survive long known-call recursion and allocation under a bounded heap. Known calls keep Chapter 21’s tail-call guarantee. Calls through C apply still have the limitation we discussed there.

Try defining a closure, entering more definitions, and then calling the old closure again. The lifetime tests push this further by forcing collections between modules. A value remaining in scope isn’t enough on its own. Both the captured data and the machine code it points to must still be there.

Exercises

Warm-up. Trace the owners of saved = make 40 across two inputs. Identify which owner protects its environment and which protects its code address.
Intermediate. Remove one runtime mapping and inspect the failure. Explain why it must be detected before calling the generated entry.
Intermediate. Add a regression where a user definition shadows a Prelude name. Show why spelling alone cannot identify its persistent slot.
Advanced. Design an isolated worker protocol that recovers from runtime errors while retaining accepted definitions. Specify what must be replayed.
Advanced. Propose a module-reclamation scheme. Account for closures inside constructors, partial applications, global values, and captured environments.

Write You a Haskell