PkmnLib_rs/src/script_implementations/wasm/script_resolver.rs

use std::any::Any;
use std::fmt::{Debug, Formatter};
use std::mem::{align_of, forget, size_of};
use std::sync::{Arc, Weak};

use hashbrown::{HashMap, HashSet};
use parking_lot::lock_api::RwLockReadGuard;
use parking_lot::{RawRwLock, RwLock};
use wasmer::{
    AsStoreMut, AsStoreRef, EngineBuilder, Extern, Features, Function, FunctionEnv, Imports, Instance, Memory, Module,
    Store, StoreMut, StoreRef, TypedFunction, Value,
};

use crate::dynamic_data::{ItemScript, Script, ScriptOwnerData, ScriptResolver};
use crate::script_implementations::wasm::export_registry::register_webassembly_funcs;
use crate::script_implementations::wasm::extern_ref::ExternRef;
use crate::script_implementations::wasm::script::WebAssemblyScript;
use crate::script_implementations::wasm::script_function_cache::ScriptFunctionCache;
use crate::script_implementations::wasm::temp_wasm_allocator::{AllocatedObject, TempWasmAllocator};
use crate::script_implementations::wasm::WebAssemblyScriptCapabilities;
use crate::static_data::Item;
use crate::{PkmnResult, ScriptCategory, StringKey, ValueIdentifiable, ValueIdentifier};

/// A WebAssembly script resolver implements the dynamic scripts functionality with WebAssembly.
pub struct WebAssemblyScriptResolver {
    /// A unique identifier so we know what value this is.
    identifier: ValueIdentifier,
    /// The global state storage of WASM.
    _store: *mut Store,
    /// The WASM modules we have loaded.
    modules: Vec<Module>,
    /// Our currently loaded WASM instances. Empty until finalize() is called, after which the loaded modules get turned
    /// into actual instances.
    instances: Vec<Instance>,

    /// This is the WASM function to load a script.
    load_script_fn: Option<TypedFunction<(u8, ExternRef<StringKey>), u32>>,

    /// The data for use in the scripting function calls.
    environment_data: Arc<WebAssemblyEnvironmentData>,
}

/// This struct allows us to index a hashmap with both a category and name of a script.
#[derive(Debug, Clone, Eq, PartialEq, Hash)]
struct ScriptCapabilitiesKey {
    /// The category for the script we're looking for capabilities for.
    category: ScriptCategory,
    /// The name of the script we're looking for capabilities for.
    script_key: StringKey,
}

impl WebAssemblyScriptResolver {
    /// Instantiates a new WebAssemblyScriptResolver.
    pub fn new() -> Box<WebAssemblyScriptResolver> {
        let compiler = wasmer::Cranelift::default();
        let mut features = Features::new();
        features.multi_value = true;
        features.reference_types = true;
        let engine = EngineBuilder::new(compiler).set_features(Some(features));
        let mut store = Box::new(Store::new(engine));
        let store_ptr: *mut Store = store.as_mut();
        forget(store);

        let environment = Arc::new(WebAssemblyEnvironmentData::new(store_ptr));
        environment.self_arc.write().replace(Arc::downgrade(&environment));

        let s = Self {
            identifier: Default::default(),
            _store: store_ptr,
            modules: Default::default(),
            instances: Default::default(),
            load_script_fn: None,
            environment_data: environment,
        };

        Box::new(s)
    }

