lust/jit/
mod.rs

1#[cfg(feature = "std")]
2pub mod codegen;
3pub mod optimizer;
4pub mod profiler;
5pub mod specialization;
6pub mod trace;
7use crate::bytecode::Value;
8use crate::VM;
9#[cfg(feature = "std")]
10pub use codegen::JitCompiler;
11#[cfg(not(feature = "std"))]
12pub struct JitCompiler;
13use alloc::{string::String, vec::Vec};
14use hashbrown::HashMap;
15pub use optimizer::TraceOptimizer;
16pub use profiler::{HotSpot, Profiler};
17pub use trace::{Trace, TraceOp, TraceRecorder};
18#[cfg(not(feature = "std"))]
19impl JitCompiler {
20    pub fn new() -> Self {
21        Self
22    }
23
24    pub fn compile_trace(
25        &mut self,
26        _trace: &Trace,
27        _trace_id: TraceId,
28        _parent: Option<TraceId>,
29        _hoisted_constants: Vec<(u8, Value)>,
30    ) -> crate::Result<CompiledTrace> {
31        Err(crate::LustError::RuntimeError {
32            message: "JIT is unavailable without the `std` feature".into(),
33        })
34    }
35}
36#[cfg(all(debug_assertions, feature = "std"))]
37#[inline]
38pub(crate) fn log<F>(message: F)
39where
40    F: FnOnce() -> String,
41{
42    println!("{}", message());
43}
44
45#[cfg(not(all(debug_assertions, feature = "std")))]
46#[inline]
47pub(crate) fn log<F>(_message: F)
48where
49    F: FnOnce() -> String,
50{
51}
52
53pub const HOT_THRESHOLD: u32 = 5;
54pub const MAX_TRACE_LENGTH: usize = 2000; // Increased to allow more loop unrolling
55pub const SIDE_EXIT_THRESHOLD: u32 = 10;
56pub const UNROLL_FACTOR: usize = 32;
57/// How many times to unroll a loop during trace recording
58pub const LOOP_UNROLL_COUNT: usize = 32;
59#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
60pub struct TraceId(pub usize);
61pub struct CompiledTrace {
62    pub id: TraceId,
63    pub entry: extern "C" fn(*mut Value, *mut VM, *const crate::bytecode::Function) -> i32,
64    pub trace: Trace,
65    pub guards: Vec<Guard>,
66    pub parent: Option<TraceId>,
67    pub side_traces: Vec<TraceId>,
68    pub leaked_constants: Vec<*const Value>,
69    pub hoisted_constants: Vec<(u8, Value)>,
70}
71
72#[derive(Debug, Clone)]
73pub struct Guard {
74    pub index: usize,
75    pub bailout_ip: usize,
76    pub kind: GuardKind,
77    pub fail_count: u32,
78    pub side_trace: Option<TraceId>,
79}
80
81#[derive(Debug, Clone)]
82pub enum GuardKind {
83    IntType {
84        register: u8,
85    },
86    FloatType {
87        register: u8,
88    },
89    BoolType {
90        register: u8,
91    },
92    Truthy {
93        register: u8,
94    },
95    Falsy {
96        register: u8,
97    },
98    ArrayBoundsCheck {
99        array_register: u8,
100        index_register: u8,
101    },
102    NestedLoop {
103        function_idx: usize,
104        loop_start_ip: usize,
105    },
106    NativeFunction {
107        register: u8,
108        expected: *const (),
109    },
110    Function {
111        register: u8,
112        function_idx: usize,
113    },
114    Closure {
115        register: u8,
116        function_idx: usize,
117        upvalues_ptr: *const (),
118    },
119}
120
121pub struct JitState {
122    pub profiler: Profiler,
123    pub traces: HashMap<TraceId, CompiledTrace>,
124    pub root_traces: HashMap<(usize, usize), TraceId>,
125    next_trace_id: usize,
126    pub enabled: bool,
127}
128
129impl JitState {
130    pub fn new() -> Self {
131        let enabled = cfg!(all(feature = "std", target_arch = "x86_64"));
132        Self {
133            profiler: Profiler::new(),
134            traces: HashMap::new(),
135            root_traces: HashMap::new(),
136            next_trace_id: 0,
137            enabled,
138        }
139    }
140
141    pub fn alloc_trace_id(&mut self) -> TraceId {
142        let id = TraceId(self.next_trace_id);
143        self.next_trace_id += 1;
144        id
145    }
146
147    pub fn check_hot(&mut self, func_idx: usize, ip: usize) -> bool {
148        if !self.enabled {
149            return false;
150        }
151
152        self.profiler.record_backedge(func_idx, ip) >= HOT_THRESHOLD
153    }
154
155    pub fn get_root_trace(&self, func_idx: usize, ip: usize) -> Option<&CompiledTrace> {
156        self.root_traces
157            .get(&(func_idx, ip))
158            .and_then(|id| self.traces.get(id))
159    }
160
161    pub fn get_trace(&self, id: TraceId) -> Option<&CompiledTrace> {
162        self.traces.get(&id)
163    }
164
165    pub fn get_trace_mut(&mut self, id: TraceId) -> Option<&mut CompiledTrace> {
166        self.traces.get_mut(&id)
167    }
168
169    pub fn store_root_trace(&mut self, func_idx: usize, ip: usize, trace: CompiledTrace) {
170        let id = trace.id;
171        self.root_traces.insert((func_idx, ip), id);
172        self.traces.insert(id, trace);
173    }
174
175    pub fn store_side_trace(&mut self, trace: CompiledTrace) {
176        let id = trace.id;
177        self.traces.insert(id, trace);
178    }
179}
180
181impl Default for JitState {
182    fn default() -> Self {
183        Self::new()
184    }
185}
lust/jit/mod.rs

lust/jit/
mod.rs
