low_luajit/lua/type/

function.rs

use core::{error::Error, ffi::{CStr, c_int, c_void}, fmt::{self, Display, Formatter}, mem, num::NonZeroI32, slice};
use crate::{LuaState, SafeFn, ffi::{self, lua_CFunction, lua_State}};

/// A safe function using the C ABI that Lua can call.
///
/// This type is technically a *subset* of [`lua_CFunction`], as it only allows safe functions.
/// Different methods will use either [`lua_CFunction`] or `CFunction` based on if they work with only safe functions or not.
///
/// It is always safe to convert a `CFunction` into a [`lua_CFunction`] (see [`SafeFn::to_raw`]).
///
/// # Protocol
///
/// To communicate properly, Lua expects `CFunction`s to follow a specific protocol:
///
/// Arguments passed to the `CFunction` are located in the stack (index 1 contains the first argument, etc.),
/// thus [`LuaState::get_top`] can be used to know the number of arguments received.
///
/// Values to return to the caller should be pushed onto the stack (first value is returned first, etc.),
/// using the returned [`AmountReturned`] to specify how many values were returned.
///
/// [`AmountReturned`] can be obtained manually using [`LuaState::return_any`] or automatically for simple values
/// using the designated methods ([`LuaState::return_number`], [`LuaState::return_string`], etc...).
///
/// # Examples
///
/// ```
/// use low_luajit::{LuaState, AmountReturned};
///
/// extern "C-unwind" fn sum(mut lua: LuaState) -> AmountReturned {
///     // Retrieve the arguments, asserting that they exist and are of the correct type.
///     let x = lua.check_number(1);
///     let y = lua.check_number(2);
///
///     // Push the value to return and return 1 since we're only returning a single value.
///     lua.push_number(x + y);
///     lua.return_any(1)
/// }
///
/// # let mut lua = unsafe { low_luajit::__test!() };
/// lua.push_cfunction(sum);
/// lua.push_number(123.0);
/// lua.push_number(456.0);
/// lua.call(2, 1);
/// assert_eq!(lua.to_number(-1), 579.0);
/// ```
///
/// Alternatively, using [`LuaState::return_number`]:
///
/// ```
/// use low_luajit::{LuaState, AmountReturned};
///
/// extern "C-unwind" fn sum(lua: LuaState) -> AmountReturned {
///     // Retrieve the arguments, asserting that they exist and are of the correct type.
///     let x = lua.check_number(1);
///     let y = lua.check_number(2);
///
///     // Just return the number (does the same as the other example, just more concise).
///     lua.return_number(x + y)
/// }
///
/// # let mut lua = unsafe { low_luajit::__test!() };
/// lua.push_cfunction(sum);
/// lua.push_number(0.5);
/// lua.push_number(0.5);
/// lua.call(2, 1);
/// assert_eq!(lua.to_number(-1), 1.0);
/// ```
///
/// [`lua_CFunction`]: ./ffi/type.lua_CFunction.html
pub type CFunction = extern "C-unwind" fn(lua: LuaState) -> AmountReturned;

impl SafeFn for CFunction {
	type RawFn = lua_CFunction;

	#[inline]
	unsafe fn from_raw(raw: lua_CFunction) -> Self {
		// SAFETY: lua_CFunction and CFunction are effectively the same.
		unsafe { mem::transmute(raw) }
	}

	#[inline]
	fn to_raw(self) -> lua_CFunction {
		// SAFETY: CFunction and lua_CFunction are effectively the same.
		unsafe { mem::transmute(self) }
	}
}

/// Simple wrapper that tells Lua how many values a [`CFunction`] is intending to return.
///
/// Without this wrapper, Lua would willingly attempt to return as many values as you say,
/// which can easily go beyond the stack and result in undefined behavior.
///
/// This wrapper can only be (safely) created by [`LuaState`],
/// which will check if the amount is valid before returning the `AmountReturned` (see [`LuaState::return_any`]).
#[repr(transparent)]
#[derive(Debug)]
#[must_use = "this `AmountReturned` should be returned to the caller"]
pub struct AmountReturned(core::ffi::c_int);

impl AmountReturned {
	/// Tells Lua that no values are intended to be returned.
	///
	/// Always safe to use since nothing gets popped.
	pub const NONE: Self = Self(0);

