PkmnLib_rs/src/script_implementations/wasm/export_registry/mod.rs

use anyhow_ext::anyhow;
use std::ffi::{c_char, CStr, CString};
use std::mem::{align_of, forget};
use std::ops::Deref;
use wasmer::{FunctionEnv, FunctionEnvMut, Imports, StoreMut};

use crate::script_implementations::wasm::extern_ref::ExternRef;
use crate::script_implementations::wasm::script_resolver::WebAssemblyEnv;
use crate::static_data::Parameter;
use crate::StringKey;

/// Dynamic data registration
mod dynamic_data;
/// Static data registration
mod static_data;
/// Handling for opaque handles passed to WASM
mod wasm_object;
/// Result types for WASM
mod wasm_result;

#[doc(inline)]
pub(super) use wasm_object::*;
#[doc(inline)]
pub use wasm_result::*;

/// Register a single function.
macro_rules! register_func_with_env {
    ($imports: ident, $store: ident, $func: ident, $env: expr) => {
        $imports.define(
            "env",
            stringify!($func),
            wasmer::Function::new_typed_with_env($store, $env, $func),
        );
    };
}

/// Utility macro to register a bunch of WASM functions easily.
macro_rules! register {
    (
        $(
            fn $name:ident($($par:ident: $par_type:ty),*$(,)?) $(-> $return:ty)? $block:block
        )*
        $(
            manual $manual_name:ident
        )?
    ) => {

        $(
            fn $name(
                $(
                $par: $par_type,
                )*
            ) $(-> $return)* $block
        )*

        pub(crate) fn register(imports: &mut wasmer::Imports,
            store: &mut wasmer::StoreMut,
            env: &wasmer::FunctionEnv<crate::script_implementations::wasm::script_resolver::WebAssemblyEnv>) {
            $(
                crate::script_implementations::wasm::export_registry::register_func_with_env!(imports, store, $name, env);
            )*
            $( $manual_name(imports, store, env) )*
        }
    };
}

pub(crate) use register;
pub(crate) use register_func_with_env;

/// Register the functions we expose to WASM.
pub(crate) fn register_webassembly_funcs(
    imports: &mut Imports,
    store: &mut StoreMut,
    env: &FunctionEnv<WebAssemblyEnv>,
) {
    register_func_with_env!(imports, store, _print, env);
    register_func_with_env!(imports, store, _error, env);

    static_data::register(imports, store, env);
    dynamic_data::register(imports, store, env);

    register_func_with_env!(imports, store, string_key_get_hash, env);
    register_func_with_env!(imports, store, string_key_get_str, env);
    register_func_with_env!(imports, store, effect_parameter_get_type, env);
    register_func_with_env!(imports, store, effect_parameter_as_bool, env);
    register_func_with_env!(imports, store, effect_parameter_as_int, env);
    register_func_with_env!(imports, store, effect_parameter_as_float, env);
    register_func_with_env!(imports, store, effect_parameter_as_string, env);
}

/// Logging function for WASM.
fn _print(env: FunctionEnvMut<WebAssemblyEnv>, p: u32) {
    unsafe {
        let mem: *mut c_char = env.data().data().get_raw_pointer(p);
        let s = CStr::from_ptr(mem);
        if let Ok(v) = s.to_str() {
            println!("{}", v);
        }
    }
}

/// Triggers when WASM panics.
#[track_caller]
fn _error(env: FunctionEnvMut<WebAssemblyEnv>, message: u32) {
    unsafe {
        let mem: *const c_char = env.data().data().get_raw_pointer(message);
        let message_str = CStr::from_ptr(mem);
        #[allow(clippy::panic)]
        if let Ok(v) = message_str.to_str() {
            panic!("WASM Error: {}", v);
        }
    }
}