    /// Get an immutable reference to the current WASM Store.
    fn store_ref(&self) -> StoreRef<'_> {
        unsafe { self._store.as_ref().unwrap().as_store_ref() }
    }

    /// Get a mutable reference to the current WASM Store.
    fn store_mut(&self) -> StoreMut<'_> {
        unsafe { self._store.as_mut().unwrap().as_store_mut() }
    }

    /// Load a compiled WASM module.
    pub fn load_wasm_from_bytes(&mut self, bytes: &[u8]) {
        // FIXME: Error handling
        let module = Module::new(&self.store_ref(), bytes).unwrap();
        self.modules.push(module);
    }

    /// Tells the script resolver we're done loading wasm modules, and to finalize the resolver.
    pub fn finalize(&mut self) {
        let mut imports = Imports::new();

        let env = FunctionEnv::new(
            &mut self.store_mut(),
            WebAssemblyEnv::new(Arc::downgrade(&self.environment_data)),
        );
        register_webassembly_funcs(&mut imports, &mut self.store_mut(), &env);
        for module in &self.modules {
            for import in module.imports() {
                if imports.get_export("env", import.name()).is_none() {
                    println!(
                        "\x1b[91mMissing import: \"{}\" with type: {:?} \x1b[0m",
                        import.name(),
                        import.ty()
                    );
                }
            }

            let instance = Instance::new(&mut self.store_mut(), module, &imports).unwrap();
            let exports = &instance.exports;

            let init_fn = exports.get_extern("_init");
            if let Some(Extern::Function(init_fn)) = init_fn {
                init_fn.call(&mut self.store_mut(), &[]).unwrap();
            }

            let mut exported_functions = self.environment_data.exported_functions.write();
            for export in exports.iter() {
                match export.1 {
                    Extern::Function(f) => {
                        exported_functions.insert(export.0.as_str().into(), f.clone());
                    }
                    Extern::Memory(m) => {
                        let _ = self.environment_data.memory.write().insert(m.clone());
                    }
                    _ => {}
                }
            }
            if let Some(m) = &self.environment_data.memory.read().as_ref() {
                m.grow(&mut self.store_mut(), 32).unwrap();
            }
            if let Some(f) = exported_functions.get::<StringKey>(&"load_script".into()) {
                self.load_script_fn = Some(f.typed(&self.store_ref()).unwrap())
            }
            if let Some(f) = exported_functions.get::<StringKey>(&"allocate_mem".into()) {
                let _ = self
                    .environment_data
                    .allocate_mem_fn
                    .write()
                    .insert(f.typed(&self.store_ref()).unwrap());

                let temp_memory_slab = self.environment_data.allocate_mem(128, 1);
                let _ = self
                    .environment_data
                    .temp_allocator
                    .write()
                    .insert(TempWasmAllocator::new(temp_memory_slab.0, temp_memory_slab.1));
            }
            self.instances.push(instance);
        }
    }

    /// Gets the data passed to every function as environment data.
    pub fn environment_data(&self) -> &Arc<WebAssemblyEnvironmentData> {
        &self.environment_data
    }
}

impl ValueIdentifiable for WebAssemblyScriptResolver {
    fn value_identifier(&self) -> ValueIdentifier {
        self.identifier
    }
}

impl ScriptResolver for WebAssemblyScriptResolver {
    fn load_script(
        &self,
        owner: ScriptOwnerData,
        category: ScriptCategory,
        script_key: &StringKey,
    ) -> PkmnResult<Option<Arc<dyn Script>>> {
        let script = self
            .load_script_fn
            .as_ref()
            .unwrap()
            .call(
                &mut self.store_mut(),
                category as u8,
                ExternRef::new_with_resolver(self, script_key),
            )
            .unwrap();
        self.environment_data.setup_script(script, category, script_key, owner)
    }

    fn load_item_script(&self, _key: &dyn Item) -> PkmnResult<Option<Arc<dyn ItemScript>>> {
        todo!()
    }
}

impl Debug for WebAssemblyScriptResolver {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        f.write_str("WebAssemblyScriptResolver")
    }
}

impl Drop for WebAssemblyScriptResolver {
    fn drop(&mut self) {
        unsafe {
            drop(Box::from_raw(self._store));
        }
    }
}

/// This data is what is passed to every function that requires access to the global runtime context.
pub struct WebAssemblyEnvironmentData {
    /// We currently have a hacky implementation of extern refs while we're waiting for ExternRef support to hit the
    /// wasm32-unknown-unknown target of Rust. As we don't want to pass raw memory pointers to WASM for security reasons,
    /// we instead keep track of all the data we've sent to WASM, and pass the ID of that data to WASM. This allows us
    /// to only operate on data we know WASM owns. We currently store this data in this continuous Vec, and give the index
    /// of the data as the ID.
    extern_ref_pointers: RwLock<Vec<*const dyn Any>>,
    /// To make sure we send the same identifier to WASM when we send the same piece of data multiple times, we have a
    /// backwards lookup on extern_ref_pointers. This allows us to get the index for a given piece of data.
    extern_ref_pointers_lookup: RwLock<HashMap<ExternRefLookupKey, usize>>,
    /// As an added security measure on our extern refs, we keep track of the types of the extern ref data we've sent.
    /// This prevents illegal arbitrary memory operations, where we expect type X, but the actual type is Y, which would
    /// allow for modifying memory we might not want to. If we get a type mismatch, we will panic, preventing this.
    extern_ref_type_lookup: RwLock<HashSet<ExternRefLookupKey>>,

