Return a slice from a function defined in cranelift that calls a rust function that returns a slice

main.rs

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

use cranelift::{
    codegen::{ir::immediates::Offset32, ir::ArgumentPurpose, 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
extern "C" fn mult(a: f32) -> [f32; 4] {
    [1.0 * a, 2.0 * a, 3.0 * a, 4.0 * a]
}

// 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 function will take a special abi param StructReturn as its first arg
    mult_sig.params.push(AbiParam::special(
        module.target_config().pointer_type(),
        ArgumentPurpose::StructReturn,
    ));
    mult_sig.params.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 ptr_ty = module.target_config().pointer_type();
    let mut testbed_sig = module.make_signature();

    //testbed function will take a special abi param StructReturn as its first arg
    testbed_sig
        .params
        .push(AbiParam::special(ptr_ty, ArgumentPurpose::StructReturn));

    ctx.func.signature = testbed_sig;

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

    builder.append_block_params_for_function_params(entry_block);
    builder.switch_to_block(entry_block);
    builder.seal_block(entry_block);

    // get return_ptr that was passed in as the first special arg StructReturn
    let return_ptr = builder.block_params(entry_block)[0];

    //declare mult function
    let mult_func = module.declare_func_in_func(mult_func, &mut builder.func);

    //setup StackSlotData to be used as a StructReturnSlot
    let stack_slot = builder.create_stack_slot(StackSlotData::new(
        StackSlotKind::StructReturnSlot,
        types::F32.bytes() * 4,
    ));

    //get stack address of StackSlotData
    let stack_slot_address = builder
        .ins()
        .stack_addr(ptr_ty, stack_slot, Offset32::new(0));

    let arg_a = builder.ins().f32const(5.0);

    //call rust function, passing in return address and arg
    let call = builder.ins().call(mult_func, &[stack_slot_address, arg_a]);
    let _ = builder.inst_results(call);

    let size = builder
        .ins()
        .iconst(types::I64, (types::F32.bytes() * 4) as i64);
    //copy returned array from rust function to return_ptr of testbed function
    builder.call_memcpy(module.target_config(), return_ptr, stack_slot_address, size);

    builder.ins().return_(&[]);
    builder.finalize();

    module
        .declare_function("test_fn", Linkage::Local, &ctx.func.signature)
        .unwrap()
}

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 array = unsafe {
        let func = mem::transmute::<_, unsafe extern "C" fn() -> [f32; 4]>(testbed);
        func()
    };

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