PkmnLib_rs/src/script_implementations/wasm/script_function_cache.rs

use std::marker::PhantomData;
use std::sync::Arc;

use parking_lot::RwLock;
use paste::paste;

use crate::dynamic_data::{Battle, DynamicLibrary, ExecutingMove, Pokemon, TurnChoice};
use crate::script_implementations::wasm::extern_ref::{ExternRef, VecExternRef};
use crate::script_implementations::wasm::script_resolver::WebAssemblyEnvironmentData;
use crate::static_data::{EffectParameter, Item, TypeIdentifier};
use crate::StringKey;
use wasmer::{FromToNativeWasmType, TypedFunction};

/// A macro to generate the script function cache a bit easier.
macro_rules! script_function_cache {
    (
        $(
            $name:ident($($par_type:ty),*$(,)?)
        )*
    ) => {
        #[derive(Default)]
        #[allow(unused_parens)]
        pub(super) struct ScriptFunctionCache {
            script_get_name: RwLock<Option<TypedFunction<(u32), u32>>>,
            dealloc_cstring: RwLock<Option<TypedFunction<(u32), ()>>>,
            $(
                $name: RwLock<Option<TypedFunction<(u32$(, $par_type)*), ()>>>,
            )*
        }

        impl ScriptFunctionCache {
            $(
                paste! {
                    #[cold]
                    #[allow(unused_parens)]
                    fn [<initialize_ $name>](&self, env: &Arc<WebAssemblyEnvironmentData>) {
                        let exported = env.exported_functions();
                        let f = exported.get::<StringKey>(&stringify!([< script_ $name >]).into());
                        if let Some(f) = f {
                            let func: TypedFunction<(u32 $(,$par_type)*), ()> = f.typed(&env.store_ref()).unwrap();
                            let _ = self.$name.write().insert(func);
                        }
                    }


                    #[allow(unused_parens)]
                    pub(crate) fn [<$name>](
                        &self,
                        env: &Arc<WebAssemblyEnvironmentData>,
                    ) -> Option<TypedFunction<(u32 $(,$par_type)*), ()>> {
                        {
                            let read_lock = self.$name.read();
                            if let Some(f) = read_lock.as_ref() {
                                return Some(f.clone());
                            }
                        }
                        self.[<initialize_ $name>](env);
                        self.$name.read().as_ref().cloned()
                    }
                }
            )*
        }
    }
}

/// Helper struct to indicate WASM client pointer type. This makes it harder to pass the wrong type
/// to a function by accident.
#[derive(Ord, PartialOrd, Eq, PartialEq)]
pub(super) struct WasmPtr<T> {
    /// The memory size
    v: u32,
    /// Phantom data to store type.
    _phantom: PhantomData<T>,
}

impl<T> Clone for WasmPtr<T> {
    fn clone(&self) -> Self {
        Self {
            v: self.v,
            _phantom: Default::default(),
        }
    }
}

impl<T> Copy for WasmPtr<T> {}

impl<T> From<u32> for WasmPtr<T> {
    fn from(v: u32) -> Self {
        Self {
            v,
            _phantom: Default::default(),
        }
    }
}

impl<T> From<WasmPtr<T>> for u32 {
    fn from(v: WasmPtr<T>) -> Self {
        v.v
    }
}

unsafe impl<T> FromToNativeWasmType for WasmPtr<T> {
    type Native = i32;

    fn from_native(native: Self::Native) -> Self {
        Self {
            v: native as u32,
            _phantom: Default::default(),
        }
    }

    fn to_native(self) -> Self::Native {
        self.v as i32
    }
}

