1#![allow(unsafe_code)]
2
3use std::collections::{HashMap, HashSet, VecDeque};
4
5use cranelift_codegen::ir::types;
6use cranelift_codegen::ir::{AbiParam, Block, Function, InstBuilder, MachMemFlags, Signature, Value as ClifValue};
7use cranelift_frontend::{FunctionBuilder, FunctionBuilderContext};
8use cranelift_module::{Linkage, Module};
9
10use crate::{Chunk, Instruction, Opcode, Value};
11
12use super::error::CompileError;
13use super::exit::{JIT_BUDGET, JIT_CONTINUE, JIT_RETURN, JIT_TRAP};
14use super::frame::{
15 JitEntry, OFF_BUDGET, OFF_CALL_DEPTH, OFF_CALL_STACK, OFF_EXIT_KIND, OFF_FUNCTION,
16 OFF_PC, OFF_REGISTER_COUNT, OFF_RETURN_REG, MAX_JIT_CALL_DEPTH,
17};
18use super::trace::{CompiledTrace, TraceKey, TraceSpan, MAX_TRACE_LENGTH};
19
20pub struct RegisterMap {
22 values: Vec<Option<ClifValue>>,
23}
24
25impl RegisterMap {
26 pub fn new(registers: usize) -> Self {
27 Self {
28 values: vec![None; registers],
29 }
30 }
31
32 pub fn get(&self, reg: usize) -> Option<ClifValue> {
33 self.values.get(reg).copied().flatten()
34 }
35
36 pub fn set(&mut self, reg: usize, value: ClifValue) {
37 if let Some(slot) = self.values.get_mut(reg) {
38 *slot = Some(value);
39 }
40 }
41
42 pub fn clear(&mut self) {
43 for slot in &mut self.values {
44 *slot = None;
45 }
46 }
47}
48
49pub struct TraceCompiler<'a> {
50 module: &'a mut cranelift_jit::JITModule,
51}
52
53impl<'a> TraceCompiler<'a> {
54 pub fn new(module: &'a mut cranelift_jit::JITModule) -> Self {
55 Self { module }
56 }
57
58 pub fn compile_trace(
59 &mut self,
60 chunk: &Chunk,
61 key: TraceKey,
62 ) -> Result<CompiledTrace, CompileError> {
63 let def = chunk.function(key.function).ok_or(CompileError::EmptyTrace {
64 function: key.function,
65 pc: key.entry_pc,
66 })?;
67 let region = collect_trace_region(chunk, key)?;
68 if region.blocks.is_empty() {
69 return Err(CompileError::EmptyTrace {
70 function: key.function,
71 pc: key.entry_pc,
72 });
73 }
74
75 let end_pc = region_max_pc(®ion, key)?;
76 let span = TraceSpan {
77 function: key.function,
78 range: key.entry_pc..end_pc.saturating_add(1),
79 };
80
81 let target = self.module.target_config();
82 let pointer_type = target.pointer_type();
83 let mut sig = Signature::new(target.default_call_conv);
84 sig.params.push(AbiParam::new(pointer_type));
85
86 let func_id = self
87 .module
88 .declare_function(
89 &format!("trace_{}_{}", key.function, key.entry_pc),
90 Linkage::Local,
91 &sig,
92 )
93 .map_err(|e| CompileError::Module(e.to_string()))?;
94
95 let mut ctx = cranelift_codegen::Context::new();
96 ctx.func = Function::with_name_signature(
97 cranelift_codegen::ir::UserFuncName::user(0, func_id.as_u32()),
98 sig,
99 );
100
101 let mut func_ctx = FunctionBuilderContext::new();
102 let mut builder = FunctionBuilder::new(&mut ctx.func, &mut func_ctx);
103 let entry = builder.create_block();
104 builder.append_block_params_for_function_params(entry);
105 builder.switch_to_block(entry);
106
107 let frame_ptr = builder.block_params(entry)[0];
108 let slots_ptr = builder.ins().load(pointer_type, MachMemFlags::new(), frame_ptr, 0);
109 let budget_ptr = field_ptr(&mut builder, frame_ptr, pointer_type, OFF_BUDGET);
110 let call_depth_ptr = field_ptr(&mut builder, frame_ptr, pointer_type, OFF_CALL_DEPTH);
111 let call_stack_ptr = {
112 let addr = field_ptr(&mut builder, frame_ptr, pointer_type, OFF_CALL_STACK);
113 builder
114 .ins()
115 .load(pointer_type, MachMemFlags::new(), addr, 0)
116 };
117
118 let supports_calls = region.blocks.iter().any(|block| {
119 block
120 .pcs
121 .iter()
122 .any(|pc| chunk.code[*pc as usize].op == Opcode::Call)
123 });
124
125 let trap_block = builder.create_block();
126 let return_block = builder.create_block();
127
128 let mut pc_to_block = HashMap::new();
129 for block in ®ion.blocks {
130 pc_to_block.insert(block.start_pc, builder.create_block());
131 }
132
133 let entry_clif = *pc_to_block
134 .get(&key.entry_pc)
135 .ok_or(CompileError::EmptyTrace {
136 function: key.function,
137 pc: key.entry_pc,
138 })?;
139 builder.ins().jump(entry_clif, &[]);
140
141 for block in ®ion.blocks {
142 let clif_block = pc_to_block[&block.start_pc];
143 builder.switch_to_block(clif_block);
144 let mut registers = RegisterMap::new(def.num_registers as usize);
145
146 for pc in &block.pcs {
147 emit_budget_tick(
148 &mut builder,
149 budget_ptr,
150 frame_ptr,
151 pointer_type,
152 *pc,
153 );
154 let instr = chunk.code[*pc as usize];
155 let mut emit_ctx = EmitContext {
156 builder: &mut builder,
157 chunk,
158 registers: &mut registers,
159 slots_ptr,
160 frame_ptr,
161 pointer_type,
162 trap_block,
163 return_block,
164 pc_to_block: &pc_to_block,
165 call_depth_ptr,
166 call_stack_ptr,
167 supports_calls,
168 };
169 emit_instruction(&mut emit_ctx, *pc, instr)?;
170 }
171 }
172
173 builder.switch_to_block(trap_block);
174 builder.seal_block(trap_block);
175 let trap_pc = region_max_pc(®ion, key)?;
176 write_exit(
177 &mut builder,
178 frame_ptr,
179 pointer_type,
180 JIT_TRAP,
181 trap_pc,
182 0,
183 );
184
185 builder.switch_to_block(return_block);
186 builder.seal_block(return_block);
187 let (last_pc, return_reg) = match last_return_in_region(chunk, ®ion) {
188 Some(pair) => pair,
189 None => (key.entry_pc, 0),
191 };
192 write_exit(
193 &mut builder,
194 frame_ptr,
195 pointer_type,
196 JIT_RETURN,
197 last_pc,
198 return_reg,
199 );
200
201 builder.seal_all_blocks();
202 builder.finalize(self.module.target_config());
203
204 self.module
205 .define_function(func_id, &mut ctx)
206 .map_err(|e| CompileError::Module(e.to_string()))?;
207 self.module.clear_context(&mut ctx);
208 self.module
209 .finalize_definitions()
210 .map_err(|e| CompileError::Module(e.to_string()))?;
211
212 let code = self.module.get_finalized_function(func_id);
213 let entry: JitEntry = unsafe { std::mem::transmute(code) };
214
215 Ok(CompiledTrace { span, entry })
216 }
217}
218
219struct TraceBlock {
220 start_pc: u32,
221 pcs: Vec<u32>,
222}
223
224struct TraceRegion {
225 blocks: Vec<TraceBlock>,
226}
227
228fn region_max_pc(region: &TraceRegion, key: TraceKey) -> Result<u32, CompileError> {
229 region
230 .blocks
231 .iter()
232 .flat_map(|b| b.pcs.iter())
233 .copied()
234 .max()
235 .ok_or(CompileError::EmptyTrace {
236 function: key.function,
237 pc: key.entry_pc,
238 })
239}
240
241fn last_return_in_region(chunk: &Chunk, region: &TraceRegion) -> Option<(u32, u32)> {
242 region
243 .blocks
244 .iter()
245 .flat_map(|b| b.pcs.iter().map(|pc| (*pc, chunk.code[*pc as usize])))
246 .rev()
247 .find(|(_, instr)| instr.op == Opcode::Return)
248 .map(|(pc, instr)| (pc, u32::from(instr.a)))
249}
250
251fn jump_target(jpc: u32, imm: i32) -> u32 {
252 (i64::from(jpc) + 1 + i64::from(imm)) as u32
253}
254
255fn branch_target(bpc: u32, imm: i32) -> u32 {
256 jump_target(bpc, imm)
257}
258
259struct EmitContext<'a, 'b> {
260 builder: &'a mut FunctionBuilder<'b>,
261 chunk: &'a Chunk,
262 registers: &'a mut RegisterMap,
263 slots_ptr: ClifValue,
264 frame_ptr: ClifValue,
265 pointer_type: types::Type,
266 trap_block: Block,
267 return_block: Block,
268 pc_to_block: &'a HashMap<u32, Block>,
269 call_depth_ptr: ClifValue,
270 call_stack_ptr: ClifValue,
271 supports_calls: bool,
272}
273
274fn emit_instruction(
275 ctx: &mut EmitContext<'_, '_>,
276 pc: u32,
277 instr: Instruction,
278) -> Result<(), CompileError> {
279 let builder = &mut *ctx.builder;
280 let chunk = ctx.chunk;
281 let registers = &mut *ctx.registers;
282 let slots_ptr = ctx.slots_ptr;
283 let frame_ptr = ctx.frame_ptr;
284 let pointer_type = ctx.pointer_type;
285 let trap_block = ctx.trap_block;
286 let return_block = ctx.return_block;
287 let supports_calls = ctx.supports_calls;
288 let pc_to_block = ctx.pc_to_block;
289 let call_depth_ptr = ctx.call_depth_ptr;
290 let call_stack_ptr = ctx.call_stack_ptr;
291 match instr.op {
292 Opcode::LoadImm => {
293 let v = builder.ins().iconst(types::I64, i64::from(instr.imm));
294 registers.set(instr.a as usize, v);
295 }
296 Opcode::LoadConst => {
297 let idx = instr.imm as u32;
298 let konst = chunk.constant(idx).ok_or(CompileError::UnsupportedOpcode {
299 opcode: Opcode::LoadConst,
300 pc,
301 })?;
302 let v = match konst {
303 Value::Int(i) => builder.ins().iconst(types::I64, *i),
304 Value::Bool(b) => builder.ins().iconst(types::I64, i64::from(*b)),
305 _ => {
306 return Err(CompileError::UnsupportedOpcode {
307 opcode: Opcode::LoadConst,
308 pc,
309 });
310 }
311 };
312 registers.set(instr.a as usize, v);
313 }
314 Opcode::Move => {
315 let v = load_reg(
316 builder,
317 registers,
318 slots_ptr,
319 pointer_type,
320 instr.b,
321 pc,
322 )?;
323 registers.set(instr.a as usize, v);
324 }
325 Opcode::Add | Opcode::Sub | Opcode::Mul | Opcode::Eq | Opcode::Lt | Opcode::Le => {
326 let lhs = load_reg(
327 builder,
328 registers,
329 slots_ptr,
330 pointer_type,
331 instr.b,
332 pc,
333 )?;
334 let rhs = load_reg(
335 builder,
336 registers,
337 slots_ptr,
338 pointer_type,
339 instr.c,
340 pc,
341 )?;
342 let result = match instr.op {
343 Opcode::Add => builder.ins().iadd(lhs, rhs),
344 Opcode::Sub => builder.ins().isub(lhs, rhs),
345 Opcode::Mul => builder.ins().imul(lhs, rhs),
346 Opcode::Eq => {
347 let cmp = builder.ins().icmp(
348 cranelift_codegen::ir::condcodes::IntCC::Equal,
349 lhs,
350 rhs,
351 );
352 bool_as_i64(builder, cmp)
353 }
354 Opcode::Lt => {
355 let cmp = builder.ins().icmp(
356 cranelift_codegen::ir::condcodes::IntCC::SignedLessThan,
357 lhs,
358 rhs,
359 );
360 bool_as_i64(builder, cmp)
361 }
362 Opcode::Le => {
363 let cmp = builder.ins().icmp(
364 cranelift_codegen::ir::condcodes::IntCC::SignedLessThanOrEqual,
365 lhs,
366 rhs,
367 );
368 bool_as_i64(builder, cmp)
369 }
370 _ => unreachable!(),
371 };
372 registers.set(instr.a as usize, result);
373 }
374 Opcode::Div | Opcode::Mod => {
375 let lhs = load_reg(
376 builder,
377 registers,
378 slots_ptr,
379 pointer_type,
380 instr.b,
381 pc,
382 )?;
383 let rhs = load_reg(
384 builder,
385 registers,
386 slots_ptr,
387 pointer_type,
388 instr.c,
389 pc,
390 )?;
391 flush_registers(builder, registers, slots_ptr, pointer_type);
392 let zero = builder.ins().iconst(types::I64, 0);
393 let is_zero = builder.ins().icmp(
394 cranelift_codegen::ir::condcodes::IntCC::Equal,
395 rhs,
396 zero,
397 );
398 let ok = builder.create_block();
399 builder.ins().brif(is_zero, trap_block, &[], ok, &[]);
400 builder.switch_to_block(ok);
401 builder.seal_block(ok);
402 let result = if instr.op == Opcode::Div {
403 builder.ins().sdiv(lhs, rhs)
404 } else {
405 builder.ins().srem(lhs, rhs)
406 };
407 registers.set(instr.a as usize, result);
408 }
409 Opcode::Neg => {
410 let v = load_reg(
411 builder,
412 registers,
413 slots_ptr,
414 pointer_type,
415 instr.b,
416 pc,
417 )?;
418 let zero = builder.ins().iconst(types::I64, 0);
419 registers.set(instr.a as usize, builder.ins().isub(zero, v));
420 }
421 Opcode::Jump => {
422 flush_registers(builder, registers, slots_ptr, pointer_type);
423 let target_pc = jump_target(pc, instr.imm);
424 branch_to_pc(
425 builder,
426 frame_ptr,
427 pointer_type,
428 pc_to_block,
429 target_pc,
430 );
431 }
432 Opcode::Branch => {
433 let cond = load_reg(
434 builder,
435 registers,
436 slots_ptr,
437 pointer_type,
438 instr.a,
439 pc,
440 )?;
441 flush_registers(builder, registers, slots_ptr, pointer_type);
442 let zero = builder.ins().iconst(types::I64, 0);
443 let is_falsy = builder.ins().icmp(
444 cranelift_codegen::ir::condcodes::IntCC::Equal,
445 cond,
446 zero,
447 );
448 let fall_pc = pc + 1;
449 let falsy_pc = branch_target(pc, instr.imm);
450 match (pc_to_block.get(&fall_pc), pc_to_block.get(&falsy_pc)) {
451 (Some(fall), Some(falsy)) => {
452 builder.ins().brif(is_falsy, *falsy, &[], *fall, &[]);
453 }
454 (Some(fall), None) => {
455 let side = builder.create_block();
456 builder.ins().brif(is_falsy, side, &[], *fall, &[]);
457 builder.switch_to_block(side);
458 builder.seal_block(side);
459 write_exit(
460 builder,
461 frame_ptr,
462 pointer_type,
463 JIT_CONTINUE,
464 falsy_pc,
465 0,
466 );
467 }
468 (None, Some(falsy)) => {
469 let side = builder.create_block();
470 builder.ins().brif(is_falsy, *falsy, &[], side, &[]);
471 builder.switch_to_block(side);
472 builder.seal_block(side);
473 write_exit(
474 builder,
475 frame_ptr,
476 pointer_type,
477 JIT_CONTINUE,
478 fall_pc,
479 0,
480 );
481 }
482 (None, None) => {
483 write_exit(
484 builder,
485 frame_ptr,
486 pointer_type,
487 JIT_CONTINUE,
488 falsy_pc,
489 0,
490 );
491 }
492 }
493 }
494 Opcode::Return => {
495 let v = load_reg(
496 builder,
497 registers,
498 slots_ptr,
499 pointer_type,
500 instr.a,
501 pc,
502 )?;
503 store_slot(builder, slots_ptr, pointer_type, instr.a, v);
504 if supports_calls {
505 emit_return(
506 builder,
507 frame_ptr,
508 pointer_type,
509 call_depth_ptr,
510 call_stack_ptr,
511 slots_ptr,
512 registers,
513 instr.a,
514 pc,
515 );
516 } else {
517 write_u32_field(builder, frame_ptr, pointer_type, OFF_PC, pc);
518 builder.ins().jump(return_block, &[]);
519 }
520 }
521 Opcode::Nop => {}
522 Opcode::Call => {
523 emit_call(
524 builder,
525 chunk,
526 registers,
527 slots_ptr,
528 frame_ptr,
529 pointer_type,
530 pc_to_block,
531 call_depth_ptr,
532 call_stack_ptr,
533 pc,
534 instr,
535 )?;
536 }
537 other => {
538 return Err(CompileError::UnsupportedOpcode { opcode: other, pc });
539 }
540 }
541 Ok(())
542}
543
544fn branch_to_pc(
545 builder: &mut FunctionBuilder<'_>,
546 frame_ptr: ClifValue,
547 pointer_type: types::Type,
548 pc_to_block: &HashMap<u32, Block>,
549 target_pc: u32,
550) {
551 if let Some(block) = pc_to_block.get(&target_pc) {
552 builder.ins().jump(*block, &[]);
553 } else {
554 write_exit(
555 builder,
556 frame_ptr,
557 pointer_type,
558 JIT_CONTINUE,
559 target_pc,
560 0,
561 );
562 }
563}
564
565fn flush_registers(
566 builder: &mut FunctionBuilder<'_>,
567 registers: &RegisterMap,
568 slots_ptr: ClifValue,
569 pointer_type: types::Type,
570) {
571 for (reg, value) in registers.values.iter().enumerate() {
572 if let Some(v) = *value {
573 store_slot(builder, slots_ptr, pointer_type, reg as u8, v);
574 }
575 }
576}
577
578fn load_reg(
579 builder: &mut FunctionBuilder<'_>,
580 registers: &mut RegisterMap,
581 slots_ptr: ClifValue,
582 pointer_type: types::Type,
583 reg: u8,
584 _pc: u32,
585) -> Result<ClifValue, CompileError> {
586 if let Some(v) = registers.get(reg as usize) {
587 return Ok(v);
588 }
589 let offset = builder.ins().iconst(pointer_type, i64::from(u32::from(reg) * 8));
590 let addr = builder.ins().iadd(slots_ptr, offset);
591 let v = builder
592 .ins()
593 .load(types::I64, MachMemFlags::new(), addr, 0);
594 registers.set(reg as usize, v);
595 Ok(v)
596}
597
598fn collect_trace_region(chunk: &Chunk, key: TraceKey) -> Result<TraceRegion, CompileError> {
599 let mut queue = VecDeque::from([key.entry_pc]);
600 let mut seen_starts = HashSet::new();
601 let mut blocks = Vec::new();
602 let mut total = 0usize;
603
604 while let Some(start) = queue.pop_front() {
605 if !seen_starts.insert(start) {
606 continue;
607 }
608 if total >= MAX_TRACE_LENGTH {
609 break;
610 }
611
612 let mut pcs = Vec::new();
613 let mut pc = start;
614 while let Some(instr) = chunk.code.get(pc as usize).copied() {
615 if is_effect_opcode(instr.op) {
616 break;
617 }
618 if instr.op == Opcode::LoadConst {
619 let idx = instr.imm as u32;
620 match chunk.constant(idx) {
621 Some(Value::Int(_) | Value::Bool(_)) => {}
622 _ => break,
623 }
624 }
625 pcs.push(pc);
626 total += 1;
627 if total >= MAX_TRACE_LENGTH {
628 break;
629 }
630
631 match instr.op {
632 Opcode::Return => break,
633 Opcode::Jump => {
634 queue.push_back(jump_target(pc, instr.imm));
635 break;
636 }
637 Opcode::Branch => {
638 queue.push_back(pc + 1);
639 queue.push_back(branch_target(pc, instr.imm));
640 break;
641 }
642 Opcode::Call => {
643 let callee = instr.imm as u32;
644 if let Some(def) = chunk.function(callee) {
645 queue.push_back(def.entry);
646 }
647 queue.push_back(pc + 1);
648 break;
649 }
650 _ => pc += 1,
651 }
652 }
653
654 if !pcs.is_empty() {
655 blocks.push(TraceBlock { start_pc: start, pcs });
656 }
657 }
658
659 Ok(TraceRegion { blocks })
660}
661
662fn is_effect_opcode(op: Opcode) -> bool {
663 matches!(
664 op,
665 Opcode::CallNative
666 | Opcode::Spawn
667 | Opcode::Yield
668 | Opcode::Sleep
669 | Opcode::Exit
670 | Opcode::SelfPid
671 | Opcode::Send
672 | Opcode::Receive
673 | Opcode::ReceiveTimeout
674 | Opcode::ReceiveMatch
675 | Opcode::ReceiveMatchImm
676 | Opcode::Ask
677 | Opcode::AskTimeout
678 | Opcode::Monitor
679 | Opcode::Demonitor
680 | Opcode::Link
681 | Opcode::Unlink
682 | Opcode::Trap
683 | Opcode::Halt
684 )
685}
686
687fn field_ptr(
688 builder: &mut FunctionBuilder<'_>,
689 base: ClifValue,
690 pointer_type: types::Type,
691 offset: i32,
692) -> ClifValue {
693 let off = builder.ins().iconst(pointer_type, i64::from(offset));
694 builder.ins().iadd(base, off)
695}
696
697fn store_slot(
698 builder: &mut FunctionBuilder<'_>,
699 slots: ClifValue,
700 pointer_type: types::Type,
701 reg: u8,
702 value: ClifValue,
703) {
704 let offset = builder.ins().iconst(pointer_type, i64::from(u32::from(reg) * 8));
705 let addr = builder.ins().iadd(slots, offset);
706 builder.ins().store(MachMemFlags::new(), value, addr, 0);
707}
708
709fn bool_as_i64(builder: &mut FunctionBuilder<'_>, cond: ClifValue) -> ClifValue {
710 let one = builder.ins().iconst(types::I64, 1);
711 let zero = builder.ins().iconst(types::I64, 0);
712 builder.ins().select(cond, one, zero)
713}
714
715fn emit_budget_tick(
716 builder: &mut FunctionBuilder<'_>,
717 budget_ptr: ClifValue,
718 frame_ptr: ClifValue,
719 pointer_type: types::Type,
720 pc: u32,
721) {
722 let budget = builder.ins().load(types::I32, MachMemFlags::new(), budget_ptr, 0);
723 let zero = builder.ins().iconst(types::I32, 0);
724 let exhausted = builder.ins().icmp(
725 cranelift_codegen::ir::condcodes::IntCC::Equal,
726 budget,
727 zero,
728 );
729 let continue_insn = builder.create_block();
730 let budget_exit = builder.create_block();
731 builder.ins().brif(exhausted, budget_exit, &[], continue_insn, &[]);
732 builder.switch_to_block(budget_exit);
733 builder.seal_block(budget_exit);
734 write_exit(builder, frame_ptr, pointer_type, JIT_BUDGET, pc, 0);
735 builder.switch_to_block(continue_insn);
736 builder.seal_block(continue_insn);
737 let one = builder.ins().iconst(types::I32, 1);
738 let new_budget = builder.ins().isub(budget, one);
739 builder.ins().store(MachMemFlags::new(), new_budget, budget_ptr, 0);
740}
741
742fn write_u32_field(
743 builder: &mut FunctionBuilder<'_>,
744 frame_ptr: ClifValue,
745 pointer_type: types::Type,
746 offset: i32,
747 value: u32,
748) {
749 let addr = field_ptr(builder, frame_ptr, pointer_type, offset);
750 let val = builder.ins().iconst(types::I32, i64::from(value));
751 builder.ins().store(MachMemFlags::new(), val, addr, 0);
752}
753
754fn write_exit(
755 builder: &mut FunctionBuilder<'_>,
756 frame_ptr: ClifValue,
757 pointer_type: types::Type,
758 kind: u32,
759 pc: u32,
760 return_reg: u32,
761) {
762 write_u32_field(builder, frame_ptr, pointer_type, OFF_EXIT_KIND, kind);
763 write_u32_field(builder, frame_ptr, pointer_type, OFF_PC, pc);
764 write_u32_field(builder, frame_ptr, pointer_type, OFF_RETURN_REG, return_reg);
765 builder.ins().return_(&[]);
766}
767
768#[allow(clippy::too_many_arguments)]
769fn emit_call(
770 builder: &mut FunctionBuilder<'_>,
771 chunk: &Chunk,
772 registers: &mut RegisterMap,
773 slots_ptr: ClifValue,
774 frame_ptr: ClifValue,
775 pointer_type: types::Type,
776 pc_to_block: &HashMap<u32, Block>,
777 call_depth_ptr: ClifValue,
778 call_stack_ptr: ClifValue,
779 pc: u32,
780 instr: Instruction,
781) -> Result<(), CompileError> {
782 let callee = instr.imm as u32;
783 let argc = instr.b;
784 let dst = instr.a;
785 let def = chunk.function(callee).ok_or(CompileError::UnsupportedOpcode {
786 opcode: Opcode::Call,
787 pc,
788 })?;
789 let callee_entry = def.entry;
790 let Some(callee_block) = pc_to_block.get(&callee_entry) else {
791 flush_registers(builder, registers, slots_ptr, pointer_type);
792 write_exit(
793 builder,
794 frame_ptr,
795 pointer_type,
796 JIT_CONTINUE,
797 pc,
798 0,
799 );
800 return Ok(());
801 };
802
803 flush_registers(builder, registers, slots_ptr, pointer_type);
804
805 let depth = builder
806 .ins()
807 .load(types::I32, MachMemFlags::new(), call_depth_ptr, 0);
808 let max_depth = builder
809 .ins()
810 .iconst(types::I32, i64::from(MAX_JIT_CALL_DEPTH as u32));
811 let too_deep = builder.ins().icmp(
812 cranelift_codegen::ir::condcodes::IntCC::SignedGreaterThanOrEqual,
813 depth,
814 max_depth,
815 );
816 let effect_exit = builder.create_block();
817 let call_body = builder.create_block();
818 builder
819 .ins()
820 .brif(too_deep, effect_exit, &[], call_body, &[]);
821 builder.switch_to_block(effect_exit);
822 builder.seal_block(effect_exit);
823 write_exit(
824 builder,
825 frame_ptr,
826 pointer_type,
827 super::exit::JIT_EFFECT,
828 pc,
829 0,
830 );
831 builder.switch_to_block(call_body);
832 builder.seal_block(call_body);
833
834 let record_bytes = builder.ins().iconst(
835 types::I32,
836 i64::from(std::mem::size_of::<super::frame::JitCallRecord>() as u32),
837 );
838 let record_off = builder.ins().imul(depth, record_bytes);
839 let record_off_ptr = builder.ins().uextend(pointer_type, record_off);
840 let record_addr = builder.ins().iadd(call_stack_ptr, record_off_ptr);
841
842 let return_pc = builder.ins().iconst(types::I32, i64::from(pc + 1));
843 builder
844 .ins()
845 .store(MachMemFlags::new(), return_pc, record_addr, 0);
846
847 let func_ptr = field_ptr(builder, frame_ptr, pointer_type, OFF_FUNCTION);
848 let caller_fn = builder
849 .ins()
850 .load(types::I32, MachMemFlags::new(), func_ptr, 0);
851 let caller_fn_addr = {
852 let off = builder.ins().iconst(pointer_type, 4);
853 builder.ins().iadd(record_addr, off)
854 };
855 builder
856 .ins()
857 .store(MachMemFlags::new(), caller_fn, caller_fn_addr, 0);
858
859 let dest = builder.ins().iconst(types::I32, i64::from(u32::from(dst)));
860 let dest_addr = {
861 let off = builder.ins().iconst(pointer_type, 8);
862 builder.ins().iadd(record_addr, off)
863 };
864 builder
865 .ins()
866 .store(MachMemFlags::new(), dest, dest_addr, 0);
867
868 let reg_ptr = field_ptr(builder, frame_ptr, pointer_type, OFF_REGISTER_COUNT);
869 let caller_regs = builder
870 .ins()
871 .load(types::I32, MachMemFlags::new(), reg_ptr, 0);
872 let caller_regs_addr = {
873 let off = builder.ins().iconst(pointer_type, 12);
874 builder.ins().iadd(record_addr, off)
875 };
876 builder
877 .ins()
878 .store(MachMemFlags::new(), caller_regs, caller_regs_addr, 0);
879
880 let one = builder.ins().iconst(types::I32, 1);
881 let new_depth = builder.ins().iadd(depth, one);
882 builder
883 .ins()
884 .store(MachMemFlags::new(), new_depth, call_depth_ptr, 0);
885
886 for i in 0..argc {
887 let src_reg = u32::from(dst) + u32::from(i);
888 let src_offset = builder.ins().iconst(pointer_type, i64::from(src_reg * 8));
889 let src_addr = builder.ins().iadd(slots_ptr, src_offset);
890 let value = builder
891 .ins()
892 .load(types::I64, MachMemFlags::new(), src_addr, 0);
893 let dst_offset = builder.ins().iconst(pointer_type, i64::from(u32::from(i) * 8));
894 let dst_addr = builder.ins().iadd(slots_ptr, dst_offset);
895 builder
896 .ins()
897 .store(MachMemFlags::new(), value, dst_addr, 0);
898 }
899
900 let callee_fn = builder.ins().iconst(types::I32, i64::from(callee));
901 builder
902 .ins()
903 .store(MachMemFlags::new(), callee_fn, func_ptr, 0);
904 let callee_regs = builder.ins().iconst(types::I32, i64::from(def.num_registers));
905 builder
906 .ins()
907 .store(MachMemFlags::new(), callee_regs, reg_ptr, 0);
908
909 registers.clear();
910 builder.ins().jump(*callee_block, &[]);
911 Ok(())
912}
913
914#[allow(clippy::too_many_arguments)]
915fn emit_return(
916 builder: &mut FunctionBuilder<'_>,
917 frame_ptr: ClifValue,
918 pointer_type: types::Type,
919 call_depth_ptr: ClifValue,
920 call_stack_ptr: ClifValue,
921 slots_ptr: ClifValue,
922 registers: &mut RegisterMap,
923 return_reg: u8,
924 pc: u32,
925) {
926 let depth = builder
927 .ins()
928 .load(types::I32, MachMemFlags::new(), call_depth_ptr, 0);
929 let zero = builder.ins().iconst(types::I32, 0);
930 let is_outer = builder.ins().icmp(
931 cranelift_codegen::ir::condcodes::IntCC::Equal,
932 depth,
933 zero,
934 );
935 let outer = builder.create_block();
936 let inner = builder.create_block();
937 builder
938 .ins()
939 .brif(is_outer, outer, &[], inner, &[]);
940 builder.switch_to_block(outer);
941 builder.seal_block(outer);
942 write_exit(
943 builder,
944 frame_ptr,
945 pointer_type,
946 JIT_RETURN,
947 pc,
948 u32::from(return_reg),
949 );
950 builder.switch_to_block(inner);
951 builder.seal_block(inner);
952
953 let one = builder.ins().iconst(types::I32, 1);
954 let new_depth = builder.ins().isub(depth, one);
955 builder
956 .ins()
957 .store(MachMemFlags::new(), new_depth, call_depth_ptr, 0);
958
959 let record_bytes = builder.ins().iconst(
960 types::I32,
961 i64::from(std::mem::size_of::<super::frame::JitCallRecord>() as u32),
962 );
963 let record_off = builder.ins().imul(new_depth, record_bytes);
964 let record_off_ptr = builder.ins().uextend(pointer_type, record_off);
965 let record_addr = builder.ins().iadd(call_stack_ptr, record_off_ptr);
966
967 let return_pc = builder
968 .ins()
969 .load(types::I32, MachMemFlags::new(), record_addr, 0);
970 let return_fn_addr = {
971 let off = builder.ins().iconst(pointer_type, 4);
972 builder.ins().iadd(record_addr, off)
973 };
974 let return_fn = builder
975 .ins()
976 .load(types::I32, MachMemFlags::new(), return_fn_addr, 0);
977 let dest_addr = {
978 let off = builder.ins().iconst(pointer_type, 8);
979 builder.ins().iadd(record_addr, off)
980 };
981 let dest_reg = builder
982 .ins()
983 .load(types::I32, MachMemFlags::new(), dest_addr, 0);
984 let caller_regs_addr = {
985 let off = builder.ins().iconst(pointer_type, 12);
986 builder.ins().iadd(record_addr, off)
987 };
988 let caller_regs = builder
989 .ins()
990 .load(types::I32, MachMemFlags::new(), caller_regs_addr, 0);
991
992 let ret_offset = builder.ins().iconst(pointer_type, i64::from(u32::from(return_reg) * 8));
993 let ret_addr = builder.ins().iadd(slots_ptr, ret_offset);
994 let ret_val = builder
995 .ins()
996 .load(types::I64, MachMemFlags::new(), ret_addr, 0);
997 let eight = builder.ins().iconst(types::I32, 8);
998 let dest_offset = builder.ins().imul(dest_reg, eight);
999 let dest_offset_ptr = builder.ins().uextend(pointer_type, dest_offset);
1000 let dest_slot = builder.ins().iadd(slots_ptr, dest_offset_ptr);
1001 builder
1002 .ins()
1003 .store(MachMemFlags::new(), ret_val, dest_slot, 0);
1004
1005 let func_ptr = field_ptr(builder, frame_ptr, pointer_type, OFF_FUNCTION);
1006 builder
1007 .ins()
1008 .store(MachMemFlags::new(), return_fn, func_ptr, 0);
1009 let reg_ptr = field_ptr(builder, frame_ptr, pointer_type, OFF_REGISTER_COUNT);
1010 builder
1011 .ins()
1012 .store(MachMemFlags::new(), caller_regs, reg_ptr, 0);
1013
1014 registers.clear();
1015 write_dynamic_exit(
1016 builder,
1017 frame_ptr,
1018 pointer_type,
1019 JIT_CONTINUE,
1020 return_pc,
1021 0,
1022 );
1023 let _ = pc;
1024}
1025
1026fn write_dynamic_exit(
1027 builder: &mut FunctionBuilder<'_>,
1028 frame_ptr: ClifValue,
1029 pointer_type: types::Type,
1030 kind: u32,
1031 pc: ClifValue,
1032 return_reg: u32,
1033) {
1034 write_u32_field(builder, frame_ptr, pointer_type, OFF_EXIT_KIND, kind);
1035 let pc_ptr = field_ptr(builder, frame_ptr, pointer_type, OFF_PC);
1036 builder.ins().store(MachMemFlags::new(), pc, pc_ptr, 0);
1037 write_u32_field(builder, frame_ptr, pointer_type, OFF_RETURN_REG, return_reg);
1038 builder.ins().return_(&[]);
1039}
1040
1041#[cfg(test)]
1042mod tests {
1043 use std::sync::Arc;
1044
1045 use crate::{bytecode::builder::ChunkBuilder, NativeTable, Opcode, Vm};
1046
1047 use super::*;
1048 use crate::jit::dispatch::force_compile;
1049 use crate::jit::trace::{JitContext, TraceKey};
1050
1051 type TestResult = Result<(), Box<dyn std::error::Error>>;
1052
1053 #[test]
1054 fn compiles_load_imm_add_return() -> TestResult {
1055 let mut b = ChunkBuilder::new("jit");
1056 b.begin_function("main", 0, 3);
1057 b.emit_load_imm(0, 41);
1058 b.emit_load_imm(1, 1);
1059 b.emit_binop(Opcode::Add, 2, 0, 1);
1060 b.emit_return(2);
1061 let chunk = Arc::new(b.finish());
1062
1063 let mut ctx = JitContext::new(chunk.clone())?;
1064 let key = TraceKey {
1065 function: 0,
1066 entry_pc: 0,
1067 };
1068 let mut compiler = TraceCompiler::new(ctx.module_mut());
1069 let trace = compiler.compile_trace(&chunk, key)?;
1070 assert!(trace.span.range.contains(&0));
1071 Ok(())
1072 }
1073
1074 #[test]
1075 fn compiles_branch_falsy_taken() -> TestResult {
1076 let mut b = ChunkBuilder::new("jit-branch");
1077 b.begin_function("main", 0, 1);
1078 let done = b.new_label();
1079 let ret = b.new_label();
1080 b.emit_load_imm(0, 0);
1081 b.emit_branch(0, done);
1082 b.emit_load_imm(0, 99);
1083 b.emit_jump(ret);
1084 b.bind_label(done);
1085 b.emit_load_imm(0, 7);
1086 b.bind_label(ret);
1087 b.emit_return(0);
1088 let chunk = Arc::new(b.finish());
1089
1090 let mut ctx = JitContext::new(chunk.clone())?;
1091 let key = TraceKey {
1092 function: 0,
1093 entry_pc: 0,
1094 };
1095 let mut compiler = TraceCompiler::new(ctx.module_mut());
1096 let trace = compiler.compile_trace(&chunk, key)?;
1097 let mut slots = vec![0i64; 1];
1098 let ret = super::super::dispatch::run_compiled_trace_ref(
1099 &trace,
1100 &mut slots,
1101 0,
1102 10_000,
1103 0,
1104 1,
1105 );
1106 assert_eq!(slots[0], 7);
1107 assert!(matches!(
1108 ret.into_reason(),
1109 super::super::exit::ExitReason::Return { .. }
1110 ));
1111 Ok(())
1112 }
1113
1114 #[test]
1115 fn compiles_unconditional_jump() -> TestResult {
1116 let mut b = ChunkBuilder::new("jit-jump");
1117 b.begin_function("main", 0, 1);
1118 let skip = b.new_label();
1119 let ret = b.new_label();
1120 b.emit_load_imm(0, 1);
1121 b.emit_jump(skip);
1122 b.emit_load_imm(0, 99);
1123 b.bind_label(skip);
1124 b.emit_load_imm(0, 41);
1125 b.bind_label(ret);
1126 b.emit_return(0);
1127 let chunk = Arc::new(b.finish());
1128
1129 let mut ctx = JitContext::new(chunk.clone())?;
1130 let key = TraceKey {
1131 function: 0,
1132 entry_pc: 0,
1133 };
1134 let mut compiler = TraceCompiler::new(ctx.module_mut());
1135 let trace = compiler.compile_trace(&chunk, key)?;
1136 let mut slots = vec![0i64; 1];
1137 let ret = super::super::dispatch::run_compiled_trace_ref(
1138 &trace,
1139 &mut slots,
1140 0,
1141 10_000,
1142 0,
1143 1,
1144 );
1145 assert!(matches!(
1146 ret.into_reason(),
1147 super::super::exit::ExitReason::Return { return_reg: 0 }
1148 ));
1149 assert_eq!(slots[0], 41);
1150 Ok(())
1151 }
1152
1153 #[test]
1154 fn compiles_intra_chunk_call() -> TestResult {
1155 let mut b = ChunkBuilder::new("jit-call");
1156 b.begin_function("double", 0, 1);
1157 b.emit_load_imm(0, 21);
1158 b.emit_return(0);
1159 b.begin_function("main", 1, 1);
1160 b.emit_load_imm(0, 0);
1161 b.emit_call(0, 0, 1);
1162 b.emit_return(0);
1163 let chunk = Arc::new(b.finish());
1164
1165 let mut ctx = JitContext::new(chunk.clone())?;
1166 let key = TraceKey {
1167 function: 1,
1168 entry_pc: chunk.functions[1].entry,
1169 };
1170 force_compile(&mut ctx, &chunk, key)?;
1171 let mut vm = Vm::new(chunk, NativeTable::empty(), key.function, &[])?;
1172 let result = super::super::dispatch::run_vm_with_jit(&mut vm, 10_000, &mut ctx);
1173 assert!(matches!(
1174 result,
1175 crate::VmResult::Complete(crate::Value::Int(21))
1176 ));
1177 Ok(())
1178 }
1179}