ableos-idl/ableos/src/wasm_jumploader/host_functions.rs

use alloc::format;
use wasmi::{
    Error, Externals, FuncInstance, FuncRef, ModuleImportResolver, RuntimeArgs, RuntimeValue,
    Signature, Trap, ValueType,
};

pub struct HostExternals {
    // counter: usize,
}

const ADD_FUNC_INDEX: usize = 0;

impl Externals for HostExternals {
    fn invoke_index(
        &mut self,
        index: usize,
        args: RuntimeArgs,
    ) -> Result<Option<RuntimeValue>, Trap> {
        match index {
            ADD_FUNC_INDEX => {
                let a: u32 = args.nth_checked(0)?;
                let b: u32 = args.nth_checked(1)?;
                let result = a + b;

                Ok(Some(RuntimeValue::I32(result as i32)))
            }
            _ => panic!("Unimplemented function at {}", index),
        }
    }
}

impl HostExternals {
    fn check_signature(&self, index: usize, signature: &Signature) -> bool {
        let (params, ret_ty): (&[ValueType], Option<ValueType>) = match index {
            ADD_FUNC_INDEX => (&[ValueType::I32, ValueType::I32], Some(ValueType::I32)),
            _ => return false,
        };
        signature.params() == params && signature.return_type() == ret_ty
    }
}

impl ModuleImportResolver for HostExternals {
    fn resolve_func(&self, field_name: &str, signature: &Signature) -> Result<FuncRef, Error> {
        let index = match field_name {
            "add" => ADD_FUNC_INDEX,
            _ => {
                return Err(Error::Instantiation(format!(
                    "Export {} not found",
                    field_name
                )))
            }
        };

        if !self.check_signature(index, signature) {
            return Err(Error::Instantiation(format!(
                "Export {} has a bad signature",
                field_name
            )));
        }

        Ok(FuncInstance::alloc_host(
            Signature::new(&[ValueType::I32, ValueType::I32][..], Some(ValueType::I32)),
            index,
        ))
    }
}