	/// Wraps `amount` around an `AmountReturned`, without doing any checks.
	///
	/// # Safety
	///
	/// `amount` must not be higher than the amount of values the Lua stack is holding at the time of returning.
	pub const unsafe fn unchecked(amount: i32) -> Self {
		Self(amount)
	}
}

/// Trait for all closures and functions that can be used with [`LuaState::dump`].
///
/// See [`LuaState::dump`] for more information.
pub trait Writer: FnMut(&LuaState, &[u8]) -> Result<(), NonZeroI32> {}
impl<T: FnMut(&LuaState, &[u8]) -> Result<(), NonZeroI32>> Writer for T {}

/*/// Trait for all closures and functions that can be used with [`LuaState::load`].
///
/// See [`LuaState::load`] for more information.
pub trait Reader<T: AsRef<[u8]>>: FnMut(&mut LuaState) -> Option<T> {}
impl<T: AsRef<[u8]>, U: FnMut(&mut LuaState) -> Option<T>> Reader<T> for U {}*/

/// All error codes that a function may return when erroring while called in protected mode.
///
/// This enum is used by the following functions:
/// - [`LuaState::pcall`].
/// - [`LuaState::cpcall`].
#[repr(i32)]
#[derive(Debug)]
pub enum PCallError {
	/// Any error during runtime.
	RuntimeError = ffi::LUA_ERRRUN,

	/// Memory allocation error (rare).
	///
	/// *([`LuaState::pcall`] will not call the error handler for these errors.)*
	MemoryError = ffi::LUA_ERRMEM,

	/// Error while calling the error handler.
	ErrorHandlerError = ffi::LUA_ERRERR,
}

impl Display for PCallError {
	#[inline]
	fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
		f.write_str(match self {
			Self::RuntimeError => "runtime error",
			Self::MemoryError => "memory allocation error",
			Self::ErrorHandlerError => "error within error handler",
		})
	}
}

impl Error for PCallError {}

impl PCallError {
	/// Converts the given `value` into a `Result`.
	///
	/// # Safety
	///
	/// `value` must either be `0` or a valid `PCallError` value.
	///
	/// # Errors
	///
	/// Returns `Err` if `value` is not `0`.
	#[inline]
	pub const unsafe fn check_raw(value: i32) -> Result<(), Self> {
		// SAFETY: The caller must only give 0 or a valid error code.
		if value == 0 { Ok(()) } else { Err(unsafe { mem::transmute::<i32, PCallError>(value) }) }
	}
}

/// All error codes that may be returned when failing to load a chunk.
///
/// This enum is used by the following functions:
/// - [`LuaState::load_bytes`].
/// - [`LuaState::load_cstring`].
/// - [`LuaState::load_file`].
#[repr(i32)]
#[derive(Debug)]
pub enum LoadError {
	/// Syntax error during pre-compilation.
	SyntaxError = ffi::LUA_ERRSYNTAX,

	/// Memory allocation error (rare).
	MemoryError = ffi::LUA_ERRMEM,

	/// *(Only returned by [`LuaState::load_file`].)*
	FileError = ffi::LUA_ERRFILE,
}

impl Display for LoadError {
	#[inline]
	fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
		f.write_str(match self {
			Self::SyntaxError => "syntax error",
			Self::MemoryError => "memory allocation error",
			Self::FileError => "file error",
		})
	}
}

impl Error for LoadError {}

impl LoadError {
	/// Converts the given `value` into a `Result`.
	///
	/// # Safety
	///
	/// `value` must either be `0` or a valid `LoadError` value.
	///
	/// # Errors
	///
	/// Returns `Err` if `value` is not `0`.
	#[inline]
	pub const unsafe fn check_raw(value: i32) -> Result<(), Self> {
		// SAFETY: The caller must only give 0 or a valid error code.
		if value == 0 { Ok(()) } else { Err(unsafe { mem::transmute::<i32, LoadError>(value) }) }
	}
}

