tidepool_codegen/
jit_machine.rs

1use tidepool_effect::{DispatchEffect, EffectContext, EffectError};
2use tidepool_eval::value::Value;
3use tidepool_repr::{CoreExpr, DataConTable};
4use cranelift_module::FuncId;
5
6use crate::context::VMContext;
7use crate::effect_machine::{CompiledEffectMachine, ConTags};
8use crate::heap_bridge;
9use crate::nursery::Nursery;
10use crate::pipeline::CodegenPipeline;
11use crate::yield_type::Yield;
12
13/// Error type for JIT compilation/execution failures.
14#[derive(Debug)]
15pub enum JitError {
16    Compilation(String),
17    MissingConTags,
18    Effect(EffectError),
19    Yield(crate::yield_type::YieldError),
20    HeapBridge(String),
21}
22
23impl std::fmt::Display for JitError {
24    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
25        match self {
26            JitError::Compilation(s) => write!(f, "JIT compilation error: {}", s),
27            JitError::MissingConTags => write!(f, "missing freer-simple constructors in DataConTable"),
28            JitError::Effect(e) => write!(f, "effect dispatch error: {}", e),
29            JitError::Yield(e) => write!(f, "yield error: {}", e),
30            JitError::HeapBridge(s) => write!(f, "heap bridge error: {}", s),
31        }
32    }
33}
34
35impl std::error::Error for JitError {}
36
37impl From<EffectError> for JitError {
38    fn from(e: EffectError) -> Self {
39        JitError::Effect(e)
40    }
41}
42
43/// High-level JIT effect machine — compile and run freer-simple effect stacks
44/// without touching raw pointers, nurseries, or stack maps.
45pub struct JitEffectMachine {
46    pipeline: CodegenPipeline,
47    nursery: Nursery,
48    tags: ConTags,
49    func_id: FuncId,
50}
51
52impl JitEffectMachine {
53    /// Compile a CoreExpr for JIT execution.
54    pub fn compile(
55        expr: &CoreExpr,
56        table: &DataConTable,
57        nursery_size: usize,
58    ) -> Result<Self, JitError> {
59        let expr = crate::datacon_env::wrap_with_datacon_env(expr, table);
60        let mut pipeline = CodegenPipeline::new(&crate::host_fns::host_fn_symbols());
61        let func_id = crate::emit::expr::compile_expr(&mut pipeline, &expr, "main")
62            .map_err(|e| JitError::Compilation(format!("{:?}", e)))?;
63        pipeline.finalize();
64
65        let tags = ConTags::from_table(table).ok_or(JitError::MissingConTags)?;
66        let nursery = Nursery::new(nursery_size);
67
68        Ok(Self {
69            pipeline,
70            nursery,
71            tags,
72            func_id,
73        })
74    }
75
76    /// Run to completion, dispatching effects through the handler HList.
77    pub fn run<U, H: DispatchEffect<U>>(
78        &mut self,
79        table: &DataConTable,
80        handlers: &mut H,
81        user: &U,
82    ) -> Result<Value, JitError> {
83        // Install registries
84        crate::debug::set_lambda_registry(self.pipeline.build_lambda_registry());
85        crate::host_fns::set_stack_map_registry(&self.pipeline.stack_maps);
86
87        let func_ptr: unsafe extern "C" fn(*mut VMContext) -> *mut u8 =
88            unsafe { std::mem::transmute(self.pipeline.get_function_ptr(self.func_id)) };
89        let vmctx = self.nursery.make_vmctx(crate::host_fns::gc_trigger);
90
91        let mut machine = CompiledEffectMachine::new(func_ptr, vmctx, self.tags);
92        let mut yield_result = machine.step();
93
94        let result = loop {
95            match yield_result {
96                Yield::Done(ptr) => {
97                    let val = unsafe { heap_bridge::heap_to_value(ptr) }
98                        .map_err(|e| JitError::HeapBridge(format!("{:?}", e)))?;
99                    break Ok(val);
100                }
101                Yield::Request {
102                    tag,
103                    request,
104                    continuation,
105                } => {
106                    let req_val = unsafe { heap_bridge::heap_to_value(request) }
107                        .map_err(|e| JitError::HeapBridge(format!("{:?}", e)))?;
108                    let cx = EffectContext::with_user(table, user);
109                    let resp_val = handlers.dispatch(tag, &req_val, &cx)?;
110                    let resp_ptr =
111                        unsafe { heap_bridge::value_to_heap(&resp_val, machine.vmctx_mut()) }
112                            .map_err(|e| JitError::HeapBridge(format!("{:?}", e)))?;
113                    yield_result = unsafe { machine.resume(continuation, resp_ptr) };
114                }
115                Yield::Error(e) => break Err(JitError::Yield(e)),
116            }
117        };
118
119        // Cleanup registries
120        crate::host_fns::clear_stack_map_registry();
121        crate::debug::clear_lambda_registry();
122
123        result
124    }
125}
tidepool_codegen/jit_machine.rs

tidepool_codegen/
jit_machine.rs
