tl_compiler/
jit.rs

1// ThinkingLanguage — Cranelift JIT Compiler
2// Compiles hot bytecode prototypes into native code.
3//
4// Strategy: baseline JIT — all type-dependent operations call runtime helpers.
5// This avoids speculative type guards while still eliminating bytecode dispatch overhead.
6
7use std::collections::HashMap;
8
9use cranelift_codegen::settings::{self, Configurable};
10use cranelift_frontend::FunctionBuilderContext;
11use cranelift_jit::{JITBuilder, JITModule};
12use cranelift_module::Module;
13
14use crate::chunk::Prototype;
15use crate::jit_runtime;
16use crate::value::VmValue;
17
18/// JIT-compiled function signature:
19/// extern "C" fn(args: *const VmValue, nargs: usize) -> VmValue
20pub type JitFn = unsafe extern "C" fn(*const VmValue, usize) -> VmValue;
21
22/// The Cranelift JIT compiler for TL functions.
23pub struct JitCompiler {
24    module: JITModule,
25    ctx: cranelift_codegen::Context,
26    builder_ctx: FunctionBuilderContext,
27    /// Cache of compiled function pointers, keyed by prototype name
28    compiled: HashMap<String, *const u8>,
29}
30
31impl JitCompiler {
32    pub fn new() -> Result<Self, String> {
33        let mut flag_builder = settings::builder();
34        flag_builder
35            .set("use_colocated_libcalls", "false")
36            .map_err(|e| format!("JIT settings error: {e}"))?;
37        flag_builder
38            .set("is_pic", "false")
39            .map_err(|e| format!("JIT settings error: {e}"))?;
40
41        let isa_builder = cranelift_codegen::isa::lookup(target_lexicon::Triple::host())
42            .map_err(|e| format!("JIT ISA error: {e}"))?;
43
44        let isa = isa_builder
45            .finish(settings::Flags::new(flag_builder))
46            .map_err(|e| format!("JIT ISA finish error: {e}"))?;
47
48        let mut builder = JITBuilder::with_isa(isa, cranelift_module::default_libcall_names());
49
50        // Register runtime helper symbols
51        builder.symbol("tl_rt_add", jit_runtime::tl_rt_add as *const u8);
52        builder.symbol("tl_rt_sub", jit_runtime::tl_rt_sub as *const u8);
53        builder.symbol("tl_rt_mul", jit_runtime::tl_rt_mul as *const u8);
54        builder.symbol("tl_rt_cmp", jit_runtime::tl_rt_cmp as *const u8);
55        builder.symbol("tl_rt_is_truthy", jit_runtime::tl_rt_is_truthy as *const u8);
56
57        let module = JITModule::new(builder);
58        let ctx = module.make_context();
59
60        Ok(JitCompiler {
61            module,
62            ctx,
63            builder_ctx: FunctionBuilderContext::new(),
64            compiled: HashMap::new(),
65        })
66    }
67
68    /// Check if a function has already been JIT-compiled.
69    pub fn get_compiled(&self, name: &str) -> Option<*const u8> {
70        self.compiled.get(name).copied()
71    }
72
73    /// Compile a prototype to native code and return the function pointer.
74    /// Returns None if the function is too complex to JIT (e.g. uses table ops).
75    pub fn compile_function(&mut self, proto: &Prototype) -> Result<Option<*const u8>, String> {
76        // For now, only JIT simple numeric functions (no table ops, no closures)
77        if !proto.upvalue_defs.is_empty() {
78            return Ok(None); // Skip closures for now
79        }
80
81        // Check for table ops or complex instructions we can't JIT
82        for &inst in &proto.code {
83            let op = crate::opcode::decode_op(inst);
84            match op {
85                crate::opcode::Op::TablePipe
86                | crate::opcode::Op::Interpolate
87                | crate::opcode::Op::NewMap
88                | crate::opcode::Op::GetMember => {
89                    return Ok(None); // Too complex for baseline JIT
90                }
91                _ => {}
92            }
93        }
94
95        // For now, return None — full Cranelift IR generation is complex
96        // The JIT infrastructure is set up and ready for incremental improvement
97        Ok(None)
98    }
99}
100
101/// JIT call count threshold — functions called more than this many times get JIT-compiled
102pub const JIT_THRESHOLD: u32 = 100;
103
104/// Tiered compilation state, tracked per prototype.
105#[derive(Debug, Default)]
106pub struct TieringState {
107    /// How many times each function has been called
108    pub call_counts: HashMap<String, u32>,
109}
110
111impl TieringState {
112    pub fn new() -> Self {
113        Self::default()
114    }
115
116    /// Record a call and return true if the function should be JIT-compiled.
117    pub fn record_call(&mut self, name: &str) -> bool {
118        let count = self.call_counts.entry(name.to_string()).or_insert(0);
119        *count += 1;
120        *count == JIT_THRESHOLD
121    }
122}
123
124#[cfg(test)]
125mod tests {
126    use super::*;
127
128    #[test]
129    fn test_jit_compiler_creation() {
130        let jit = JitCompiler::new();
131        assert!(jit.is_ok(), "JIT compiler should initialize");
132    }
133
134    #[test]
135    fn test_tiering_state() {
136        let mut state = TieringState::new();
137        for _ in 0..JIT_THRESHOLD - 1 {
138            assert!(!state.record_call("test_fn"));
139        }
140        // The threshold-th call should trigger JIT
141        assert!(state.record_call("test_fn"));
142        // After threshold, no more triggers
143        assert!(!state.record_call("test_fn"));
144    }
145}
tl_compiler/jit.rs

tl_compiler/
jit.rs