/// Function related functions.
impl LuaState {
	/// Finishes the currently called [`CFunction`], returning the last `amount` elements that were pushed to the stack to the caller.
	///
	/// The `LuaState` wrapper is consumed to prevent invalidating the `AmountReturned` (the Lua state itself is left intact).
	///
	/// # Panics
	///
	/// Panics if `amount` is negative or higher than the amount of elements currently held in the stack (see [`LuaState::get_top`]).
	///
	/// # Examples
	///
	/// ```
	/// use low_luajit::{LuaState, AmountReturned};
	///
	/// extern "C-unwind" fn example(mut lua: LuaState) -> AmountReturned {
	///     lua.push_number(12.3); // First value returned.
	///     lua.push_string("abc"); // Second value returned.
	///     lua.return_any(2)
	/// }
	///
	/// # let mut lua = unsafe { low_luajit::__test!() };
	/// lua.push_cfunction(example);
	/// lua.call(0, 2);
	/// assert_eq!(lua.to_number(1), 12.3);
	/// assert_eq!(lua.to_string(2), Some("abc"));
	/// ```
	#[inline]
	pub fn return_any(self, amount: i32) -> AmountReturned {
		assert!(amount >= 0, "cannot return a negative amount of values");
		assert!(amount <= self.get_top(), "tried to return {amount} values but the stack currently holds {} values", self.get_top());

		// SAFETY: amount should not be higher than the amount of values currently present in the stack.
		unsafe { AmountReturned::unchecked(amount) }
	}

	/// Returns `true` if the value at the given acceptable `index` is a function (either [`CFunction`] or Lua function), `false` otherwise.
	#[inline]
	#[must_use]
	pub fn is_function(&self, index: i32) -> bool {
		// SAFETY: self.as_ptr() is always valid.
		unsafe { ffi::lua_isfunction(self.as_ptr(), index) != 0 }
	}

	/// Returns `true` if the value at the given acceptable `index` is a [`CFunction`], `false` otherwise.
	#[inline]
	#[must_use]
	pub fn is_cfunction(&self, index: i32) -> bool {
		// SAFETY: self.as_ptr() is always valid.
		unsafe { ffi::lua_iscfunction(self.as_ptr(), index) != 0 }
	}

	/// Returns `true` if the value at the given acceptable `index` is a Lua function, `false` otherwise.
	#[inline]
	#[must_use]
	pub fn is_lua_function(&self, index: i32) -> bool {
		self.is_function(index) && !self.is_cfunction(index)
	}

	/// Returns the [`lua_CFunction`] at the given `index`, or `None` if the value is anything else.
	///
	/// [`lua_CFunction`] is returned instead of [`CFunction`] as this function cannot determine its safety.
	/// Additionally, this function may sometimes return `None` for some C functions even if [`LuaState::is_cfunction`] returns `true`.
	///
	/// If you just want to call a function, you probably want to use [`LuaState::call`], [`LuaState::pcall`], or [`LuaState::cpcall`] instead.
	///
	/// # Examples
	///
	/// ```
	/// # use low_luajit::{LuaState, AmountReturned};
	/// # extern "C-unwind" fn example(lua: LuaState) -> AmountReturned {
	/// let f = lua.to_cfunction(-1).unwrap();
	/// // SAFETY: This function is safe to call.
	/// unsafe { f(lua.as_ptr()) };
	/// # AmountReturned::NONE }
	/// ```
	#[inline]
	#[must_use]
	pub fn to_cfunction(&self, index: i32) -> Option<lua_CFunction> {
		// SAFETY: self.as_ptr() is always valid.
		unsafe { ffi::lua_tocfunction(self.as_ptr(), index) }
	}

	/// Pushes the given [`CFunction`] `func` onto the stack.
	///
	/// To associate values with `func`, see [`LuaState::push_cclosure`].
	#[inline]
	pub fn push_cfunction(&mut self, func: CFunction) {
		// SAFETY: self.as_ptr() and transmuting from CFunction to lua_CFunction are always valid.
		unsafe { ffi::lua_pushcfunction(self.as_ptr(), func.to_raw()); }
	}

