Load & Store slice cranelift

main.rs

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

use cranelift::{
    codegen::{
        self,
        ir::{immediates::Offset32, Inst},
    },
    frontend::{FunctionBuilder, FunctionBuilderContext},
    prelude::{types, AbiParam, InstBuilder, MemFlags, Type, Value},
};
use cranelift_jit::{JITBuilder, JITModule};
use cranelift_module::{Linkage, Module};

fn array_load(src_ptr: Value, index: u32, ty: Type, builder: &mut FunctionBuilder) -> Value {
    builder.ins().load(
        ty,
        MemFlags::trusted(),
        src_ptr,
        Offset32::new((index * ty.bytes()) as i32),
    )
}

fn array_set(dest_ptr: Value, new_val: Value, index: u32,ty: Type, builder: &mut FunctionBuilder) -> Inst {
    builder.ins().store(
        MemFlags::trusted(),
        new_val,
        dest_ptr,
        Offset32::new((index * ty.bytes()) as i32),
    )
}

fn main() {
    let builder = JITBuilder::new(cranelift_module::default_libcall_names());
    let mut builder_context = FunctionBuilderContext::new();
    let mut module = JITModule::new(builder);
    let mut ctx = module.make_context();
    let ptr_ty = module.target_config().pointer_type();

    let mut fn_sig = module.make_signature();
    fn_sig
        .params
        .extend([AbiParam::new(types::F64), AbiParam::new(ptr_ty)]);
    ctx.func.signature = fn_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);
    let mult_val_ptr = builder.block_params(entry_block)[0];
    let array_ptr = builder.block_params(entry_block)[1];

    for i in 0..4 {
        let val = array_load(array_ptr, i, types::F64, &mut builder);
        let new_val = builder.ins().fmul(mult_val_ptr, val);
        array_set(array_ptr, new_val, i, types::F64, &mut builder);
    }

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

    println!("{}", builder.func.display(None));

    builder.finalize();

    let test_fn_id = module
        .declare_function("test_fn", Linkage::Export, &ctx.func.signature)
        .unwrap();

    module
        .define_function(
            test_fn_id,
            &mut ctx,
            &mut codegen::binemit::NullTrapSink {},
            &mut codegen::binemit::NullStackMapSink {},
        )
        .unwrap();
    module.finalize_definitions();
    let func: *const u8 = module.get_finalized_function(test_fn_id);

    //Heap Allocated
    let array = Box::new([1.0f64, 2.0, 3.0, 4.0]);
    println!("before {:?}", array);
    let array_ptr = Box::into_raw(array) as *mut f64;
    let result: &[f64] = unsafe {
        let func = std::mem::transmute::<_, unsafe extern "C" fn(f64, *mut f64) -> ()>(func);
        func(5.0, array_ptr);
        std::slice::from_raw_parts(array_ptr, 4)
    };
    println!("after {:?}", result);

    //Stack Allocated
    let mut array = [1.0f64, 2.0, 3.0, 4.0];
    println!("before {:?}", array);
    unsafe {
        let func = std::mem::transmute::<_, unsafe extern "C" fn(f64, &mut [f64; 4]) -> ()>(func);
        func(5.0, &mut array);
    };
    println!("after {:?}", array);
}