1use anyhow::{anyhow, Result};
3use core::mem;
4use cranelift_codegen::data_value::DataValue;
5use cranelift_codegen::ir::{
6 ExternalName, Function, InstBuilder, Signature, UserExternalName, UserFuncName,
7};
8use cranelift_codegen::isa::{OwnedTargetIsa, TargetIsa};
9use cranelift_codegen::{ir, settings, CodegenError, Context};
10use cranelift_control::ControlPlane;
11use cranelift_frontend::{FunctionBuilder, FunctionBuilderContext};
12use cranelift_jit::{JITBuilder, JITModule};
13use cranelift_module::{FuncId, Linkage, Module, ModuleError};
14use cranelift_native::builder_with_options;
15use cranelift_reader::TestFile;
16use pulley_interpreter::interp as pulley;
17use std::cmp::max;
18use std::collections::hash_map::Entry;
19use std::collections::HashMap;
20use std::ptr::NonNull;
21use target_lexicon::Architecture;
22use thiserror::Error;
23
24const TESTFILE_NAMESPACE: u32 = 0;
25
26#[derive(Debug)]
28struct DefinedFunction {
29 new_name: UserExternalName,
38
39 signature: ir::Signature,
41
42 func_id: FuncId,
44}
45
46pub struct TestFileCompiler {
74 module: JITModule,
75 ctx: Context,
76
77 defined_functions: HashMap<UserFuncName, DefinedFunction>,
81
82 trampolines: HashMap<Signature, UserFuncName>,
88}
89
90impl TestFileCompiler {
91 pub fn new(isa: OwnedTargetIsa) -> Self {
95 let mut builder = JITBuilder::with_isa(isa, cranelift_module::default_libcall_names());
96 let _ = &mut builder; #[cfg(target_arch = "x86_64")]
98 {
99 builder.symbol_lookup_fn(Box::new(|name| {
100 if name == "__cranelift_x86_pshufb" {
101 Some(__cranelift_x86_pshufb as *const u8)
102 } else {
103 None
104 }
105 }));
106 }
107
108 #[cfg(unix)]
113 {
114 unsafe extern "C" {
115 safe fn ceilf(f: f32) -> f32;
116 }
117 let f = 1.2_f32;
118 assert_eq!(f.ceil(), ceilf(f));
119 }
120
121 let module = JITModule::new(builder);
122 let ctx = module.make_context();
123
124 Self {
125 module,
126 ctx,
127 defined_functions: HashMap::new(),
128 trampolines: HashMap::new(),
129 }
130 }
131
132 pub fn with_host_isa(flags: settings::Flags) -> Result<Self> {
134 let builder =
135 builder_with_options(true).expect("Unable to build a TargetIsa for the current host");
136 let isa = builder.finish(flags)?;
137 Ok(Self::new(isa))
138 }
139
140 pub fn with_default_host_isa() -> Result<Self> {
143 let flags = settings::Flags::new(settings::builder());
144 Self::with_host_isa(flags)
145 }
146
147 pub fn add_functions(
150 &mut self,
151 functions: &[Function],
152 ctrl_planes: Vec<ControlPlane>,
153 ) -> Result<()> {
154 for func in functions {
156 self.declare_function(func)?;
157 }
158
159 let ctrl_planes = ctrl_planes
160 .into_iter()
161 .chain(std::iter::repeat(ControlPlane::default()));
162
163 for (func, ref mut ctrl_plane) in functions.iter().zip(ctrl_planes) {
165 self.define_function(func.clone(), ctrl_plane)?;
166 self.create_trampoline_for_function(func, ctrl_plane)?;
167 }
168
169 Ok(())
170 }
171
172 pub fn add_testfile(&mut self, testfile: &TestFile) -> Result<()> {
175 let functions = testfile
176 .functions
177 .iter()
178 .map(|(f, _)| f)
179 .cloned()
180 .collect::<Vec<_>>();
181
182 self.add_functions(&functions[..], Vec::new())?;
183 Ok(())
184 }
185
186 pub fn declare_function(&mut self, func: &Function) -> Result<()> {
188 let next_id = self.defined_functions.len() as u32;
189 match self.defined_functions.entry(func.name.clone()) {
190 Entry::Occupied(_) => {
191 anyhow::bail!("Duplicate function with name {} found!", &func.name)
192 }
193 Entry::Vacant(v) => {
194 let name = func.name.to_string();
195 let func_id =
196 self.module
197 .declare_function(&name, Linkage::Local, &func.signature)?;
198
199 v.insert(DefinedFunction {
200 new_name: UserExternalName::new(TESTFILE_NAMESPACE, next_id),
201 signature: func.signature.clone(),
202 func_id,
203 });
204 }
205 };
206
207 Ok(())
208 }
209
210 fn apply_func_rename(
215 &self,
216 mut func: Function,
217 defined_func: &DefinedFunction,
218 ) -> Result<Function> {
219 let func_original_name = func.name;
221 func.name = UserFuncName::User(defined_func.new_name.clone());
222
223 let mut redefines = Vec::with_capacity(func.dfg.ext_funcs.len());
226 for (ext_ref, ext_func) in &func.dfg.ext_funcs {
227 let old_name = match &ext_func.name {
228 ExternalName::TestCase(tc) => UserFuncName::Testcase(tc.clone()),
229 ExternalName::User(username) => {
230 UserFuncName::User(func.params.user_named_funcs()[*username].clone())
231 }
232 _ => continue,
234 };
235
236 let target_df = self.defined_functions.get(&old_name).ok_or(anyhow!(
237 "Undeclared function {} is referenced by {}!",
238 &old_name,
239 &func_original_name
240 ))?;
241
242 redefines.push((ext_ref, target_df.new_name.clone()));
243 }
244
245 for (ext_ref, new_name) in redefines.into_iter() {
247 let new_name_ref = func.params.ensure_user_func_name(new_name);
249
250 func.dfg.ext_funcs[ext_ref].name = ExternalName::User(new_name_ref);
252 }
253
254 Ok(func)
255 }
256
257 pub fn define_function(&mut self, func: Function, ctrl_plane: &mut ControlPlane) -> Result<()> {
259 let defined_func = self
260 .defined_functions
261 .get(&func.name)
262 .ok_or(anyhow!("Undeclared function {} found!", &func.name))?;
263
264 self.ctx.func = self.apply_func_rename(func, defined_func)?;
265 self.module.define_function_with_control_plane(
266 defined_func.func_id,
267 &mut self.ctx,
268 ctrl_plane,
269 )?;
270 self.module.clear_context(&mut self.ctx);
271 Ok(())
272 }
273
274 pub fn create_trampoline_for_function(
276 &mut self,
277 func: &Function,
278 ctrl_plane: &mut ControlPlane,
279 ) -> Result<()> {
280 if !self.defined_functions.contains_key(&func.name) {
281 anyhow::bail!("Undeclared function {} found!", &func.name);
282 }
283
284 if self.trampolines.contains_key(&func.signature) {
286 return Ok(());
287 }
288
289 let name = UserFuncName::user(TESTFILE_NAMESPACE, self.defined_functions.len() as u32);
291 let trampoline = make_trampoline(name.clone(), &func.signature, self.module.isa());
292
293 self.declare_function(&trampoline)?;
294 self.define_function(trampoline, ctrl_plane)?;
295
296 self.trampolines.insert(func.signature.clone(), name);
297
298 Ok(())
299 }
300
301 pub fn compile(mut self) -> Result<CompiledTestFile, CompilationError> {
303 self.module.finalize_definitions()?;
307
308 Ok(CompiledTestFile {
309 module: Some(self.module),
310 defined_functions: self.defined_functions,
311 trampolines: self.trampolines,
312 })
313 }
314}
315
316pub struct CompiledTestFile {
318 module: Option<JITModule>,
321
322 defined_functions: HashMap<UserFuncName, DefinedFunction>,
325
326 trampolines: HashMap<Signature, UserFuncName>,
329}
330
331impl CompiledTestFile {
332 pub fn get_trampoline(&self, func: &Function) -> Option<Trampoline> {
336 let defined_func = self.defined_functions.get(&func.name)?;
337 let trampoline_id = self
338 .trampolines
339 .get(&func.signature)
340 .and_then(|name| self.defined_functions.get(name))
341 .map(|df| df.func_id)?;
342 Some(Trampoline {
343 module: self.module.as_ref()?,
344 func_id: defined_func.func_id,
345 func_signature: &defined_func.signature,
346 trampoline_id,
347 })
348 }
349}
350
351impl Drop for CompiledTestFile {
352 fn drop(&mut self) {
353 unsafe { self.module.take().unwrap().free_memory() }
356 }
357}
358
359pub struct Trampoline<'a> {
361 module: &'a JITModule,
362 func_id: FuncId,
363 func_signature: &'a Signature,
364 trampoline_id: FuncId,
365}
366
367impl<'a> Trampoline<'a> {
368 pub fn call(&self, arguments: &[DataValue]) -> Vec<DataValue> {
370 let mut values = UnboxedValues::make_arguments(arguments, &self.func_signature);
371 let arguments_address = values.as_mut_ptr();
372
373 let function_ptr = self.module.get_finalized_function(self.func_id);
374 let trampoline_ptr = self.module.get_finalized_function(self.trampoline_id);
375
376 unsafe {
377 self.call_raw(trampoline_ptr, function_ptr, arguments_address);
378 }
379
380 values.collect_returns(&self.func_signature)
381 }
382
383 unsafe fn call_raw(
384 &self,
385 trampoline_ptr: *const u8,
386 function_ptr: *const u8,
387 arguments_address: *mut u128,
388 ) {
389 match self.module.isa().triple().architecture {
390 Architecture::Pulley32
393 | Architecture::Pulley64
394 | Architecture::Pulley32be
395 | Architecture::Pulley64be => {
396 let mut state = pulley::Vm::new();
397 state.call(
398 NonNull::new(trampoline_ptr.cast_mut()).unwrap(),
399 &[
400 pulley::XRegVal::new_ptr(function_ptr.cast_mut()).into(),
401 pulley::XRegVal::new_ptr(arguments_address).into(),
402 ],
403 [],
404 );
405 }
406
407 _ => {
409 let callable_trampoline: fn(*const u8, *mut u128) -> () =
410 unsafe { mem::transmute(trampoline_ptr) };
411 callable_trampoline(function_ptr, arguments_address);
412 }
413 }
414 }
415}
416
417#[derive(Error, Debug)]
419pub enum CompilationError {
420 #[error("Cranelift codegen error")]
422 CodegenError(#[from] CodegenError),
423 #[error("Module error")]
425 ModuleError(#[from] ModuleError),
426 #[error("Memory mapping error")]
428 IoError(#[from] std::io::Error),
429}
430
431struct UnboxedValues(Vec<u128>);
434
435impl UnboxedValues {
436 const SLOT_SIZE: usize = 16;
441
442 pub fn make_arguments(arguments: &[DataValue], signature: &ir::Signature) -> Self {
445 assert_eq!(arguments.len(), signature.params.len());
446 let mut values_vec = vec![0; max(signature.params.len(), signature.returns.len())];
447
448 for ((arg, slot), param) in arguments.iter().zip(&mut values_vec).zip(&signature.params) {
450 assert!(
451 arg.ty() == param.value_type || arg.is_vector(),
452 "argument type mismatch: {} != {}",
453 arg.ty(),
454 param.value_type
455 );
456 unsafe {
457 arg.write_value_to(slot);
458 }
459 }
460
461 Self(values_vec)
462 }
463
464 pub fn as_mut_ptr(&mut self) -> *mut u128 {
466 self.0.as_mut_ptr()
467 }
468
469 pub fn collect_returns(&self, signature: &ir::Signature) -> Vec<DataValue> {
472 assert!(self.0.len() >= signature.returns.len());
473 let mut returns = Vec::with_capacity(signature.returns.len());
474
475 for (slot, param) in self.0.iter().zip(&signature.returns) {
477 let value = unsafe { DataValue::read_value_from(slot, param.value_type) };
478 returns.push(value);
479 }
480
481 returns
482 }
483}
484
485fn make_trampoline(name: UserFuncName, signature: &ir::Signature, isa: &dyn TargetIsa) -> Function {
491 let pointer_type = isa.pointer_type();
493 let mut wrapper_sig = ir::Signature::new(isa.frontend_config().default_call_conv);
494 wrapper_sig.params.push(ir::AbiParam::new(pointer_type)); wrapper_sig.params.push(ir::AbiParam::new(pointer_type)); let mut func = ir::Function::with_name_signature(name, wrapper_sig);
498
499 let mut builder_context = FunctionBuilderContext::new();
501 let mut builder = FunctionBuilder::new(&mut func, &mut builder_context);
502 let block0 = builder.create_block();
503 builder.append_block_params_for_function_params(block0);
504 builder.switch_to_block(block0);
505 builder.seal_block(block0);
506
507 let (callee_value, values_vec_ptr_val) = {
509 let params = builder.func.dfg.block_params(block0);
510 (params[0], params[1])
511 };
512
513 let callee_args = signature
515 .params
516 .iter()
517 .enumerate()
518 .map(|(i, param)| {
519 let mut flags = ir::MemFlags::trusted();
521 if param.value_type.is_vector() {
522 flags.set_endianness(ir::Endianness::Little);
523 }
524
525 builder.ins().load(
527 param.value_type,
528 flags,
529 values_vec_ptr_val,
530 (i * UnboxedValues::SLOT_SIZE) as i32,
531 )
532 })
533 .collect::<Vec<_>>();
534
535 let new_sig = builder.import_signature(signature.clone());
537 let call = builder
538 .ins()
539 .call_indirect(new_sig, callee_value, &callee_args);
540
541 let results = builder.func.dfg.inst_results(call).to_vec();
543 for ((i, value), param) in results.iter().enumerate().zip(&signature.returns) {
544 let mut flags = ir::MemFlags::trusted();
546 if param.value_type.is_vector() {
547 flags.set_endianness(ir::Endianness::Little);
548 }
549 builder.ins().store(
551 flags,
552 *value,
553 values_vec_ptr_val,
554 (i * UnboxedValues::SLOT_SIZE) as i32,
555 );
556 }
557
558 builder.ins().return_(&[]);
559 builder.finalize();
560
561 func
562}
563
564#[cfg(target_arch = "x86_64")]
565use std::arch::x86_64::__m128i;
566#[cfg(target_arch = "x86_64")]
567#[expect(
568 improper_ctypes_definitions,
569 reason = "manually verified to work for now"
570)]
571extern "C" fn __cranelift_x86_pshufb(a: __m128i, b: __m128i) -> __m128i {
572 union U {
573 reg: __m128i,
574 mem: [u8; 16],
575 }
576
577 unsafe {
578 let a = U { reg: a }.mem;
579 let b = U { reg: b }.mem;
580
581 let select = |arr: &[u8; 16], byte: u8| {
582 if byte & 0x80 != 0 {
583 0x00
584 } else {
585 arr[(byte & 0xf) as usize]
586 }
587 };
588
589 U {
590 mem: [
591 select(&a, b[0]),
592 select(&a, b[1]),
593 select(&a, b[2]),
594 select(&a, b[3]),
595 select(&a, b[4]),
596 select(&a, b[5]),
597 select(&a, b[6]),
598 select(&a, b[7]),
599 select(&a, b[8]),
600 select(&a, b[9]),
601 select(&a, b[10]),
602 select(&a, b[11]),
603 select(&a, b[12]),
604 select(&a, b[13]),
605 select(&a, b[14]),
606 select(&a, b[15]),
607 ],
608 }
609 .reg
610 }
611}
612
613#[cfg(test)]
614mod test {
615 use super::*;
616 use cranelift_reader::{parse_functions, parse_test, ParseOptions};
617
618 fn parse(code: &str) -> Function {
619 parse_functions(code).unwrap().into_iter().nth(0).unwrap()
620 }
621
622 #[test]
623 fn nop() {
624 if cranelift_native::builder().is_err() {
626 return;
627 }
628 let code = String::from(
629 "
630 test run
631 function %test() -> i8 {
632 block0:
633 nop
634 v1 = iconst.i8 -1
635 return v1
636 }",
637 );
638 let ctrl_plane = &mut ControlPlane::default();
639
640 let test_file = parse_test(code.as_str(), ParseOptions::default()).unwrap();
642 assert_eq!(1, test_file.functions.len());
643 let function = test_file.functions[0].0.clone();
644
645 let mut compiler = TestFileCompiler::with_default_host_isa().unwrap();
647 compiler.declare_function(&function).unwrap();
648 compiler
649 .define_function(function.clone(), ctrl_plane)
650 .unwrap();
651 compiler
652 .create_trampoline_for_function(&function, ctrl_plane)
653 .unwrap();
654 let compiled = compiler.compile().unwrap();
655 let trampoline = compiled.get_trampoline(&function).unwrap();
656 let returned = trampoline.call(&[]);
657 assert_eq!(returned, vec![DataValue::I8(-1)])
658 }
659
660 #[test]
661 fn trampolines() {
662 if cranelift_native::builder().is_err() {
664 return;
665 }
666 let function = parse(
667 "
668 function %test(f32, i8, i64x2, i8) -> f32x4, i64 {
669 block0(v0: f32, v1: i8, v2: i64x2, v3: i8):
670 v4 = vconst.f32x4 [0x0.1 0x0.2 0x0.3 0x0.4]
671 v5 = iconst.i64 -1
672 return v4, v5
673 }",
674 );
675
676 let compiler = TestFileCompiler::with_default_host_isa().unwrap();
677 let trampoline = make_trampoline(
678 UserFuncName::user(0, 0),
679 &function.signature,
680 compiler.module.isa(),
681 );
682 println!("{trampoline}");
683 assert!(format!("{trampoline}").ends_with(
684 "sig0 = (f32, i8, i64x2, i8) -> f32x4, i64 fast
685
686block0(v0: i64, v1: i64):
687 v2 = load.f32 notrap aligned v1
688 v3 = load.i8 notrap aligned v1+16
689 v4 = load.i64x2 notrap aligned little v1+32
690 v5 = load.i8 notrap aligned v1+48
691 v6, v7 = call_indirect sig0, v0(v2, v3, v4, v5)
692 store notrap aligned little v6, v1
693 store notrap aligned v7, v1+16
694 return
695}
696"
697 ));
698 }
699}