1use std::collections::{HashMap, HashSet};
2
3use crate::debug_info::DebugInfo;
4use crate::vm::{OpCode, Program};
5
6use super::ir::SsaTrace;
7use super::liveness::{boxed_load_site_count, boxed_store_site_count};
8use super::recorder::{RecordedTrace, TraceRecordError, record_trace};
9
10#[derive(Clone, Copy, Debug, PartialEq, Eq)]
11pub struct JitConfig {
12 pub enabled: bool,
13 pub hot_loop_threshold: u32,
14 pub max_trace_len: usize,
15}
16
17impl Default for JitConfig {
18 fn default() -> Self {
19 Self {
20 enabled: native_jit_supported(),
21 hot_loop_threshold: 8,
22 max_trace_len: 256,
23 }
24 }
25}
26
27fn native_jit_supported() -> bool {
28 (cfg!(target_arch = "x86_64")
29 && (cfg!(target_os = "windows") || (cfg!(unix) && !cfg!(target_os = "macos"))))
30 || (cfg!(target_arch = "aarch64")
31 && (cfg!(target_os = "linux") || cfg!(target_os = "macos")))
32}
33
34#[derive(Clone, Debug, PartialEq, Eq, Hash)]
35pub enum JitTraceTerminal {
36 LoopBack,
37 BranchExit,
38 Halt,
39}
40
41#[derive(Clone, Debug, PartialEq, Eq)]
42pub enum JitNyiReason {
43 UnsupportedArch,
44 HotLoopThresholdZero,
45 UnsupportedOpcode(u8),
46 UnsupportedTrace(String),
47 InvalidJumpTarget { target: usize },
48 InvalidImmediate(&'static str),
49 TraceTooLong { limit: usize },
50 MissingTerminal,
51}
52
53impl JitNyiReason {
54 pub fn message(&self) -> String {
55 match self {
56 JitNyiReason::UnsupportedArch => {
57 "target architecture is not x86_64-unix-non-macos/x86_64-windows/aarch64-linux/aarch64-macos".to_string()
58 }
59 JitNyiReason::HotLoopThresholdZero => "hot_loop_threshold must be > 0".to_string(),
60 JitNyiReason::UnsupportedOpcode(op) => format!("unsupported opcode 0x{op:02X}"),
61 JitNyiReason::UnsupportedTrace(detail) => detail.clone(),
62 JitNyiReason::InvalidJumpTarget { target } => {
63 format!("jump target {target} is invalid or out of bytecode bounds")
64 }
65 JitNyiReason::InvalidImmediate(kind) => {
66 format!("failed to decode immediate operand for {kind}")
67 }
68 JitNyiReason::TraceTooLong { limit } => {
69 format!("trace length exceeded configured limit {limit}")
70 }
71 JitNyiReason::MissingTerminal => {
72 "trace recorder reached end without loopback/ret terminal".to_string()
73 }
74 }
75 }
76}
77
78#[derive(Clone, Debug, PartialEq)]
79pub struct JitTrace {
80 pub id: usize,
81 pub root_ip: usize,
82 pub start_line: Option<u32>,
83 pub has_call: bool,
84 pub has_yielding_call: bool,
85 pub op_names: Vec<String>,
86 pub terminal: JitTraceTerminal,
87 pub executions: u64,
88 pub(crate) ssa: SsaTrace,
89}
90
91impl JitTrace {
92 pub fn op_names(&self) -> &[String] {
93 &self.op_names
94 }
95
96 pub fn ssa_text(&self) -> String {
97 self.ssa.render_text()
98 }
99
100 pub fn ssa_block_count(&self) -> usize {
101 self.ssa.blocks.len()
102 }
103
104 pub fn ssa_exit_count(&self) -> usize {
105 self.ssa.exits.len()
106 }
107
108 pub fn boxed_load_site_count(&self) -> u64 {
109 boxed_load_site_count(&self.ssa)
110 }
111
112 pub fn boxed_store_site_count(&self) -> u64 {
113 boxed_store_site_count(&self.ssa)
114 }
115}
116
117#[derive(Clone, Debug, PartialEq, Eq)]
118pub struct JitAttempt {
119 pub root_ip: usize,
120 pub line: Option<u32>,
121 pub result: Result<usize, JitNyiReason>,
122}
123
124#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
125pub struct JitMetrics {
126 pub boxed_load_site_count: u64,
127 pub boxed_store_site_count: u64,
128 pub trace_exit_count: u64,
129 pub native_loop_back_count: u64,
130 pub helper_fallback_count: u64,
131 pub native_trace_exec_count: u64,
132}
133
134impl JitMetrics {
135 pub fn guard_exit_count(self) -> u64 {
136 self.trace_exit_count
137 }
138}
139
140#[derive(Clone, Debug, PartialEq)]
141pub struct JitSnapshot {
142 pub arch: &'static str,
143 pub config: JitConfig,
144 pub traces: Vec<JitTrace>,
145 pub attempts: Vec<JitAttempt>,
146 pub metrics: JitMetrics,
147 pub nyi_reference: Vec<JitNyiDoc>,
148}
149
150#[derive(Clone, Debug, PartialEq, Eq)]
151pub struct JitNyiDoc {
152 pub item: &'static str,
153 pub reason: &'static str,
154}
155
156pub struct TraceJitEngine {
157 config: JitConfig,
158 hot_counts: HashMap<usize, u32>,
159 compiled_by_root: HashMap<usize, usize>,
160 blocked_roots: HashSet<usize>,
161 loop_headers: Option<HashSet<usize>>,
162 traces: Vec<JitTrace>,
163 attempts: Vec<JitAttempt>,
164}
165
166impl Default for TraceJitEngine {
167 fn default() -> Self {
168 Self::new(JitConfig::default())
169 }
170}
171
172impl TraceJitEngine {
173 pub fn new(config: JitConfig) -> Self {
174 Self {
175 config,
176 hot_counts: HashMap::new(),
177 compiled_by_root: HashMap::new(),
178 blocked_roots: HashSet::new(),
179 loop_headers: None,
180 traces: Vec::new(),
181 attempts: Vec::new(),
182 }
183 }
184
185 pub fn config(&self) -> &JitConfig {
186 &self.config
187 }
188
189 pub fn set_config(&mut self, config: JitConfig) {
190 self.config = config;
191 self.hot_counts.clear();
192 self.compiled_by_root.clear();
193 self.blocked_roots.clear();
194 self.loop_headers = None;
195 self.traces.clear();
196 self.attempts.clear();
197 }
198
199 pub fn observe_hot_ip(&mut self, ip: usize, program: &Program) -> Option<usize> {
200 if !self.config.enabled || !native_jit_supported() {
201 return None;
202 }
203 if let Some(&trace_id) = self.compiled_by_root.get(&ip) {
204 return Some(trace_id);
205 }
206 if self.blocked_roots.contains(&ip) || !self.is_loop_header(program, ip) {
207 return None;
208 }
209
210 let count = self.hot_counts.entry(ip).or_insert(0);
211 *count = count.saturating_add(1);
212 if *count < self.config.hot_loop_threshold {
213 return None;
214 }
215
216 let line = program
217 .debug
218 .as_ref()
219 .and_then(|debug| debug.line_for_offset(ip));
220 let result = if self.config.hot_loop_threshold == 0 {
221 Err(JitNyiReason::HotLoopThresholdZero)
222 } else {
223 self.compile_trace(program, ip)
224 };
225 self.finish_attempt(ip, line, result)
226 }
227
228 pub fn trace_clone(&self, trace_id: usize) -> Option<JitTrace> {
229 self.traces.get(trace_id).cloned()
230 }
231
232 pub fn observe_exit_ip(&mut self, ip: usize, program: &Program) -> Option<usize> {
233 if !self.config.enabled || !native_jit_supported() {
234 return None;
235 }
236 if self.config.hot_loop_threshold > 1 {
237 return None;
238 }
239 if let Some(&trace_id) = self.compiled_by_root.get(&ip) {
240 return Some(trace_id);
241 }
242 if self.blocked_roots.contains(&ip) {
243 return None;
244 }
245
246 let line = program
247 .debug
248 .as_ref()
249 .and_then(|debug| debug.line_for_offset(ip));
250 let result = if self.config.hot_loop_threshold == 0 {
251 Err(JitNyiReason::HotLoopThresholdZero)
252 } else {
253 self.compile_trace(program, ip)
254 };
255 self.finish_attempt(ip, line, result)
256 }
257
258 pub fn trace_has_call(&self, trace_id: usize) -> bool {
259 self.traces
260 .get(trace_id)
261 .is_some_and(|trace| trace.has_call)
262 }
263
264 pub fn compiled_trace_for_ip(&self, ip: usize) -> Option<usize> {
265 self.compiled_by_root.get(&ip).copied()
266 }
267
268 pub fn mark_trace_executed(&mut self, trace_id: usize) {
269 if let Some(trace) = self.traces.get_mut(trace_id) {
270 trace.executions = trace.executions.saturating_add(1);
271 }
272 }
273
274 pub(crate) fn block_trace(&mut self, trace_id: usize) {
275 if let Some(trace) = self.traces.get(trace_id) {
276 self.compiled_by_root.remove(&trace.root_ip);
277 self.blocked_roots.insert(trace.root_ip);
278 }
279 }
280
281 pub fn snapshot(&self, runtime_metrics: JitMetrics) -> JitSnapshot {
282 JitSnapshot {
283 arch: std::env::consts::ARCH,
284 config: self.config,
285 traces: self.traces.clone(),
286 attempts: self.attempts.clone(),
287 metrics: self.aggregate_metrics(runtime_metrics),
288 nyi_reference: nyi_reference(),
289 }
290 }
291
292 pub fn dump_text(&self, debug: Option<&DebugInfo>, runtime_metrics: JitMetrics) -> String {
293 let mut out = String::new();
294 let metrics = self.aggregate_metrics(runtime_metrics);
295 out.push_str("trace-jit:\n");
296 out.push_str(&format!(" arch: {}\n", std::env::consts::ARCH));
297 out.push_str(&format!(" enabled: {}\n", self.config.enabled));
298 out.push_str(&format!(
299 " hot_loop_threshold: {}\n",
300 self.config.hot_loop_threshold
301 ));
302 out.push_str(&format!(" max_trace_len: {}\n", self.config.max_trace_len));
303 out.push_str(&format!(" compiled traces: {}\n", self.traces.len()));
304 out.push_str(&format!(" compile attempts: {}\n", self.attempts.len()));
305 out.push_str(&format!(
306 " boxed value sites: loads={} stores={}\n",
307 metrics.boxed_load_site_count, metrics.boxed_store_site_count
308 ));
309 out.push_str(&format!(
310 " trace exits: total={} guard_like={} loop_backs={}\n",
311 metrics.trace_exit_count,
312 metrics.guard_exit_count(),
313 metrics.native_loop_back_count
314 ));
315 out.push_str(&format!(
316 " interpreter fallbacks: {}\n",
317 metrics.helper_fallback_count
318 ));
319
320 for trace in &self.traces {
321 let line = trace
322 .start_line
323 .map(|value| value.to_string())
324 .unwrap_or_else(|| "-".to_string());
325 let source = debug
326 .and_then(|info| trace.start_line.and_then(|l| info.source_line(l)))
327 .unwrap_or_default();
328 out.push_str(&format!(
329 " trace#{} root_ip={} line={} terminal={:?} ops={} executions={}\n",
330 trace.id,
331 trace.root_ip,
332 line,
333 trace.terminal,
334 trace.op_names.len(),
335 trace.executions
336 ));
337 if !source.is_empty() {
338 out.push_str(&format!(" source: {}\n", source.trim()));
339 }
340 out.push_str(" ops:");
341 for op in &trace.op_names {
342 out.push_str(&format!(" {}", op));
343 }
344 out.push('\n');
345 out.push_str(&format!(
346 " ssa: blocks={} exits={}\n",
347 trace.ssa.blocks.len(),
348 trace.ssa.exits.len()
349 ));
350 for line in trace.ssa.render_text().lines() {
351 out.push_str(" ");
352 out.push_str(line);
353 out.push('\n');
354 }
355 }
356
357 let mut nyi = 0usize;
358 for attempt in &self.attempts {
359 if let Err(reason) = &attempt.result {
360 nyi = nyi.saturating_add(1);
361 let line = attempt
362 .line
363 .map(|value| value.to_string())
364 .unwrap_or_else(|| "-".to_string());
365 out.push_str(&format!(
366 " nyi root_ip={} line={} reason={}\n",
367 attempt.root_ip,
368 line,
369 reason.message()
370 ));
371 }
372 }
373 out.push_str(&format!(" nyi attempts: {nyi}\n"));
374
375 out.push_str(" nyi reference:\n");
376 for doc in nyi_reference() {
377 out.push_str(&format!(" - {}: {}\n", doc.item, doc.reason));
378 }
379
380 out
381 }
382
383 fn finish_attempt(
384 &mut self,
385 ip: usize,
386 line: Option<u32>,
387 result: Result<usize, JitNyiReason>,
388 ) -> Option<usize> {
389 match result {
390 Ok(trace_id) => {
391 self.attempts.push(JitAttempt {
392 root_ip: ip,
393 line,
394 result: Ok(trace_id),
395 });
396 self.compiled_by_root.insert(ip, trace_id);
397 Some(trace_id)
398 }
399 Err(reason) => {
400 self.attempts.push(JitAttempt {
401 root_ip: ip,
402 line,
403 result: Err(reason),
404 });
405 self.blocked_roots.insert(ip);
406 None
407 }
408 }
409 }
410
411 fn compile_trace(&mut self, program: &Program, root_ip: usize) -> Result<usize, JitNyiReason> {
412 let recorded = record_trace(program, root_ip, self.config.max_trace_len).map_err(to_nyi)?;
413 let id = self.traces.len();
414 let start_line = program
415 .debug
416 .as_ref()
417 .and_then(|debug| debug.line_for_offset(root_ip));
418 let trace = build_jit_trace(id, root_ip, start_line, recorded);
419 self.traces.push(trace);
420 Ok(id)
421 }
422
423 fn is_loop_header(&mut self, program: &Program, ip: usize) -> bool {
424 if self.loop_headers.is_none() {
425 self.loop_headers = Some(scan_loop_headers(program));
426 }
427 self.loop_headers
428 .as_ref()
429 .is_some_and(|headers| headers.contains(&ip))
430 }
431
432 fn aggregate_metrics(&self, mut runtime_metrics: JitMetrics) -> JitMetrics {
433 for trace in &self.traces {
434 runtime_metrics.boxed_load_site_count = runtime_metrics
435 .boxed_load_site_count
436 .saturating_add(boxed_load_site_count(&trace.ssa));
437 runtime_metrics.boxed_store_site_count = runtime_metrics
438 .boxed_store_site_count
439 .saturating_add(boxed_store_site_count(&trace.ssa));
440 if trace.terminal == JitTraceTerminal::LoopBack {
441 runtime_metrics.native_loop_back_count =
442 runtime_metrics.native_loop_back_count.saturating_add(1);
443 }
444 }
445 runtime_metrics
446 }
447}
448
449fn build_jit_trace(
450 id: usize,
451 root_ip: usize,
452 start_line: Option<u32>,
453 recorded: RecordedTrace,
454) -> JitTrace {
455 JitTrace {
456 id,
457 root_ip,
458 start_line,
459 has_call: recorded.has_call,
460 has_yielding_call: recorded.has_yielding_call,
461 op_names: recorded.op_names,
462 terminal: recorded.terminal,
463 executions: 0,
464 ssa: recorded.ssa,
465 }
466}
467
468fn to_nyi(err: TraceRecordError) -> JitNyiReason {
469 match err {
470 TraceRecordError::UnsupportedOpcode(op) => JitNyiReason::UnsupportedOpcode(op),
471 TraceRecordError::UnsupportedTrace(detail) => JitNyiReason::UnsupportedTrace(detail),
472 TraceRecordError::InvalidJumpTarget { target } => {
473 JitNyiReason::InvalidJumpTarget { target }
474 }
475 TraceRecordError::InvalidImmediate(kind) => JitNyiReason::InvalidImmediate(kind),
476 TraceRecordError::TraceTooLong { limit } => JitNyiReason::TraceTooLong { limit },
477 TraceRecordError::MissingTerminal => JitNyiReason::MissingTerminal,
478 TraceRecordError::InvalidLocal(_)
479 | TraceRecordError::StackUnderflow
480 | TraceRecordError::TypeMismatch { .. }
481 | TraceRecordError::LiveStackOnBackedge { .. }
482 | TraceRecordError::InvalidIr(_) => JitNyiReason::UnsupportedTrace(err.to_string()),
483 }
484}
485
486fn scan_loop_headers(program: &Program) -> HashSet<usize> {
487 let mut headers = HashSet::new();
488 let code = &program.code;
489 let mut ip = 0usize;
490
491 while ip < code.len() {
492 let opcode = code[ip];
493 let instr_ip = ip;
494 ip = ip.saturating_add(1);
495 match opcode {
496 x if x == OpCode::Ldc as u8 => {
497 if read_u32(code, &mut ip).is_none() {
498 break;
499 }
500 }
501 x if x == OpCode::Br as u8 || x == OpCode::Brfalse as u8 => {
502 let Some(target_u32) = read_u32(code, &mut ip) else {
503 break;
504 };
505 let target = target_u32 as usize;
506 if target <= instr_ip {
507 headers.insert(target);
508 }
509 }
510 x if x == OpCode::Ldloc as u8 || x == OpCode::Stloc as u8 => {
511 if read_u8(code, &mut ip).is_none() {
512 break;
513 }
514 }
515 x if x == OpCode::Call as u8 => {
516 if read_u16(code, &mut ip).is_none() {
517 break;
518 }
519 if read_u8(code, &mut ip).is_none() {
520 break;
521 }
522 }
523 _ => {}
524 }
525 }
526
527 headers
528}
529
530fn read_u8(code: &[u8], ip: &mut usize) -> Option<u8> {
531 let value = *code.get(*ip)?;
532 *ip = ip.saturating_add(1);
533 Some(value)
534}
535
536fn read_u16(code: &[u8], ip: &mut usize) -> Option<u16> {
537 if ip.saturating_add(2) > code.len() {
538 return None;
539 }
540 let bytes = [code[*ip], code[*ip + 1]];
541 *ip = ip.saturating_add(2);
542 Some(u16::from_le_bytes(bytes))
543}
544
545fn read_u32(code: &[u8], ip: &mut usize) -> Option<u32> {
546 if ip.saturating_add(4) > code.len() {
547 return None;
548 }
549 let bytes = [code[*ip], code[*ip + 1], code[*ip + 2], code[*ip + 3]];
550 *ip = ip.saturating_add(4);
551 Some(u32::from_le_bytes(bytes))
552}
553
554fn nyi_reference() -> Vec<JitNyiDoc> {
555 vec![
556 JitNyiDoc {
557 item: "Oversized traces",
558 reason: "trace recording stops at max_trace_len",
559 },
560 JitNyiDoc {
561 item: "Unsupported native JIT targets",
562 reason: "native emission currently supports x86_64 on windows plus unix non-macos, and aarch64 on linux/macos",
563 },
564 ]
565}
566
567#[cfg(test)]
568mod tests {
569 use super::*;
570 use crate::{BytecodeBuilder, Value};
571
572 #[test]
573 fn scan_loop_headers_finds_backward_targets() {
574 let mut bc = BytecodeBuilder::new();
575 let root_ip = bc.position();
576 bc.ldc(0);
577 let branch_ip = bc.position();
578 bc.br(root_ip);
579 let program = Program::new(vec![Value::Int(1)], bc.finish());
580
581 let headers = scan_loop_headers(&program);
582 assert!(headers.contains(&(root_ip as usize)));
583 assert!(!headers.contains(&(branch_ip as usize)));
584 }
585}