lust/jit/
mod.rs

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

lust/jit/
mod.rs