/// Gets the hash value of a StringKey.
fn string_key_get_hash(env: FunctionEnvMut<WebAssemblyEnv>, string_key: ExternRef<StringKey>) -> WasmResult<u32> {
    let value = get_value!(string_key, env);
    wasm_ok(value.hash())
}

/// Get a null-terminated C string from a StringKey. Note that this involves a copy into WASM
/// memory, so this is relatively heavy.
fn string_key_get_str(env: FunctionEnvMut<WebAssemblyEnv>, string_key: ExternRef<StringKey>) -> WasmResult<u32> {
    let value = get_value!(string_key, env);
    let string_key = value.str();
    let wasm_string_ptr = try_wasm!(
        env.data()
            .data()
            .allocate_mem((string_key.len() + 1) as u32, align_of::<CString>() as u32),
        env
    );
    let mut wasm_string: Vec<u8> =
        unsafe { Vec::from_raw_parts(wasm_string_ptr.0, string_key.len() + 1, string_key.len() + 1) };
    wasm_string.resize(string_key.len() + 1, 0);
    string_key.as_bytes().clone_into(&mut wasm_string);
    wasm_string.insert(string_key.len(), 0_u8);
    forget(wasm_string);
    println!("Returning v: {:?}", wasm_string_ptr.1);
    wasm_ok(wasm_string_ptr.1)
}

/// Gets the type of an EffectParameter
fn effect_parameter_get_type(env: FunctionEnvMut<WebAssemblyEnv>, parameter: ExternRef<Parameter>) -> WasmResult<u8> {
    let value = get_value_arc!(parameter, env);
    wasm_ok(match value.deref() {
        Parameter::Bool(_) => 1,
        Parameter::Int(_) => 2,
        Parameter::Float(_) => 3,
        Parameter::String(_) => 4,
    })
}

/// Gets the inner bool data of an EffectParameter. Panics if it's not a bool.
fn effect_parameter_as_bool(env: FunctionEnvMut<WebAssemblyEnv>, parameter: ExternRef<Parameter>) -> WasmResult<u8> {
    let value = get_value_arc!(parameter, env);
    match value.deref() {
        Parameter::Bool(b) => wasm_ok(<u8 as From<bool>>::from(*b)),
        _ => wasm_err::<u8>(anyhow!("Unexpected parameter type. Expected bool, got {}", value), &env),
    }
}

/// Gets the inner int data of an EffectParameter. Panics if it's not an int.
fn effect_parameter_as_int(env: FunctionEnvMut<WebAssemblyEnv>, parameter: ExternRef<Parameter>) -> WasmResult<i64> {
    let value = get_value_arc!(parameter, env);
    match value.deref() {
        Parameter::Int(i) => wasm_ok(*i),
        _ => wasm_err::<i64>(anyhow!("Unexpected parameter type. Expected int, got {}", value), &env),
    }
}

/// Gets the inner float data of an EffectParameter. Panics if it's not a float.
fn effect_parameter_as_float(env: FunctionEnvMut<WebAssemblyEnv>, parameter: ExternRef<Parameter>) -> WasmResult<f32> {
    let value = get_value_arc!(parameter, env);
    match value.deref() {
        Parameter::Float(f) => wasm_ok(*f),
        _ => wasm_err::<f32>(
            anyhow!("Unexpected parameter type. Expected float, got {}", value),
            &env,
        ),
    }
}

/// Gets the inner string data of an EffectParameter. Panics if it's not a string.
fn effect_parameter_as_string(
    env: FunctionEnvMut<WebAssemblyEnv>,
    parameter: ExternRef<Parameter>,
) -> WasmResult<ExternRef<StringKey>> {
    let value = get_value_arc!(parameter, env);
    match value.deref() {
        Parameter::String(s) => wasm_ok(ExternRef::<StringKey>::func_new(&env, s.clone().into())),
        _ => wasm_err::<ExternRef<StringKey>>(
            anyhow!("Unexpected parameter type. Expected string, got {}", value),
            &env,
        ),
    }
}