script_function_cache! {
    stack()
    on_remove()
    on_initialize(ExternRef<DynamicLibrary>, VecExternRef<EffectParameter>)
    on_before_turn(ExternRef<TurnChoice>)
    change_speed(ExternRef<TurnChoice>, WasmPtr<u32>)
    change_priority(ExternRef<TurnChoice>, WasmPtr<i8>)
    change_move(ExternRef<TurnChoice>, WasmPtr<ExternRef<StringKey>>)
    change_number_of_hits(ExternRef<TurnChoice>, WasmPtr<u8>)
    prevent_move(ExternRef<ExecutingMove>, WasmPtr<bool>)
    fail_move(ExternRef<ExecutingMove>, WasmPtr<bool>)
    stop_before_move(ExternRef<ExecutingMove>, WasmPtr<bool>)
    on_before_move(ExternRef<ExecutingMove>)
    fail_incoming_move(ExternRef<ExecutingMove>, ExternRef<Pokemon>, WasmPtr<bool>)
    is_invulnerable(ExternRef<ExecutingMove>, ExternRef<Pokemon>, WasmPtr<bool>)
    on_move_miss(ExternRef<ExecutingMove>, ExternRef<Pokemon>)
    change_move_type(ExternRef<ExecutingMove>, ExternRef<Pokemon>, u8, WasmPtr<TypeIdentifier>)
    change_effectiveness(ExternRef<ExecutingMove>, ExternRef<Pokemon>, u8, WasmPtr<f32>)
    block_critical(ExternRef<ExecutingMove>, ExternRef<Pokemon>, u8, WasmPtr<bool>)
    block_incoming_critical(ExternRef<ExecutingMove>, ExternRef<Pokemon>, u8, WasmPtr<bool>)
    change_accuracy(ExternRef<ExecutingMove>, ExternRef<Pokemon>, u8, WasmPtr<u8>)
    change_critical_stage(ExternRef<ExecutingMove>, ExternRef<Pokemon>, u8, WasmPtr<u8>)
    change_critical_modifier(ExternRef<ExecutingMove>, ExternRef<Pokemon>, u8, WasmPtr<f32>)
    change_stab_modifier(ExternRef<ExecutingMove>, ExternRef<Pokemon>, u8, WasmPtr<f32>)
    change_base_power(ExternRef<ExecutingMove>, ExternRef<Pokemon>, u8, WasmPtr<u8>)
    bypass_defensive_stat_boost(ExternRef<ExecutingMove>, ExternRef<Pokemon>, u8, WasmPtr<bool>)
    bypass_offensive_stat_boost(ExternRef<ExecutingMove>, ExternRef<Pokemon>, u8, WasmPtr<bool>)
    change_offensive_stat_value(ExternRef<ExecutingMove>, ExternRef<Pokemon>, u8, WasmPtr<u32>)
    change_defensive_stat_value(ExternRef<ExecutingMove>, ExternRef<Pokemon>, u8, WasmPtr<u32>)
    change_damage_stat_modifier(ExternRef<ExecutingMove>, ExternRef<Pokemon>, u8, WasmPtr<f32>)
    change_damage_modifier(ExternRef<ExecutingMove>, ExternRef<Pokemon>, u8, WasmPtr<f32>)
    change_damage(ExternRef<ExecutingMove>, ExternRef<Pokemon>, u8, WasmPtr<u32>)
    change_incoming_damage(ExternRef<ExecutingMove>, ExternRef<Pokemon>, u8, WasmPtr<u32>)
    on_incoming_hit(ExternRef<ExecutingMove>, ExternRef<Pokemon>, u8)
    on_opponent_faints(ExternRef<ExecutingMove>, ExternRef<Pokemon>, u8)
    prevent_stat_boost_change(ExternRef<Pokemon>, u8, i8, u8, WasmPtr<bool>)
    prevent_secondary_effect(ExternRef<ExecutingMove>, ExternRef<Pokemon>, u8, WasmPtr<bool>)
    change_effect_chance(ExternRef<ExecutingMove>, ExternRef<Pokemon>, u8, WasmPtr<f32>)
    change_incoming_effect_chance(ExternRef<ExecutingMove>, ExternRef<Pokemon>, u8, WasmPtr<f32>)
    on_secondary_effect(ExternRef<ExecutingMove>, ExternRef<Pokemon>, u8)
    on_after_hits(ExternRef<ExecutingMove>, ExternRef<Pokemon>)
    prevent_self_switch(ExternRef<TurnChoice>, WasmPtr<bool>)
    prevent_opponent_switch(ExternRef<TurnChoice>, WasmPtr<bool>)
    on_fail(ExternRef<Pokemon>)
    on_opponent_fail(ExternRef<Pokemon>)
    prevent_self_run_away(ExternRef<TurnChoice>, WasmPtr<bool>)
    prevent_opponent_run_away(ExternRef<TurnChoice>, WasmPtr<bool>)
    on_end_turn()
    on_damage(ExternRef<Pokemon>, u8, u32, u32)
    on_faint(ExternRef<Pokemon>, u8)
    on_switch_in(ExternRef<Pokemon>)
    on_after_held_item_consume(ExternRef<Pokemon>, ExternRef<Item>)
    change_experience_gained(ExternRef<Pokemon>, ExternRef<Pokemon>, WasmPtr<u32>)
    share_experience(ExternRef<Pokemon>, ExternRef<Pokemon>, WasmPtr<bool>)
    block_weather(ExternRef<Battle>, WasmPtr<bool>)
    change_capture_rate_bonus(ExternRef<Pokemon>, ExternRef<Item>, WasmPtr<u8>)
}

impl ScriptFunctionCache {
    pub(crate) fn script_get_name(&self, env: &Arc<WebAssemblyEnvironmentData>) -> Option<TypedFunction<u32, u32>> {
        {
            let read_lock = self.script_get_name.read();
            if let Some(f) = read_lock.as_ref() {
                return Some(f.clone());
            }
        }
        {
            let exported = env.exported_functions();
            let f = exported.get::<StringKey>(&"script_get_name".into());
            if let Some(f) = f {
                let func: TypedFunction<u32, u32> = f.typed(&env.store_ref()).unwrap();
                let _ = self.script_get_name.write().insert(func);
            }
        }
        self.script_get_name.read().as_ref().cloned()
    }

    pub(crate) fn dealloc_cstring(&self, env: &Arc<WebAssemblyEnvironmentData>) -> Option<TypedFunction<u32, ()>> {
        {
            let read_lock = self.dealloc_cstring.read();
            if let Some(f) = read_lock.as_ref() {
                return Some(f.clone());
            }
        }
        {
            let exported = env.exported_functions();
            let f = exported.get::<StringKey>(&"dealloc_cstring".into());
            if let Some(f) = f {
                let func: TypedFunction<u32, ()> = f.typed(&env.store_ref()).unwrap();
                let _ = self.dealloc_cstring.write().insert(func);
            }
        }
        self.dealloc_cstring.read().as_ref().cloned()
    }
}