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::Delegate
683 | Opcode::Trap
684 | Opcode::Halt
685 )
686}
687
688fn field_ptr(
689 builder: &mut FunctionBuilder<'_>,
690 base: ClifValue,
691 pointer_type: types::Type,
692 offset: i32,
693) -> ClifValue {
694 let off = builder.ins().iconst(pointer_type, i64::from(offset));
695 builder.ins().iadd(base, off)
696}
697
698fn store_slot(
699 builder: &mut FunctionBuilder<'_>,
700 slots: ClifValue,
701 pointer_type: types::Type,
702 reg: u8,
703 value: ClifValue,
704) {
705 let offset = builder.ins().iconst(pointer_type, i64::from(u32::from(reg) * 8));
706 let addr = builder.ins().iadd(slots, offset);
707 builder.ins().store(MachMemFlags::new(), value, addr, 0);
708}
709
710fn bool_as_i64(builder: &mut FunctionBuilder<'_>, cond: ClifValue) -> ClifValue {
711 let one = builder.ins().iconst(types::I64, 1);
712 let zero = builder.ins().iconst(types::I64, 0);
713 builder.ins().select(cond, one, zero)
714}
715
716fn emit_budget_tick(
717 builder: &mut FunctionBuilder<'_>,
718 budget_ptr: ClifValue,
719 frame_ptr: ClifValue,
720 pointer_type: types::Type,
721 pc: u32,
722) {
723 let budget = builder.ins().load(types::I32, MachMemFlags::new(), budget_ptr, 0);
724 let zero = builder.ins().iconst(types::I32, 0);
725 let exhausted = builder.ins().icmp(
726 cranelift_codegen::ir::condcodes::IntCC::Equal,
727 budget,
728 zero,
729 );
730 let continue_insn = builder.create_block();
731 let budget_exit = builder.create_block();
732 builder.ins().brif(exhausted, budget_exit, &[], continue_insn, &[]);
733 builder.switch_to_block(budget_exit);
734 builder.seal_block(budget_exit);
735 write_exit(builder, frame_ptr, pointer_type, JIT_BUDGET, pc, 0);
736 builder.switch_to_block(continue_insn);
737 builder.seal_block(continue_insn);
738 let one = builder.ins().iconst(types::I32, 1);
739 let new_budget = builder.ins().isub(budget, one);
740 builder.ins().store(MachMemFlags::new(), new_budget, budget_ptr, 0);
741}
742
743fn write_u32_field(
744 builder: &mut FunctionBuilder<'_>,
745 frame_ptr: ClifValue,
746 pointer_type: types::Type,
747 offset: i32,
748 value: u32,
749) {
750 let addr = field_ptr(builder, frame_ptr, pointer_type, offset);
751 let val = builder.ins().iconst(types::I32, i64::from(value));
752 builder.ins().store(MachMemFlags::new(), val, addr, 0);
753}
754
755fn write_exit(
756 builder: &mut FunctionBuilder<'_>,
757 frame_ptr: ClifValue,
758 pointer_type: types::Type,
759 kind: u32,
760 pc: u32,
761 return_reg: u32,
762) {
763 write_u32_field(builder, frame_ptr, pointer_type, OFF_EXIT_KIND, kind);
764 write_u32_field(builder, frame_ptr, pointer_type, OFF_PC, pc);
765 write_u32_field(builder, frame_ptr, pointer_type, OFF_RETURN_REG, return_reg);
766 builder.ins().return_(&[]);
767}
768
769#[allow(clippy::too_many_arguments)]
770fn emit_call(
771 builder: &mut FunctionBuilder<'_>,
772 chunk: &Chunk,
773 registers: &mut RegisterMap,
774 slots_ptr: ClifValue,
775 frame_ptr: ClifValue,
776 pointer_type: types::Type,
777 pc_to_block: &HashMap<u32, Block>,
778 call_depth_ptr: ClifValue,
779 call_stack_ptr: ClifValue,
780 pc: u32,
781 instr: Instruction,
782) -> Result<(), CompileError> {
783 let callee = instr.imm as u32;
784 let argc = instr.b;
785 let dst = instr.a;
786 let def = chunk.function(callee).ok_or(CompileError::UnsupportedOpcode {
787 opcode: Opcode::Call,
788 pc,
789 })?;
790 let callee_entry = def.entry;
791 let Some(callee_block) = pc_to_block.get(&callee_entry) else {
792 flush_registers(builder, registers, slots_ptr, pointer_type);
793 write_exit(
794 builder,
795 frame_ptr,
796 pointer_type,
797 JIT_CONTINUE,
798 pc,
799 0,
800 );
801 return Ok(());
802 };
803
804 flush_registers(builder, registers, slots_ptr, pointer_type);
805
806 let depth = builder
807 .ins()
808 .load(types::I32, MachMemFlags::new(), call_depth_ptr, 0);
809 let max_depth = builder
810 .ins()
811 .iconst(types::I32, i64::from(MAX_JIT_CALL_DEPTH as u32));
812 let too_deep = builder.ins().icmp(
813 cranelift_codegen::ir::condcodes::IntCC::SignedGreaterThanOrEqual,
814 depth,
815 max_depth,
816 );
817 let effect_exit = builder.create_block();
818 let call_body = builder.create_block();
819 builder
820 .ins()
821 .brif(too_deep, effect_exit, &[], call_body, &[]);
822 builder.switch_to_block(effect_exit);
823 builder.seal_block(effect_exit);
824 write_exit(
825 builder,
826 frame_ptr,
827 pointer_type,
828 super::exit::JIT_EFFECT,
829 pc,
830 0,
831 );
832 builder.switch_to_block(call_body);
833 builder.seal_block(call_body);
834
835 let record_bytes = builder.ins().iconst(
836 types::I32,
837 i64::from(std::mem::size_of::<super::frame::JitCallRecord>() as u32),
838 );
839 let record_off = builder.ins().imul(depth, record_bytes);
840 let record_off_ptr = builder.ins().uextend(pointer_type, record_off);
841 let record_addr = builder.ins().iadd(call_stack_ptr, record_off_ptr);
842
843 let return_pc = builder.ins().iconst(types::I32, i64::from(pc + 1));
844 builder
845 .ins()
846 .store(MachMemFlags::new(), return_pc, record_addr, 0);
847
848 let func_ptr = field_ptr(builder, frame_ptr, pointer_type, OFF_FUNCTION);
849 let caller_fn = builder
850 .ins()
851 .load(types::I32, MachMemFlags::new(), func_ptr, 0);
852 let caller_fn_addr = {
853 let off = builder.ins().iconst(pointer_type, 4);
854 builder.ins().iadd(record_addr, off)
855 };
856 builder
857 .ins()
858 .store(MachMemFlags::new(), caller_fn, caller_fn_addr, 0);
859
860 let dest = builder.ins().iconst(types::I32, i64::from(u32::from(dst)));
861 let dest_addr = {
862 let off = builder.ins().iconst(pointer_type, 8);
863 builder.ins().iadd(record_addr, off)
864 };
865 builder
866 .ins()
867 .store(MachMemFlags::new(), dest, dest_addr, 0);
868
869 let reg_ptr = field_ptr(builder, frame_ptr, pointer_type, OFF_REGISTER_COUNT);
870 let caller_regs = builder
871 .ins()
872 .load(types::I32, MachMemFlags::new(), reg_ptr, 0);
873 let caller_regs_addr = {
874 let off = builder.ins().iconst(pointer_type, 12);
875 builder.ins().iadd(record_addr, off)
876 };
877 builder
878 .ins()
879 .store(MachMemFlags::new(), caller_regs, caller_regs_addr, 0);
880
881 let one = builder.ins().iconst(types::I32, 1);
882 let new_depth = builder.ins().iadd(depth, one);
883 builder
884 .ins()
885 .store(MachMemFlags::new(), new_depth, call_depth_ptr, 0);
886
887 for i in 0..argc {
888 let src_reg = u32::from(dst) + u32::from(i);
889 let src_offset = builder.ins().iconst(pointer_type, i64::from(src_reg * 8));
890 let src_addr = builder.ins().iadd(slots_ptr, src_offset);
891 let value = builder
892 .ins()
893 .load(types::I64, MachMemFlags::new(), src_addr, 0);
894 let dst_offset = builder.ins().iconst(pointer_type, i64::from(u32::from(i) * 8));
895 let dst_addr = builder.ins().iadd(slots_ptr, dst_offset);
896 builder
897 .ins()
898 .store(MachMemFlags::new(), value, dst_addr, 0);
899 }
900
901 let callee_fn = builder.ins().iconst(types::I32, i64::from(callee));
902 builder
903 .ins()
904 .store(MachMemFlags::new(), callee_fn, func_ptr, 0);
905 let callee_regs = builder.ins().iconst(types::I32, i64::from(def.num_registers));
906 builder
907 .ins()
908 .store(MachMemFlags::new(), callee_regs, reg_ptr, 0);
909
910 registers.clear();
911 builder.ins().jump(*callee_block, &[]);
912 Ok(())
913}
914
915#[allow(clippy::too_many_arguments)]
916fn emit_return(
917 builder: &mut FunctionBuilder<'_>,
918 frame_ptr: ClifValue,
919 pointer_type: types::Type,
920 call_depth_ptr: ClifValue,
921 call_stack_ptr: ClifValue,
922 slots_ptr: ClifValue,
923 registers: &mut RegisterMap,
924 return_reg: u8,
925 pc: u32,
926) {
927 let depth = builder
928 .ins()
929 .load(types::I32, MachMemFlags::new(), call_depth_ptr, 0);
930 let zero = builder.ins().iconst(types::I32, 0);
931 let is_outer = builder.ins().icmp(
932 cranelift_codegen::ir::condcodes::IntCC::Equal,
933 depth,
934 zero,
935 );
936 let outer = builder.create_block();
937 let inner = builder.create_block();
938 builder
939 .ins()
940 .brif(is_outer, outer, &[], inner, &[]);
941 builder.switch_to_block(outer);
942 builder.seal_block(outer);
943 write_exit(
944 builder,
945 frame_ptr,
946 pointer_type,
947 JIT_RETURN,
948 pc,
949 u32::from(return_reg),
950 );
951 builder.switch_to_block(inner);
952 builder.seal_block(inner);
953
954 let one = builder.ins().iconst(types::I32, 1);
955 let new_depth = builder.ins().isub(depth, one);
956 builder
957 .ins()
958 .store(MachMemFlags::new(), new_depth, call_depth_ptr, 0);
959
960 let record_bytes = builder.ins().iconst(
961 types::I32,
962 i64::from(std::mem::size_of::<super::frame::JitCallRecord>() as u32),
963 );
964 let record_off = builder.ins().imul(new_depth, record_bytes);
965 let record_off_ptr = builder.ins().uextend(pointer_type, record_off);
966 let record_addr = builder.ins().iadd(call_stack_ptr, record_off_ptr);
967
968 let return_pc = builder
969 .ins()
970 .load(types::I32, MachMemFlags::new(), record_addr, 0);
971 let return_fn_addr = {
972 let off = builder.ins().iconst(pointer_type, 4);
973 builder.ins().iadd(record_addr, off)
974 };
975 let return_fn = builder
976 .ins()
977 .load(types::I32, MachMemFlags::new(), return_fn_addr, 0);
978 let dest_addr = {
979 let off = builder.ins().iconst(pointer_type, 8);
980 builder.ins().iadd(record_addr, off)
981 };
982 let dest_reg = builder
983 .ins()
984 .load(types::I32, MachMemFlags::new(), dest_addr, 0);
985 let caller_regs_addr = {
986 let off = builder.ins().iconst(pointer_type, 12);
987 builder.ins().iadd(record_addr, off)
988 };
989 let caller_regs = builder
990 .ins()
991 .load(types::I32, MachMemFlags::new(), caller_regs_addr, 0);
992
993 let ret_offset = builder.ins().iconst(pointer_type, i64::from(u32::from(return_reg) * 8));
994 let ret_addr = builder.ins().iadd(slots_ptr, ret_offset);
995 let ret_val = builder
996 .ins()
997 .load(types::I64, MachMemFlags::new(), ret_addr, 0);
998 let eight = builder.ins().iconst(types::I32, 8);
999 let dest_offset = builder.ins().imul(dest_reg, eight);
1000 let dest_offset_ptr = builder.ins().uextend(pointer_type, dest_offset);
1001 let dest_slot = builder.ins().iadd(slots_ptr, dest_offset_ptr);
1002 builder
1003 .ins()
1004 .store(MachMemFlags::new(), ret_val, dest_slot, 0);
1005
1006 let func_ptr = field_ptr(builder, frame_ptr, pointer_type, OFF_FUNCTION);
1007 builder
1008 .ins()
1009 .store(MachMemFlags::new(), return_fn, func_ptr, 0);
1010 let reg_ptr = field_ptr(builder, frame_ptr, pointer_type, OFF_REGISTER_COUNT);
1011 builder
1012 .ins()
1013 .store(MachMemFlags::new(), caller_regs, reg_ptr, 0);
1014
1015 registers.clear();
1016 write_dynamic_exit(
1017 builder,
1018 frame_ptr,
1019 pointer_type,
1020 JIT_CONTINUE,
1021 return_pc,
1022 0,
1023 );
1024 let _ = pc;
1025}
1026
1027fn write_dynamic_exit(
1028 builder: &mut FunctionBuilder<'_>,
1029 frame_ptr: ClifValue,
1030 pointer_type: types::Type,
1031 kind: u32,
1032 pc: ClifValue,
1033 return_reg: u32,
1034) {
1035 write_u32_field(builder, frame_ptr, pointer_type, OFF_EXIT_KIND, kind);
1036 let pc_ptr = field_ptr(builder, frame_ptr, pointer_type, OFF_PC);
1037 builder.ins().store(MachMemFlags::new(), pc, pc_ptr, 0);
1038 write_u32_field(builder, frame_ptr, pointer_type, OFF_RETURN_REG, return_reg);
1039 builder.ins().return_(&[]);
1040}
1041
1042#[cfg(test)]
1043mod tests {
1044 use std::sync::Arc;
1045
1046 use crate::{bytecode::builder::ChunkBuilder, NativeTable, Opcode, Vm};
1047
1048 use super::*;
1049 use crate::jit::dispatch::force_compile;
1050 use crate::jit::trace::{JitContext, TraceKey};
1051
1052 type TestResult = Result<(), Box<dyn std::error::Error>>;
1053
1054 #[test]
1055 fn compiles_load_imm_add_return() -> TestResult {
1056 let mut b = ChunkBuilder::new("jit");
1057 b.begin_function("main", 0, 3);
1058 b.emit_load_imm(0, 41);
1059 b.emit_load_imm(1, 1);
1060 b.emit_binop(Opcode::Add, 2, 0, 1);
1061 b.emit_return(2);
1062 let chunk = Arc::new(b.finish());
1063
1064 let mut ctx = JitContext::new(chunk.clone())?;
1065 let key = TraceKey {
1066 function: 0,
1067 entry_pc: 0,
1068 };
1069 let mut compiler = TraceCompiler::new(ctx.module_mut());
1070 let trace = compiler.compile_trace(&chunk, key)?;
1071 assert!(trace.span.range.contains(&0));
1072 Ok(())
1073 }
1074
1075 #[test]
1076 fn compiles_branch_falsy_taken() -> TestResult {
1077 let mut b = ChunkBuilder::new("jit-branch");
1078 b.begin_function("main", 0, 1);
1079 let done = b.new_label();
1080 let ret = b.new_label();
1081 b.emit_load_imm(0, 0);
1082 b.emit_branch(0, done);
1083 b.emit_load_imm(0, 99);
1084 b.emit_jump(ret);
1085 b.bind_label(done);
1086 b.emit_load_imm(0, 7);
1087 b.bind_label(ret);
1088 b.emit_return(0);
1089 let chunk = Arc::new(b.finish());
1090
1091 let mut ctx = JitContext::new(chunk.clone())?;
1092 let key = TraceKey {
1093 function: 0,
1094 entry_pc: 0,
1095 };
1096 let mut compiler = TraceCompiler::new(ctx.module_mut());
1097 let trace = compiler.compile_trace(&chunk, key)?;
1098 let mut slots = vec![0i64; 1];
1099 let ret = super::super::dispatch::run_compiled_trace_ref(
1100 &trace,
1101 &mut slots,
1102 0,
1103 10_000,
1104 0,
1105 1,
1106 );
1107 assert_eq!(slots[0], 7);
1108 assert!(matches!(
1109 ret.into_reason(),
1110 super::super::exit::ExitReason::Return { .. }
1111 ));
1112 Ok(())
1113 }
1114
1115 #[test]
1116 fn compiles_unconditional_jump() -> TestResult {
1117 let mut b = ChunkBuilder::new("jit-jump");
1118 b.begin_function("main", 0, 1);
1119 let skip = b.new_label();
1120 let ret = b.new_label();
1121 b.emit_load_imm(0, 1);
1122 b.emit_jump(skip);
1123 b.emit_load_imm(0, 99);
1124 b.bind_label(skip);
1125 b.emit_load_imm(0, 41);
1126 b.bind_label(ret);
1127 b.emit_return(0);
1128 let chunk = Arc::new(b.finish());
1129
1130 let mut ctx = JitContext::new(chunk.clone())?;
1131 let key = TraceKey {
1132 function: 0,
1133 entry_pc: 0,
1134 };
1135 let mut compiler = TraceCompiler::new(ctx.module_mut());
1136 let trace = compiler.compile_trace(&chunk, key)?;
1137 let mut slots = vec![0i64; 1];
1138 let ret = super::super::dispatch::run_compiled_trace_ref(
1139 &trace,
1140 &mut slots,
1141 0,
1142 10_000,
1143 0,
1144 1,
1145 );
1146 assert!(matches!(
1147 ret.into_reason(),
1148 super::super::exit::ExitReason::Return { return_reg: 0 }
1149 ));
1150 assert_eq!(slots[0], 41);
1151 Ok(())
1152 }
1153
1154 #[test]
1155 fn compiles_intra_chunk_call() -> TestResult {
1156 let mut b = ChunkBuilder::new("jit-call");
1157 b.begin_function("double", 0, 1);
1158 b.emit_load_imm(0, 21);
1159 b.emit_return(0);
1160 b.begin_function("main", 1, 1);
1161 b.emit_load_imm(0, 0);
1162 b.emit_call(0, 0, 1);
1163 b.emit_return(0);
1164 let chunk = Arc::new(b.finish());
1165
1166 let mut ctx = JitContext::new(chunk.clone())?;
1167 let key = TraceKey {
1168 function: 1,
1169 entry_pc: chunk.functions[1].entry,
1170 };
1171 force_compile(&mut ctx, &chunk, key)?;
1172 let mut vm = Vm::new(chunk, NativeTable::empty(), key.function, &[])?;
1173 let result = super::super::dispatch::run_vm_with_jit(&mut vm, 10_000, &mut ctx);
1174 assert!(matches!(
1175 result,
1176 crate::VmResult::Complete(crate::Value::Int(21))
1177 ));
1178 Ok(())
1179 }
1180}