lust/jit/
mod.rs

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

lust/jit/
mod.rs
