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, instance: Option, memory: Option, imports: HashMap, exports: Exports, exported_functions: HashMap, load_script_fn: Option), u32>>, allocate_mem_fn: Option>, script_capabilities: RwLock>>, extern_ref_pointers: RwLock>, extern_ref_pointers_lookup: RwLock>, extern_ref_type_lookup: RwLock>, } #[derive(Debug, Clone, Eq, PartialEq, Hash)] struct ScriptCapabilitiesKey { category: ScriptCategory, script_key: StringKey, } impl WebAssemblyScriptResolver { /// Instantiates a new WebAssemblyScriptResolver. pub fn new() -> Box { 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>(&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>(&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>> { 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 as *mut HashSet, self as *const WebAssemblyScriptResolver as *mut WebAssemblyScriptResolver, script_key.clone(), )))) } fn load_item_script(&self, _key: &Item) -> PkmnResult>> { 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 {}