Skip to main content

laminar_core/compiler/
jit.rs

1//! Cranelift JIT compilation context.
2//!
3//! [`JitContext`] owns the Cranelift [`JITModule`] and provides shared
4//! compilation infrastructure for [`ExprCompiler`](super::expr::ExprCompiler).
5
6use std::sync::Arc;
7
8use cranelift_codegen::isa::TargetIsa;
9use cranelift_codegen::settings::{self, Configurable};
10use cranelift_frontend::FunctionBuilderContext;
11use cranelift_jit::{JITBuilder, JITModule};
12use cranelift_module::Module;
13
14use super::error::CompileError;
15
16/// Owns the Cranelift JIT module and builder context used for compiling expressions.
17///
18/// Create one per compilation session (typically per query). Each compiled function
19/// gets a unique name via [`next_func_name`](Self::next_func_name).
20pub struct JitContext {
21    module: JITModule,
22    builder_ctx: FunctionBuilderContext,
23    func_counter: u32,
24}
25
26impl JitContext {
27    /// Creates a new JIT context targeting the host CPU with `opt_level = speed`.
28    ///
29    /// # Errors
30    ///
31    /// Returns [`CompileError::Cranelift`] if the native ISA cannot be detected.
32    pub fn new() -> Result<Self, CompileError> {
33        let isa = Self::create_isa()?;
34        let builder = JITBuilder::with_isa(isa, cranelift_module::default_libcall_names());
35        let module = JITModule::new(builder);
36
37        Ok(Self {
38            module,
39            builder_ctx: FunctionBuilderContext::new(),
40            func_counter: 0,
41        })
42    }
43
44    /// Returns a mutable reference to the underlying [`JITModule`].
45    pub fn module(&mut self) -> &mut JITModule {
46        &mut self.module
47    }
48
49    /// Returns a mutable reference to the shared [`FunctionBuilderContext`].
50    pub fn builder_ctx(&mut self) -> &mut FunctionBuilderContext {
51        &mut self.builder_ctx
52    }
53
54    /// Returns the target ISA used by this context.
55    pub fn isa(&self) -> &dyn TargetIsa {
56        self.module.isa()
57    }
58
59    /// Generates a unique function name with the given prefix.
60    pub fn next_func_name(&mut self, prefix: &str) -> String {
61        let id = self.func_counter;
62        self.func_counter += 1;
63        format!("{prefix}_{id}")
64    }
65
66    /// Creates a host-native ISA with speed optimization.
67    fn create_isa() -> Result<Arc<dyn TargetIsa>, CompileError> {
68        let mut flag_builder = settings::builder();
69        flag_builder
70            .set("opt_level", "speed")
71            .expect("valid opt_level setting");
72        let isa_builder = cranelift_native::builder()
73            .map_err(|e| CompileError::UnsupportedExpr(format!("native ISA: {e}")))?;
74        let flags = settings::Flags::new(flag_builder);
75        Ok(isa_builder.finish(flags).expect("valid ISA flags"))
76    }
77}
78
79#[cfg(test)]
80mod tests {
81    use super::*;
82
83    #[test]
84    fn jit_context_creation() {
85        let ctx = JitContext::new().unwrap();
86        let isa = ctx.isa();
87        // Should detect some valid target triple.
88        let triple = isa.triple().to_string();
89        assert!(!triple.is_empty());
90    }
91
92    #[test]
93    fn unique_func_names() {
94        let mut ctx = JitContext::new().unwrap();
95        let n1 = ctx.next_func_name("filter");
96        let n2 = ctx.next_func_name("filter");
97        let n3 = ctx.next_func_name("scalar");
98        assert_eq!(n1, "filter_0");
99        assert_eq!(n2, "filter_1");
100        assert_eq!(n3, "scalar_2");
101    }
102
103    #[test]
104    fn module_accessible() {
105        let mut ctx = JitContext::new().unwrap();
106        // Just verify we can access the module without panicking.
107        let _m = ctx.module();
108    }
109
110    #[test]
111    fn builder_ctx_accessible() {
112        let mut ctx = JitContext::new().unwrap();
113        let _b = ctx.builder_ctx();
114    }
115}