Setup cranelift to build a function that calls a rust function

main.rs

/*
[dependencies]
cranelift = "0.76.0"
cranelift-module = "0.76.0"
cranelift-jit = "0.76.0"
*/

use cranelift::{codegen::Context, prelude::*};
use cranelift_jit::{JITBuilder, JITModule};
use cranelift_module::{FuncId, Linkage, Module};
use std::mem;

// Rust function that we will call from cranelift jit
fn mult(a: f32, b: f32) -> f32 {
    a * b
}

// Declare mult function in jit and get cranelift callable funcid
fn mult_func(module: &mut JITModule) -> FuncId {
    let mut mult_sig = module.make_signature();
    mult_sig.params.push(AbiParam::new(types::F32));
    mult_sig.params.push(AbiParam::new(types::F32));
    mult_sig.returns.push(AbiParam::new(types::F32));

    module
        .declare_function("mult", Linkage::Import, &mult_sig)
        .unwrap()
}

// Setup a testbed function to call mult_func from
fn build_testbed_func(mult_func: FuncId, ctx: &mut Context, module: &mut JITModule) -> FuncId {
    //Cranelift uses FunctionBuilderContext to reuse dynamic allocations between compiling multiple functions.
    let mut builder_context = FunctionBuilderContext::new();

    let mut testbed_sig = module.make_signature();
    testbed_sig.returns.push(AbiParam::new(types::F32));

    let testbed_func = module
        .declare_function("b", Linkage::Local, &testbed_sig)
        .unwrap();

    ctx.func.signature = testbed_sig;
    ctx.func.name = ExternalName::user(0, testbed_func.as_u32());

    let mut builder = FunctionBuilder::new(&mut ctx.func, &mut builder_context);
    let entry_block = builder.create_block();

    builder.switch_to_block(entry_block);
    let local_func = module.declare_func_in_func(mult_func, &mut builder.func);
    let arg_a = builder.ins().f32const(5.5);
    let arg_b = builder.ins().f32const(6.6);
    let call = builder.ins().call(local_func, &[arg_a, arg_b]);
    let value = {
        let results = builder.inst_results(call);
        assert_eq!(results.len(), 1);
        results[0].clone()
    };
    builder.ins().return_(&[value]);
    builder.seal_all_blocks();
    builder.finalize();
    testbed_func
}

fn main() {
    let mut builder = JITBuilder::new(cranelift_module::default_libcall_names());

    // Add symbol for mult function to builder
    builder.symbol("mult", mult as *const u8);

    // The Module holds information about all functions and data objects defined in the current JIT
    let mut module = JITModule::new(builder);

    // This is the main Context object for compiling functions.
    let mut ctx = module.make_context();

    // Declare mult function in jit and get cranelift callable funcid
    let mult_func = mult_func(&mut module);

    // Setup a testbed function to call mult_func from
    let testbed_func = build_testbed_func(mult_func, &mut ctx, &mut module);

    // Define testbed function to jit. This finishes compilation.
    module
        .define_function(
            testbed_func,
            &mut ctx,
            &mut codegen::binemit::NullTrapSink {},
            &mut codegen::binemit::NullStackMapSink {},
        )
        .unwrap();

    // Now that compilation is finished, we can clear out the context state.
    module.clear_context(&mut ctx);

    // Perform linking.
    module.finalize_definitions();

    // Get a raw pointer to the generated code.
    let testbed = module.get_finalized_function(testbed_func);

    // Cast testbed function to a rust function pointer type.
    let testbed_ptr = unsafe { mem::transmute::<_, fn() -> f32>(testbed) };

    // Call testbed function
    println!("result: {}", testbed_ptr());
}