	/// Pushes the given [`CFunction`] `func` onto the stack, while also associating `n` values to it (called upvalues).
	///
	/// To associate values with `func`, simply push said values before calling this function;
	/// they will be popped but will be accessible from within `func` at the pseudo-indices returned by [`LuaState::upvalue_index`].
	///
	/// If `n` is `0` this function will behave exactly as [`LuaState::push_cfunction`].
	///
	/// # Panics
	///
	/// Panics if the stack currently contains less than `n` values.
	#[inline]
	pub fn push_cclosure(&mut self, func: CFunction, n: u8) {
		assert!(i32::from(n) <= self.get_top(), "tried to push a `CFunction` with {n} upvalues but the stack currently holds {} values", self.get_top());

		// SAFETY: self.as_ptr() is always valid, the raw version of func is safe to call by Lua, and the stack contains enough values to pop.
		unsafe { ffi::lua_pushcclosure(self.as_ptr(), func.to_raw(), i32::from(n)); }
	}

	/// Sets the given [`CFunction`] `func` as a global named `name`.
	///
	/// # Lua Errors
	///
	/// Indexing the global environment may trigger the `__newindex` metamethod which could error.
	#[inline]
	pub fn register_cfunction(&mut self, name: &CStr, func: CFunction) {
		// SAFETY: self.as_ptr() and name's pointer are always valid, and the raw version of func is safe to call by Lua.
		unsafe { ffi::lua_register(self.as_ptr(), name.as_ptr(), func.to_raw()); }
	}

	/// When passed as the number of expected results for the called functions in [`LuaState::call`] or [`LuaState::pcall`],
	/// it will tell Lua to push to the stack exactly as many values as the called function returned.
	pub const MULTI_RETURN: i32 = ffi::LUA_MULTRET;

	/// Calls a function.
	///
	/// To properly call a function, follow these steps:
	/// 1. Push the function onto the stack.
	/// 2. Push the arguments of the function (as many as `n_args`), pushing the first arg first and so on.
	/// 3. Call this function, which will pop the function and arguments and push the results (as many as `n_results`) in the same order.
	///
	/// If `n_results` is [`LuaState::MULTI_RETURN`] (`-1`), all the results from the function are pushed to the stack.
	/// If `n_results` is higher than the amount of values returned by the function, additional `nil`s will be pushed to compensate.
	///
	/// # Panics
	///
	/// Panics if the stack does not contain the function itself and its arguments.
	///
	/// # Lua Errors
	///
	/// The called function itself may error, to catch any error consider using [`LuaState::pcall`] or [`LuaState::cpcall`] instead.
	///
	/// Additionally, this function will error if the first element popped (the function to call) is not a function.
	#[inline]
	pub fn call(&mut self, n_args: i32, n_results: i32) {
		let top = self.get_top();
		assert!(top > 0, "tried to call a function despite having an empty stack");
		assert!(n_args < top, "cannot pass {n_args} arguments to the called function, as there are only {top} elements");

		// SAFETY: self.as_ptr() is always valid and the stack contains enough values to pop.
		unsafe { ffi::lua_call(self.as_ptr(), n_args, n_results); }
	}

	/// Calls a function in protected mode.
	///
	/// See [`LuaState::call`] for how to properly call a function.
	///
	/// If the called function does not error, this function behaves exactly like [`LuaState::call`].
	///
	/// `error_func` should either be `0` or the index to an error handler function,
	/// which will be called on error and its returned value will be treated as the error object instead of the original error message.
	/// *(Pesudo-indices cannot be used for `error_func`.)*
	///
	/// # Panics
	///
	/// Panics if the stack does not contain the function itself and its arguments.
	///
	/// # Errors
	///
	/// If the called function triggers a Lua error,
	/// `Err` will be returned with the appropriate [`PCallError`] and the error object is pushed onto of the stack.
	///
	/// Additionally, this function will error if the first element popped (the function to call)
	/// or the element at index `error_func` are not functions.
	#[inline]
	pub fn pcall(&mut self, n_args: i32, n_results: i32, error_func: i32) -> Result<(), PCallError> {
		let top = self.get_top();
		assert!(top > 0, "tried to pcall a function despite having an empty stack");
		assert!(n_args < top, "cannot pass {n_args} arguments to the pcalled function, as there are only {top} elements");

		// SAFETY: self.as_ptr() is always valid, the stack contains enough values to pop, and the returned value is either 0 or a valid PCallError.
		unsafe { PCallError::check_raw(ffi::lua_pcall(self.as_ptr(), n_args, n_results, error_func)) }
	}

