PkmnLib_rs/src/script_implementations/wasm/script_resolver.rs

use std::fmt::{Debug, Formatter};
use std::sync::Arc;

use hashbrown::{HashMap, HashSet};
use parking_lot::RwLock;
use unique_type_id::UniqueTypeId;
use wasmer::{
    Cranelift, Exports, Extern, Features, Function, ImportObject, Instance, Memory, Module, NativeFunc, Store,
    Universal, UniversalEngine, Value, WasmerEnv,
};

use crate::dynamic_data::{ItemScript, Script, 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::WebAssemblyScriptCapabilities;
use crate::static_data::Item;
use crate::{PkmnResult, ScriptCategory, StringKey};

/// A WebAssembly script resolver implements the dynamic scripts functionality with WebAssembly.
pub struct WebAssemblyScriptResolver {
    engine: UniversalEngine,
    store: Store,
    module: Option<Module>,
    instance: Option<Instance>,
    memory: Option<Memory>,
    imports: HashMap<String, Function>,
    exports: Exports,
    exported_functions: HashMap<StringKey, Function>,
    load_script_fn: Option<NativeFunc<(u8, ExternRef<StringKey>), u32>>,
    allocate_mem_fn: Option<NativeFunc<(u32, u32), u32>>,
    script_capabilities: RwLock<HashMap<ScriptCapabilitiesKey, HashSet<WebAssemblyScriptCapabilities>>>,

    extern_ref_pointers: RwLock<Vec<*const u8>>,
    extern_ref_pointers_lookup: RwLock<HashMap<*const u8, u32>>,
    extern_ref_type_lookup: RwLock<HashMap<*const u8, u64>>,
}

#[derive(Debug, Clone, Eq, PartialEq, Hash)]
struct ScriptCapabilitiesKey {
    category: ScriptCategory,
    script_key: StringKey,
}

impl WebAssemblyScriptResolver {
    /// Instantiates a new WebAssemblyScriptResolver.
    pub fn new() -> Box<WebAssemblyScriptResolver> {
        let config = Cranelift::default();
        let mut features = Features::new();
        features.multi_value = true;
        features.reference_types = true;
        let universal = Universal::new(config).features(features);
        let engine = universal.engine();
        let store = Store::new(&engine);
        let s = Self {
            engine,
            store,
            module: Default::default(),
            instance: Default::default(),
            memory: Default::default(),
            imports: Default::default(),
            exports: Default::default(),
            exported_functions: Default::default(),
            load_script_fn: None,
            allocate_mem_fn: None,
            script_capabilities: Default::default(),
            extern_ref_pointers: Default::default(),
            extern_ref_pointers_lookup: Default::default(),
            extern_ref_type_lookup: Default::default(),
        };
        Box::new(s)
    }

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

        self.finalize();
    }

    /// Initialise all the data we need.
    fn finalize(&mut self) {
        let mut imports = ImportObject::new();
        let mut exports = Exports::new();

        let env = WebAssemblyEnv {
            resolver: self as *const WebAssemblyScriptResolver,
        };
        register_webassembly_funcs(&mut exports, &self.store, env);
        imports.register("env", exports);

        self.instance = Some(Instance::new(&self.module.as_ref().unwrap(), &imports).unwrap());
        let exports = &self.instance.as_ref().unwrap().exports;
        for export in exports.iter() {
            match export.1 {
                Extern::Function(f) => {
                    self.exported_functions.insert(export.0.as_str().into(), f.clone());
                }
                Extern::Memory(m) => {
                    self.memory = Some(m.clone());
                }
                _ => {}
            }
        }
        if let Some(m) = &self.memory {
            m.grow(32).unwrap();
        }
        if let Some(f) = self.exported_functions.get(&"load_script".into()) {
            self.load_script_fn = Some(f.native().unwrap())
        }
        if let Some(f) = self.exported_functions.get(&"allocate_mem".into()) {
            self.allocate_mem_fn = Some(f.native().unwrap())
        }
    }

    /// Gets the internal WASM memory.
    pub fn memory(&self) -> &Memory {
        self.memory.as_ref().unwrap()
    }

    /// 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_ref_index<T: UniqueTypeId<u64>>(&self, value: &T) -> u32 {
        let ptr = value as *const T as *const u8;
        if let Some(v) = self.extern_ref_pointers_lookup.read().get(&ptr) {
            return *v as u32;
        }
        let index = {
            let mut extern_ref_guard = self.extern_ref_pointers.write();
            extern_ref_guard.push(ptr);
            extern_ref_guard.len() as u32
        };
        self.extern_ref_pointers_lookup.write().insert(ptr, index);
        self.extern_ref_type_lookup.write().insert(ptr, T::id().0);
        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<T: UniqueTypeId<u64>>(&self, index: u32) -> &T {
        let read_guard = self.extern_ref_pointers.read();
        let ptr = read_guard.get((index - 1) as usize).unwrap();
        let expected_type_id = self.extern_ref_type_lookup.read()[&ptr];
        if expected_type_id != T::id().0 {
            panic!("Extern ref was accessed with wrong type");
        }

        unsafe { (*ptr as *const T).as_ref().unwrap() }
    }

    /// 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) -> (*const u8, u32) {
        let wasm_ptr = self.allocate_mem_fn.as_ref().unwrap().call(size, align).unwrap();
        unsafe {
            (
                self.memory.as_ref().unwrap().data_ptr().offset(wasm_ptr as isize),
                wasm_ptr,
            )
        }
    }
}

impl ScriptResolver for WebAssemblyScriptResolver {
    fn load_script(
        &self,
        owner: *const u8,
        category: ScriptCategory,
        script_key: &StringKey,
    ) -> PkmnResult<Option<Arc<dyn Script>>> {
        let script = self
            .load_script_fn
            .as_ref()
            .unwrap()
            .call(
                category as u8,
                ExternRef::from_index(self.get_extern_ref_index(script_key)),
            )
            .unwrap();
        if script == 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.get(&"get_script_capabilities".into()) {
                    let res = get_cap.call(&[Value::I32(script as i32)]).unwrap();
                    let ptr = (self.memory.as_ref().unwrap().data_ptr() 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.offset(i as isize));
                    }
                }
            }
            self.script_capabilities.write().insert(key.clone(), capabilities);
        }

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

        Ok(Some(Arc::new(WebAssemblyScript::new(
            owner as *mut u8,
            script,
            capabilities as *const HashSet<WebAssemblyScriptCapabilities>
                as *mut HashSet<WebAssemblyScriptCapabilities>,
            self as *const WebAssemblyScriptResolver as *mut WebAssemblyScriptResolver,
            script_key.clone(),
        ))))
    }

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

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

#[derive(Clone)]
pub(crate) struct WebAssemblyEnv {
    pub resolver: *const WebAssemblyScriptResolver,
}

impl WebAssemblyEnv {
    pub fn resolver(&self) -> &WebAssemblyScriptResolver {
        unsafe { self.resolver.as_ref().unwrap() }
    }
}

unsafe impl Sync for WebAssemblyEnv {}

unsafe impl Send for WebAssemblyEnv {}

impl WasmerEnv for WebAssemblyEnv {}