strafe-project/lib/roblox_emulator/src/runner/script_signal.rs

use std::{cell::RefCell,rc::Rc};

use mlua::UserDataFields;

#[derive(Clone)]
struct FunctionList{
	functions:Vec<mlua::Function>,
}
impl FunctionList{
	pub fn new()->Self{
		Self{
			functions:Vec::new(),
		}
	}
	// This eats the Lua error
	pub fn fire(self,args:&mlua::MultiValue){
		// Make a copy of the list in case Lua attempts to modify it during the loop
		for function in self.functions{
			//wee let's allocate for our function calls
			if let Err(e)=function.call::<mlua::MultiValue>(args.clone()){
				println!("Script Signal Error: {e}");
			}
		}
	}
}
#[derive(Clone)]
struct RcFunctionList{
	functions:Rc<RefCell<FunctionList>>,
}
impl RcFunctionList{
	pub fn new()->Self{
		Self{
			functions:Rc::new(RefCell::new(FunctionList::new())),
		}
	}
	pub fn fire(&self,args:&mlua::MultiValue){
		// Make a copy of the list in case Lua attempts to modify it during the loop
		self.functions.borrow().clone().fire(args)
	}
}
#[derive(Clone)]
pub struct ScriptSignal{
	// Emulate the garbage roblox api.
	// ScriptConnection should not exist.
	// :Disconnect should be a method on ScriptSignal, and this would be avoided entirely.
	connections:RcFunctionList,
	once:RcFunctionList,
	wait:Rc<RefCell<Vec<mlua::Thread>>>,
}
pub struct ScriptConnection{
	connection:RcFunctionList,
	function:mlua::Function,
}
impl ScriptSignal{
	pub fn new()->Self{
		Self{
			connections:RcFunctionList::new(),
			once:RcFunctionList::new(),
			wait:Rc::new(RefCell::new(Vec::new())),
		}
	}
	pub fn fire(&self,args:&mlua::MultiValue){
		self.connections.fire(args);
		//Replace the FunctionList with an empty one and drop the borrow
		let once=std::mem::replace(&mut *self.once.functions.borrow_mut(),FunctionList::new());
		once.fire(args);
		//resume threads waiting for this signal
		let threads=std::mem::replace(&mut *self.wait.borrow_mut(),Vec::new());
		for thread in threads{
			if let Err(e)=thread.resume::<mlua::MultiValue>(args.clone()){
				println!("Script Signal thread resume Error: {e}");
			}
		}
	}
	pub fn connect(&self,function:mlua::Function)->ScriptConnection{
		self.connections.functions.borrow_mut().functions.push(function.clone());
		ScriptConnection{
			connection:self.connections.clone(),
			function,
		}
	}
	pub fn once(&self,function:mlua::Function)->ScriptConnection{
		self.once.functions.borrow_mut().functions.push(function.clone());
		ScriptConnection{
			connection:self.once.clone(),
			function,
		}
	}
	pub fn wait(&self,thread:mlua::Thread){
		self.wait.borrow_mut().push(thread);
	}
}
impl ScriptConnection{
	pub fn position(&self)->Option<usize>{
		self.connection.functions.borrow().functions.iter().position(|function|function==&self.function)
	}
}

impl mlua::UserData for ScriptSignal{
	fn add_methods<M:mlua::UserDataMethods<Self>>(methods:&mut M){
		methods.add_method("Connect",|_lua,this,f:mlua::Function|
			Ok(this.connect(f))
		);
		methods.add_method("Once",|_lua,this,f:mlua::Function|
			Ok(this.once(f))
		);
		// Fire is not allowed to be called from Lua
		// methods.add_method("Fire",|_lua,this,args:mlua::MultiValue|
		// 	Ok(this.fire(args))
		// );
	}
}
impl mlua::FromLua for ScriptSignal{
	fn from_lua(value:mlua::Value,_lua:&mlua::Lua)->Result<Self,mlua::Error>{
		match value{
			mlua::Value::UserData(ud)=>Ok(ud.borrow::<Self>()?.clone()),
			other=>Err(mlua::Error::runtime(format!("Expected {} got {:?}",stringify!(ScriptSignal),other))),
		}
	}
}

impl mlua::UserData for ScriptConnection{
	fn add_fields<F:mlua::UserDataFields<Self>>(fields:&mut F){
		fields.add_field_method_get("Connected",|_,this|{
			Ok(this.position().is_some())
		});
	}
	fn add_methods<M:mlua::UserDataMethods<Self>>(methods:&mut M){
		methods.add_method("Disconnect",|_,this,_:()|{
			if let Some(index)=this.position(){
				this.connection.functions.borrow_mut().functions.remove(index);
			}
			Ok(())
		});
	}
}

fn wait_thread(lua:&mlua::Lua,this:ScriptSignal)->Result<(),mlua::Error>{
	Ok(this.wait(lua.current_thread()))
}

// This is used to avoid calling coroutine.yield from the rust side.
const LUA_WAIT:&str=
"local coroutine_yield=coroutine.yield
local wait_thread=wait_thread
return function(signal)
	wait_thread(signal)
	return coroutine_yield()
end";

pub fn set_globals(lua:&mlua::Lua,globals:&mlua::Table)->Result<(),mlua::Error>{
	let coroutine_table=globals.get::<mlua::Table>("coroutine")?;
	let wait_thread=lua.create_function(wait_thread)?;

	//create wait function environment
	let wait_env=lua.create_table()?;
	wait_env.raw_set("coroutine",coroutine_table)?;
	wait_env.raw_set("wait_thread",wait_thread)?;

	//construct wait function from Lua code
	let wait=lua.load(LUA_WAIT)
	.set_name("wait")
	.set_environment(wait_env)
	.call::<mlua::Function>(())?;

	lua.register_userdata_type::<ScriptSignal>(|reg|{
		reg.add_field("Wait",wait);
		mlua::UserData::register(reg);
	})?;

	Ok(())
}