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