	/// Calls the function `func` in protected mode with a single argument: a light userdata containing `ud`.
	///
	/// Any value returned by `func` will be discarded (the stack is not changed).
	///
	/// # Errors
	///
	/// If `func` triggers a Lua error,
	/// `Err` will be returned with the appropriate [`PCallError`] and the error object is pushed onto of the stack.
	#[inline]
	pub fn cpcall<T>(&mut self, func: CFunction, ud: *mut T) -> Result<(), PCallError> {
		// SAFETY: self.as_ptr() is always valid, func is safe to call, ud is only stored and the returned value is either 0 or a valid PCallError.
		unsafe { PCallError::check_raw(ffi::lua_cpcall(self.as_ptr(), func.to_raw(), ud.cast())) }
	}

	/// Dumps the Lua function at the top of the stack as a binary chunk ([`CFunction`]s cannot be dumped).
	///
	/// Dumping is done by sending fragments of the chunk to the `writer` function, which should combine them together into the full chunk.
	///
	/// If an error happens during dumping, `Err` is returned with any non-zero number.
	///
	/// Reloading the chunk results in a function equivalent to the one originally dumped (see [`LuaState::load_bytes`]).
	///
	/// # Errors
	///
	/// If the top element is not a Lua function (or the stack is empty), this function will error returning `1`.
	///
	/// `writer` may error and return any other non-zero number (`1` included), based on its implementation.
	///
	/// # Examples
	///
	/// ```
	/// # use low_luajit::{LuaState, AmountReturned};
	/// # extern "C-unwind" fn example(lua: LuaState) -> AmountReturned {
	/// if lua.is_lua_function(-1) {
	///     let mut chunk = Vec::new();
	///     lua.dump(|_, fragment| {
	///         chunk.extend_from_slice(fragment);
	///         Ok(())
	///     }).expect("the given writer never errors");
	///     // Use the chunk in some way...
	/// }
	/// # AmountReturned::NONE }
	/// ```
	#[inline]
	pub fn dump<F: Writer>(&self, mut writer: F) -> Result<(), NonZeroI32> {
		unsafe extern "C-unwind" fn middleware<F: Writer>(lua: *mut lua_State, ptr: *const c_void, size: usize, writer: *mut c_void) -> c_int {
			// SAFETY: writer should be a valid pointer to the writer of the upper function, the lua_State should always be valid,
			// ptr should be non-null and valid for size amount of values, and Result<(), NonZeroI32> can be transmuted to i32.
			unsafe { mem::transmute((&mut *writer.cast::<F>())(&LuaState::from_ptr(lua), slice::from_raw_parts(ptr.cast(), size))) }
		}

		// SAFETY: self.as_ptr() is always valid, lua_dump calls middleware properly, and i32 can be transmuted to Result<(), NonZeroI32>.
		unsafe { mem::transmute(ffi::lua_dump(self.as_ptr(), middleware::<F>, (&raw mut writer).cast())) }
	}

	/// Loads the given `bytes` as a Lua chunk named `chunk_name`, pushing the resulting function onto the stack (if successful).
	///
	/// *(The resulting function is only pushed, not called.)*
	///
	/// While `chunk_name` can be any value, there are some common naming conventions:
	/// - If the name is the path to a file, it should be prefixed with `'@'`. *([`LuaState::load_file`] does this.)*
	/// - If the name is a custom identifier for the function, it should be prefixed with `'='`. *(All [`CFunction`]s are named `"=[C]"`.)*
	/// - Anything else will be interpreted as if being the function's source itself. *([`LuaState::load_cstring`] does this.)*
	///
	/// # Errors
	///
	/// If unable to load the given `bytes`,
	/// `Err` will be returned with the appropriate [`LoadError`] and the error message is pushed onto the stack.
	#[inline]
	pub fn load_bytes(&mut self, bytes: &[u8], chunk_name: &CStr) -> Result<(), LoadError> {
		// SAFETY: self.as_ptr(), bytes and chunk_name are always valid, u8 can be casted as c_char (they're both 1 byte),
		// the length is correct as it's returned from the array itself, and the returned value is either 0 or a valid LoadError.
		unsafe { LoadError::check_raw(ffi::luaL_loadbuffer(self.as_ptr(), bytes.as_ptr().cast(), bytes.len(), chunk_name.as_ptr())) }
	}