    /// Additional security for data slices passed to WASM.
    extern_vec_ref_lookup: RwLock<HashMap<usize, Vec<usize>>>,

    /// The memory inside of the WASM container.
    memory: RwLock<Option<Memory>>,

    /// This is a map of all the functions that WASM gives us.
    exported_functions: RwLock<HashMap<StringKey, Function>>,

    /// A cache of all script functions for faster calls.
    script_function_cache: ScriptFunctionCache,

    /// This is the WASM function to allocate memory inside the WASM container.
    allocate_mem_fn: RwLock<Option<TypedFunction<(u32, u32), u32>>>,

    /// An allocator for quick short lifetime allocations within WASM.
    temp_allocator: RwLock<Option<TempWasmAllocator>>,

    /// The WASM store.
    store: *mut Store,

    /// Script capabilities tell us which functions are implemented on a given script. This allows us to skip unneeded
    /// WASM calls.
    script_capabilities: RwLock<HashMap<ScriptCapabilitiesKey, Arc<HashSet<WebAssemblyScriptCapabilities>>>>,

    /// A weak reference to ourselves.
    self_arc: RwLock<Option<Weak<Self>>>,

    /// A lookup from WASM memory pointer to their actual script wrappers.
    loaded_scripts: RwLock<HashMap<u32, Weak<WebAssemblyScript>>>,
}

/// A quick lookup so we can find the extern ref of the value.
#[allow(clippy::vtable_address_comparisons)]
#[derive(Clone, Eq, PartialEq, Hash)]
struct ExternRefLookupKey {
    /// The raw pointer to the data
    pub ptr: *const dyn Any,
    /// Whether or not the reference is a Vec
    pub is_vec: bool,
}

impl WebAssemblyEnvironmentData {
    /// Instantiates new Environment data with a given store.
    pub(crate) fn new(store: *mut Store) -> Self {
        Self {
            extern_ref_pointers: Default::default(),
            extern_ref_pointers_lookup: Default::default(),
            extern_ref_type_lookup: Default::default(),
            extern_vec_ref_lookup: Default::default(),
            memory: Default::default(),
            exported_functions: Default::default(),
            script_function_cache: Default::default(),
            allocate_mem_fn: Default::default(),
            temp_allocator: Default::default(),
            store,
            script_capabilities: Default::default(),
            self_arc: Default::default(),
            loaded_scripts: Default::default(),
        }
    }

    /// This returns the memory of the WASM container.
    pub fn memory(&self) -> *mut u8 {
        self.memory.read().as_ref().unwrap().view(&self.store_ref()).data_ptr()
    }

    /// Return a pointer to something inside the WASM memory.
    pub fn get_raw_pointer<T>(&self, offset: u32) -> *mut T {
        unsafe {
            self.memory
                .read()
                .as_ref()
                .unwrap()
                .view(&self.store_ref())
                .data_ptr()
                .offset(offset as isize) as *mut T
        }
    }

