path: root/src/script.rs

use std::path::Path;

use rustc_hash::FxHashMap;

use crate::{buffer::Buffer, file::FileManager, input::InputManager};

pub struct ScriptManager<'a> {
	scripts: &'a [ScriptDeclaration],
	pub active_scripts: Vec<usize>,
	script_states: FxHashMap<usize, Box<[u8]>>,
}

pub struct WasmRunner {
	pub engine: wasmi::Engine,
	pub linker: wasmi::Linker<&'static ScriptContext<'static>>,
}

pub struct ScriptContext<'a> {
	pub buffer: &'a mut Buffer,
	pub input: &'a InputManager,
}

type StaticScriptFunction = Option<fn(&mut ScriptContext<'_>, &mut Box<[u8]>)>;

pub enum ScriptDeclaration {
	Static {
		start: StaticScriptFunction,
		update: StaticScriptFunction,
		end: StaticScriptFunction,
	},
	Wasm {
		path: Box<Path>,
	},
}

impl<'a> ScriptContext<'a> {
	pub fn exit(&mut self) -> ! {
		let _ = self.buffer.finish();
		std::process::exit(0);
	}
}

fn get_script_state(states: &mut FxHashMap<usize, Box<[u8]>>, script_id: usize) -> &mut Box<[u8]> {
	states.entry(script_id).or_insert(Box::from([]))
}

impl Default for WasmRunner {
	fn default() -> Self {
		Self::new()
	}
}

impl WasmRunner {
	pub fn new() -> Self {
		let engine = wasmi::Engine::default();
		let linker = wasmi::Linker::new(&engine);
		Self { engine, linker }
	}

	pub fn compile(&self, data: &[u8]) -> Result<wasmi::Module, wasmi::Error> {
		wasmi::Module::new(&self.engine, data)
	}

	pub fn run_function(
		&self,
		module: &wasmi::Module,
		context: &mut ScriptContext<'_>,
		function_name: &str,
	) -> Result<(), wasmi::Error> {
		// safety: the store won't be kept around after this function ends. Only the
		//         lifetimes are being transmuted
		let mut store = wasmi::Store::new(&self.engine, unsafe {
			std::mem::transmute::<&ScriptContext<'_>, &ScriptContext<'_>>(context)
		});
		let instance = self.linker.instantiate_and_start(&mut store, module)?;
		let function = instance.get_func(&store, function_name).ok_or_else(|| {
			wasmi::Error::new(format!("No function named {function_name} was exported"))
		})?;
		function.call(&mut store, &[], &mut [])
	}
}

impl<'a> ScriptManager<'a> {
	pub fn new(scripts: &'a [ScriptDeclaration]) -> Self {
		Self {
			scripts,
			active_scripts: Vec::new(),
			script_states: FxHashMap::default(),
		}
	}

	pub fn activate_script(
		&mut self,
		script_id: usize,
		context: &mut ScriptContext<'_>,
		wasm_runner: &WasmRunner,
		files: &mut FileManager,
	) {
		let Some(script_declaration) = self.scripts.get(script_id) else {
			unreachable!("The active script does not exist");
		};
		match script_declaration {
			ScriptDeclaration::Static { start, .. } => {
				start.map(|s| {
					s(
						context,
						get_script_state(&mut self.script_states, script_id),
					)
				});
			}
			ScriptDeclaration::Wasm { path } => {
				let module = files
					.load_file(path.clone(), &wasm_runner.engine)
					.unwrap()
					.unwrap_wasm_module();
				wasm_runner.run_function(module, context, "start");
			}
		}
		self.active_scripts.push(script_id);
	}

	pub fn deactivate_script(
		&mut self,
		script_id: usize,
		context: &mut ScriptContext<'_>,
		wasm_runner: &WasmRunner,
		files: &mut FileManager,
	) {
		self.active_scripts.retain(|element| {
			if *element != script_id {
				let Some(script) = self.scripts.get(script_id) else {
					unreachable!("The active script does not exist");
				};
				match script {
					ScriptDeclaration::Static { end, .. } => {
						end.map(|s| {
							s(
								context,
								get_script_state(&mut self.script_states, script_id),
							);
						});
					}
					ScriptDeclaration::Wasm { path } => {
						let module = files
							.load_file(path.clone(), &wasm_runner.engine)
							.unwrap()
							.unwrap_wasm_module();
						wasm_runner.run_function(module, context, "end");
					}
				};

				true
			} else {
				false
			}
		});
	}

	pub fn run_scripts(
		&mut self,
		context: &mut ScriptContext<'_>,
		wasm_runner: &WasmRunner,
		files: &mut FileManager,
	) {
		puffin::profile_function!();
		for script_id in &self.active_scripts {
			let Some(script) = self.scripts.get(*script_id) else {
				unreachable!("The active script does not exist");
			};
			match script {
				ScriptDeclaration::Static { update, .. } => {
					if let Some(s) = update {
						s(
							context,
							get_script_state(&mut self.script_states, *script_id),
						)
					}
				}
				ScriptDeclaration::Wasm { path } => {
					let module = files
						.load_file(path.clone(), &wasm_runner.engine)
						.unwrap()
						.unwrap_wasm_module();
					wasm_runner.run_function(module, context, "update");
				}
			}
		}
	}

	pub fn reset_script(
		&mut self,
		script_id: usize,
		context: &mut ScriptContext<'_>,
		wasm_runner: &WasmRunner,
		files: &mut FileManager,
	) {
		self.activate_script(script_id, context, wasm_runner, files);
		self.deactivate_script(script_id, context, wasm_runner, files);
	}
}