	/// Loads the given `cstring` as a Lua chunk, pushing the resulting function onto the stack (if successful).
	///
	/// *(The resulting function is only pushed, not called.)*
	///
	/// The chunk will be named after the given `cstring` itself.
	///
	/// # Errors
	///
	/// If unable to load the given `cstring`,
	/// `Err` will be returned with the appropriate [`LoadError`] and the error message is pushed onto the stack.
	#[inline]
	pub fn load_cstring(&mut self, cstring: &CStr) -> Result<(), LoadError> {
		// SAFETY: self.as_ptr() and cstring's pointer are always valid and the returned value is either 0 or a valid LoadError.
		unsafe { LoadError::check_raw(ffi::luaL_loadstring(self.as_ptr(), cstring.as_ptr())) }
	}

	/// Attempts to open the file named `filename` and load it as a Lua chunk, pushing the resulting function onto the stack (if successful).
	///
	/// *(The resulting function is only pushed, not called.)*
	///
	/// The chunk will be named after the given `filename` itself.
	///
	/// # Errors
	///
	/// If unable to open the file or load its contents,
	/// `Err` will be returned with the appropriate [`LoadError`] and the error message is pushed onto the stack.
	#[inline]
	pub fn load_file(&mut self, filename: &CStr) -> Result<(), LoadError> {
		// SAFETY: self.as_ptr() and filename's pointer are always valid and the returned value is either 0 or a valid LoadError.
		unsafe { LoadError::check_raw(ffi::luaL_loadfile(self.as_ptr(), filename.as_ptr())) }
	}

	/*#[inline] TODO: this doesn't work, also it might be over the top? maybe it should just accept the chunk as a simple variable?
	pub fn load<T: AsRef<[u8]>, F: Reader<T>>(&mut self, mut reader: F, chunk_name: &CStr) {
		unsafe extern "C-unwind" fn middleware<T: AsRef<[u8]>, F: Reader<T>>(lua: *mut lua_State, reader: *mut c_void, size: *mut usize) -> *const c_char {
			unsafe { (&mut *reader.cast::<F>())(&mut LuaState::from_ptr(lua)).map_or(ptr::null(), |fragment| {
				*size = fragment.as_ref().len();
				(&raw const fragment).cast()
			}) }
		}

		unsafe { ffi::lua_load(self.as_ptr(), middleware::<T, F>, (&raw mut reader).cast(), chunk_name.as_ptr()); }
	}*/

	/// Finishes the currently called [`CFunction`], returning the given `func` to the caller.
	///
	/// The `LuaState` wrapper is consumed to prevent invalidating the `AmountReturned` (the Lua state itself is left intact).
	#[inline]
	pub fn return_cfunction(mut self, func: CFunction) -> AmountReturned {
		self.push_cfunction(func);

		// SAFETY: There should always be at least one element in the stack.
		unsafe { AmountReturned::unchecked(1) }
	}

	/// Finishes the currently called [`CFunction`], returning the given `func` to the caller while also associating `n` values to it
	/// (see [`LuaState::push_cclosure`]).
	///
	/// The `LuaState` wrapper is consumed to prevent invalidating the `AmountReturned` (the Lua state itself is left intact).
	///
	/// # Panics
	///
	/// Panics if the stack currently contains less than `n` values.
	#[inline]
	pub fn return_cclosure(mut self, func: CFunction, n: u8) -> AmountReturned {
		self.push_cclosure(func, n);

		// SAFETY: There should always be at least one element in the stack.
		unsafe { AmountReturned::unchecked(1) }
	}

	/// If `condition` is `false`, [`LuaState::arg_error`] is called using `arg_n` and `extra_msg`.
	/// Otherwise nothing happens.
	///
	/// # Lua Errors
	///
	/// Intentionally errors if `condition` is `false`.
	#[inline]
	pub fn arg_check(&self, condition: bool, arg_n: i32, extra_msg: &CStr) {
		if !condition {
			self.arg_error(arg_n, extra_msg)
		}
	}
}