    /// This returns the functions exported from WASM.
    pub fn exported_functions(&self) -> RwLockReadGuard<'_, RawRwLock, HashMap<StringKey, Function>> {
        self.exported_functions.read()
    }

    ///
    pub(super) fn script_function_cache(&self) -> &ScriptFunctionCache {
        &self.script_function_cache
    }

    /// Allocates memory inside the WASM container with a given size and alignment. This memory is
    /// owned by WASM, and is how we can pass memory references that the host allocated to WASM.
    /// The return is a tuple containing both the actual pointer to the memory (usable by the host),
    /// and the WASM offset to the memory (usable by the client).
    pub fn allocate_mem(&self, size: u32, align: u32) -> (*mut u8, u32) {
        let wasm_ptr = self
            .allocate_mem_fn
            .read()
            .as_ref()
            .unwrap()
            .call(&mut self.store_mut(), size, align)
            .unwrap();
        unsafe { (self.memory().offset(wasm_ptr as isize), wasm_ptr) }
    }

    /// Allocates memory inside the WASM container with a given size and alignment. This memory is
    /// owned by WASM, and is how we can pass memory references that the host allocated to WASM.
    /// The return is a tuple containing both the actual pointer to the memory (usable by the host),
    /// and the WASM offset to the memory (usable by the client).
    pub fn allocate_mem_typed<T>(&self) -> (*mut u8, u32) {
        let wasm_ptr = self
            .allocate_mem_fn
            .read()
            .as_ref()
            .unwrap()
            .call(&mut self.store_mut(), size_of::<T>() as u32, align_of::<T>() as u32)
            .unwrap();
        unsafe { (self.memory().offset(wasm_ptr as isize), wasm_ptr) }
    }

    /// Allocates a raw array in wasm, and copy the values from a given slice to it. Returns the
    /// pointer in wasm memory.
    pub fn copy_value_vec_to_wasm<T>(&self, v: &[T]) -> u32 {
        let wasm_ptr = self
            .allocate_mem_fn
            .read()
            .as_ref()
            .unwrap()
            .call(
                &mut self.store_mut(),
                (size_of::<T>() * v.len()) as u32,
                align_of::<T>() as u32,
            )
            .unwrap();
        unsafe {
            let raw = self.memory().offset(wasm_ptr as isize) as *mut T;
            v.as_ptr().copy_to(raw, v.len());
        }
        wasm_ptr
    }

    /// Allocate a piece of memory inside WASM with a very short lifespan. This is mainly used for
    /// rapid allocation of script function parameters, where WASM needs to write to a specific
    /// pointer.
    pub(super) fn allocate_temp<T>(&self, value: T) -> AllocatedObject<T> {
        self.temp_allocator.read().as_ref().unwrap().alloc::<T>(value)
    }

    /// Get a numeric value from any given value. This is not a true Extern Ref from WASM, as this
    /// is not supported by our current WASM platform (Rust). Instead, this is simply a way to not
    /// have to send arbitrary pointer values back and forth with WASM. Only values WASM can actually
    /// access can be touched through this, and we ensure the value is the correct type. In the future,
    /// when extern refs get actually properly implemented at compile time we might want to get rid
    /// of this code.
    #[inline(always)]
    pub fn get_extern_ref_index<T: ?Sized>(&self, value: &dyn Any) -> usize {
        self.get_extern_ref_from_identifier::<T>(value, false)
    }

    /// Get a numeric value from any given value. This is not a true Extern Ref from WASM, as this
    /// is not supported by our current WASM platform (Rust). Instead, this is simply a way to not
    /// have to send arbitrary pointer values back and forth with WASM. Only values WASM can actually
    /// access can be touched through this, and we ensure the value is the correct type. In the future,
    /// when extern refs get actually properly implemented at compile time we might want to get rid
    /// of this code.
    pub fn get_extern_vec_ref_index<T: 'static>(&self, value: &Vec<T>) -> usize {
        let mut vec = Vec::with_capacity(value.len());
        for v in value {
            vec.push(self.get_extern_ref_index::<T>(v));
        }
        let p = self.get_extern_ref_from_identifier::<Vec<T>>(value, true);
        self.extern_vec_ref_lookup.write().insert(p, vec);
        p
    }

    /// Get an extern ref belonging to a vector we have passed to WASM.
    pub fn get_extern_vec_ref_extern_ref(&self, extern_vec_ref: usize, index: usize) -> usize {
        let r = self.extern_vec_ref_lookup.read();
        let v = r.get(&extern_vec_ref).unwrap();
        v[index]
    }

    /// Gets the extern ref index belonging to a specific pointer. If none exists, this will create
    /// a new one.
    #[inline(always)]
    fn get_extern_ref_from_identifier<T: ?Sized>(&self, value: &dyn Any, is_vec: bool) -> usize {
        if let Some(v) = self.extern_ref_pointers_lookup.read().get(&ExternRefLookupKey {
            ptr: value as *const dyn Any,
            is_vec,
        }) {
            return *v;
        }
        let index = {
            let mut extern_ref_guard = self.extern_ref_pointers.write();
            extern_ref_guard.push(value as *const dyn Any);
            extern_ref_guard.len()
        };
        self.extern_ref_pointers_lookup.write().insert(
            ExternRefLookupKey {
                ptr: value as *const dyn Any,
                is_vec,
            },
            index,
        );
        self.extern_ref_type_lookup.write().insert(ExternRefLookupKey {
            ptr: value as *const dyn Any,
            is_vec,
        });
        index
    }

    /// Gets a value from the extern ref lookup. This turns an earlier registered index back into
    /// its proper value, validates its type, and returns the value.
    pub fn get_extern_ref_value<'a, T: ?Sized>(&self, index: usize) -> &'a dyn Any {
        let read_guard = self.extern_ref_pointers.read();
        let ptr = read_guard.get(index - 1).unwrap();
        if self
            .extern_ref_type_lookup
            .read()
            .get(&ExternRefLookupKey {
                ptr: *ptr,
                is_vec: false,
            })
            .is_none()
        {
            panic!(
                "Extern ref was accessed with wrong type. Requested type {}, but this was not the type the extern ref was stored with.",
                std::any::type_name::<T>()
            );
        }

        unsafe { ptr.as_ref().unwrap() }
    }

    /// The WASM store.
    pub fn store_ref(&self) -> StoreRef<'_> {
        unsafe { self.store.as_ref().unwrap().as_store_ref() }
    }

    /// The mutable WASM store.
    pub fn store_mut(&self) -> StoreMut<'_> {
        unsafe { self.store.as_mut().unwrap().as_store_mut() }
    }

    /// Find a loaded script based on the pointer in WASM memory.
    pub(crate) fn get_loaded_script(&self, wasm_ptr: u32) -> Option<Arc<WebAssemblyScript>> {
        if let Some(script) = self.loaded_scripts.read().get(&wasm_ptr) {
            script.upgrade()
        } else {
            None
        }
    }

    /// Wrap a script pointer in WASM memory into a host managed script.  
    pub fn setup_script(
        &self,
        script_ptr: u32,
        category: ScriptCategory,
        script_key: &StringKey,
        owner: ScriptOwnerData,
    ) -> PkmnResult<Option<Arc<dyn Script>>> {
        if script_ptr == 0 {
            return Ok(None);
        }

        let key = ScriptCapabilitiesKey {
            category,
            script_key: script_key.clone(),
        };

        if !self.script_capabilities.read().contains_key(&key) {
            let mut capabilities = HashSet::new();
            unsafe {
                if let Some(get_cap) = self
                    .exported_functions
                    .read()
                    .get::<StringKey>(&"get_script_capabilities".into())
                {
                    let res = get_cap
                        .call(&mut self.store_mut(), &[Value::I32(script_ptr as i32)])
                        .unwrap();
                    let ptr =
                        (self.memory() as *const WebAssemblyScriptCapabilities).offset(res[0].i32().unwrap() as isize);
                    let length = res[1].i32().unwrap() as usize;
                    for i in 0..length {
                        capabilities.insert(*ptr.add(i));
                    }
                }
            }
            self.script_capabilities
                .write()
                .insert(key.clone(), Arc::new(capabilities));
        }

        let read_guard = self.script_capabilities.read();
        let capabilities = read_guard.get(&key).unwrap();

        let script = Arc::new(WebAssemblyScript::new(
            owner,
            script_ptr,
            capabilities.clone(),
            self.self_arc.read().as_ref().unwrap().upgrade().unwrap(),
            script_key.clone(),
        ));

        self.loaded_scripts.write().insert(script_ptr, Arc::downgrade(&script));
        Ok(Some(script))
    }
}

/// The runtime environment for script execution. This is passed to most of the host functions being called.
#[derive(Clone)]
pub(crate) struct WebAssemblyEnv {
    /// A pointer to the WebAssemblyScriptResolver belonging to the current script environment.
    data: Weak<WebAssemblyEnvironmentData>,
}

impl WebAssemblyEnv {
    /// Instantiates a new Environment with the requested data.
    pub fn new(data: Weak<WebAssemblyEnvironmentData>) -> Self {
        Self { data }
    }

    /// Get the actual data belonging to the current context.
    pub fn data(&self) -> Arc<WebAssemblyEnvironmentData> {
        self.data.upgrade().unwrap()
    }
}

unsafe impl Sync for WebAssemblyEnv {}

unsafe impl Send for WebAssemblyEnv {}

unsafe impl Sync for WebAssemblyEnvironmentData {}

unsafe impl Send for WebAssemblyEnvironmentData {}