low_luajit/lua/type/

table.rs

use core::{ffi::CStr, marker::PhantomData, ops::Not};
use crate::{LuaState, ffi, map_ptr};

macro_rules! iterator_impl {
	($iterator:ident, $assert_err:literal; $(impla {$($impl_fn:item)*};)? $($iter_fn:item)+) => {
		impl $iterator<'_> {
			/// Constructs a new iterator over the table at the given `index`.
			///
			/// # Panics
			///
			/// Panics if the element at the given `index` is not a table.
			#[inline]
			pub fn new(lua: &LuaState, index: i32) -> Self {
				let iterator = Self { lua: lua.as_ptr(), _marker: PhantomData };

				let mut lua = iterator.lua();
				lua.push_value(index);
				lua.push_nil();
				lua.push_nil();
				assert!(lua.is_table(-3), $assert_err, index);

				iterator
			}

			/// Grants special access to the Lua state for the iterator. Do not expose to the public API!
			const fn lua(&self) -> LuaState {
				// SAFETY: self.lua is always valid and the originating LuaState should be borrowed.
				// While usually this would not be safe enough (as the user can still *read* from the state),
				// this iterator does not interact with elements it has not pushed itself and said elements are all referenced in a table.
				unsafe { LuaState::from_ptr(self.lua) }
			}

			$($($impl_fn)*)?
		}

		impl Iterator for $iterator<'_> {
			type Item = (i32, i32);

			$($iter_fn)+
		}

		impl Drop for $iterator<'_> {
			#[inline]
			fn drop(&mut self) {
				// SAFETY: The stack should always hold at least 3 elements until now.
				unsafe { self.lua().pop_unchecked(3); }
			}
		}
	};
}

/// Iterator over all key-value pairs of a specified table.
///
/// This iterator internally calls [`LuaState::next`] for you,
/// returning for each iteration the indices to the key and value, respectively.
/// (Currently, the returned indices are always `-2` and `-1`, this may change in the future.)
///
/// Once created, three elements will be pushed onto the stack *(the table, the key, and the value)*.
/// Each iteration overwrites the previous key-value pair's elements with the new pair's.
/// All three elements will be automatically popped once the iterator is *dropped*.
///
/// After traversing the entire table, the iterator can be reused to reiterate over the entire table again and again if needed.
///
/// Multiple `Pairs` (and [`IPairs`]) may exist simultaneously for nested iterations,
/// but the iterators should be used and dropped in the order in which they were created:
/// doing otherwise is considered a logic error and will result in unexpected behavior *(not UB)*.
///
/// Finishing the loop without properly dropping the iterator is also considered a logic error (as the elements would not be popped).
///
/// Lastly, getting the key's value using a function that can replace the number stored in the stack with a string
/// (like [`LuaState::to_string`]) may cause a logic error if said replacement occurs.
#[repr(transparent)]
#[derive(Debug)]
#[must_use = "iterators are lazy and do nothing unless consumed"]
pub struct Pairs<'a> {
	lua: *mut ffi::lua_State,
	_marker: PhantomData<&'a LuaState>,
}

iterator_impl!(
	Pairs,
	"cannot iterate over all key-value pairs of the element at index {} as it is not a table";

	#[inline]
	fn next(&mut self) -> Option<Self::Item> {
		let mut lua = self.lua();
		// SAFETY: The stack should always hold at least 3 elements until this Pairs is dropped.
		unsafe { lua.pop_unchecked(1); }

		// SAFETY: The element at index -2 should always be a table.
		if unsafe { lua.next_unchecked(-2) } {
			Some((-2, -1))
		} else {
			lua.push_nil();
			lua.push_nil();
			None
		}
	}
);

/// Iterator over the indexed key-value pairs of a specified table.
///
/// This iterator returns for each iteration the indices to the key and value, respectively.
/// (Currently, the returned indices are always `-2` and `-1`, this may change in the future.)
/// Just like Lua's `ipairs` function, it starts at key `1` and increments it for each iteration;
/// stopping once the key does not contain a value in the table.
///
/// Once created, three elements will be pushed onto the stack *(the table, the value, and the key)*.
/// Each iteration overwrites the previous key-value pair's elements with the new pair's.
/// All three elements will be automatically popped once the iterator is *dropped*.
///
/// Calling `next()` after the iterator ended will still cause it to increment the key and read the next key-value pair,
/// which could be used to skip over gaps.
///
/// Multiple `IPairs` (and [`Pairs`]) may exist simultaneously for nested iterations,
/// but the iterators should be used and dropped in the order in which they were created:
/// doing otherwise is considered a logic error and will result in unexpected behavior *(not UB)*.
///
/// Finishing the loop without properly dropping the iterator is also considered a logic error (as the elements would not be popped).
#[repr(transparent)]
#[derive(Debug)]
#[must_use = "iterators are lazy and do nothing unless consumed"]
pub struct IPairs<'a> {
	lua: *mut ffi::lua_State,
	_marker: PhantomData<&'a LuaState>,
}

iterator_impl!(
	IPairs,
	"cannot iterate over the indexed key-value pairs of the element at index {} as it is not a table";

	#[inline]
	#[allow(clippy::cast_possible_truncation)]
	fn next(&mut self) -> Option<Self::Item> {
		let mut lua = self.lua();
		let i = lua.to_integer(-2) + 1;
		// SAFETY: The stack should always hold at least 3 elements until this IPairs is dropped.
		unsafe { lua.pop_unchecked(2); }

		lua.push_integer(i);
		// SAFETY: lua.as_ptr() is always valid and the element at index -2 should always be a table.
		unsafe { ffi::lua_rawgeti(lua.as_ptr(), -2, i as i32); }
		lua.is_nil(-1).not().then_some((-2, -1))
	}

	#[inline]
	fn size_hint(&self) -> (usize, Option<usize>) {
		(0, Some(self.lua().object_length(-3)))
	}
);

impl IPairs<'_> {
	/// Returns the index that the iterator has last read.
	///
	/// If the iterator has not done its first iteration yet, `0` will be returned.
	#[inline]
	#[must_use]
	pub fn current_index(&self) -> usize {
		self.lua().to_integer(-2).cast_unsigned()
	}
}

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

	/// Creates a new empty table and pushes it onto the stack.
	///
	/// To pre-allocate elements when creating a new table, see [`LuaState::create_table`].
	#[inline]
	pub fn new_table(&mut self) {
		// SAFETY: self.as_ptr() is always valid.
		unsafe { ffi::lua_newtable(self.as_ptr()); }
	}

	/// Creates a new empty table and pushes it onto the stack.
	///
	/// The table will have `n_arr` array elements and `n_rec` non-array elements pre-allocated.
	///
	/// To create a new table without pre-allocating, see [`LuaState::new_table`].
	///
	/// # Lua Errors
	///
	/// Calling this function way too many times may trigger a stack overflow.
	///
	/// Incredibly high values for `n_arr` or `n_rec` may instantly trigger a table overflow.
	#[inline]
	pub fn create_table(&mut self, n_arr: i32, n_rec: i32) {
		// SAFETY: self.as_ptr() is always valid.
		unsafe { ffi::lua_createtable(self.as_ptr(), n_arr, n_rec); }
	}

	/// Does the equivalent of the Lua code `t[key]`, where `t` is the value located at the given valid `index`.
	///
	/// The result is pushed onto the stack (even if `nil`).
	///
	/// For a function that does the same but allows any value instead of `key`, see [`LuaState::get_table`].
	///
	/// # Lua Errors
	///
	/// Not every value `t` can be indexed: for example, if `t` is `nil` this function will error.
	///
	/// Indexing `t` may also trigger the `__index` metamethod which could error.
	#[inline]
	pub fn get_field(&mut self, index: i32, key: &CStr) {
		// SAFETY: self.as_ptr() is always valid and key's pointer is valid for the duration of this function.
		unsafe { ffi::lua_getfield(self.as_ptr(), index, key.as_ptr()); }
	}

	/// Does the equivalent of the Lua code `t[key] = v`, where `t` is the value located at the given valid `index`
	/// and `v` is the value at the top of the stack (which is popped).
	///
	/// For a function that does the same but allows any value instead of `key`, see [`LuaState::set_table`].
	///
	/// # Panics
	///
	/// Panics if the stack is empty.
	///
	/// # Lua Errors
	///
	/// Not every value `t` can be indexed: for example, if `t` is `nil` this function will error.
	///
	/// Indexing `t` may also trigger the `__newindex` metamethod which could error.
	#[inline]
	pub fn set_field(&mut self, index: i32, key: &CStr) {
		assert!(self.get_top() > 0, "tried to set field {key:?} for the table at index {index} despite having an empty stack");

		// SAFETY: self.as_ptr() is always valid, key's pointer is valid for the duration of this function, and the stack is not empty.
		unsafe { ffi::lua_setfield(self.as_ptr(), index, key.as_ptr()); }
	}

	/// Gets the global named `name` and pushes it onto the stack (even if `nil`).
	///
	/// # Lua Errors
	///
	/// Indexing the global environment may trigger the `__index` metamethod which could error.
	#[inline]
	pub fn get_global(&mut self, name: &CStr) {
		// SAFETY: self.as_ptr() is always valid and name's pointer is valid for the duration of this function.
		unsafe { ffi::lua_getglobal(self.as_ptr(), name.as_ptr()); }
	}

	/// Pops the top element and sets it as the value of the global named `name`.
	///
	/// # Panics
	///
	/// Panics if the stack is empty.
	///
	/// # Lua Errors
	///
	/// Indexing the global environment may trigger the `__newindex` metamethod which could error.
	#[inline]
	pub fn set_global(&mut self, name: &CStr) {
		assert!(self.get_top() > 0, "tried to set global {name:?} despite having an empty stack");

		// SAFETY: self.as_ptr() and name's pointer are always valid and the stack is not empty.
		unsafe { ffi::lua_setglobal(self.as_ptr(), name.as_ptr()); }
	}

	/// Does the equivalent of the Lua code `t[k]`, where `t` is the value located at the given valid `index`
	/// and `k` is the value at the top of the stack (which is popped).
	///
	/// The result is pushed onto the stack, replacing `k` (even if `nil`).
	///
	/// # Panics
	///
	/// Panics if the stack is empty.
	///
	/// # Lua Errors
	///
	/// Not every value `t` can be indexed: for example, if `t` is `nil` this function will error.
	///
	/// Indexing `t` may also trigger the `__index` metamethod which could error.
	#[inline]
	pub fn get_table(&mut self, index: i32) {
		assert!(self.get_top() > 0, "tried to get a value from the table at index {index} despite having an empty stack");

		// SAFETY: self.as_ptr() is always valid and the stack is not empty.
		unsafe { ffi::lua_gettable(self.as_ptr(), index); }
	}

	/// Does the equivalent of the Lua code `t[k] = v`, where `t` is the value located at the given valid `index`
	/// and `v` and `k` are the values at the top of the stack (which are popped from the stack in said order).
	///
	/// # Panics
	///
	/// Panics if the stack does not contain both `v` and `k` (thus contains less than 2 elements).
	///
	/// # Lua Errors
	///
	/// Not every value `t` or `k` can be used: if `t` or `k` are `nil` this function will error.
	///
	/// Indexing `t` may also trigger the `__newindex` metamethod which could error.
	#[inline]
	pub fn set_table(&mut self, index: i32) {
		assert!(2 <= self.get_top(), "tried to set a value for the table at index {index} despite not having `v` and `k` in the stack");

		// SAFETY: self.as_ptr() is always valid and v and k are in the stack.
		unsafe { ffi::lua_settable(self.as_ptr(), index); }
	}

	/// Accesses a table similarly to [`LuaState::get_table`], but without triggering any metamethod (see [`LuaState::get_table`]).
	///
	/// # Panics
	///
	/// Panics if the element at the given valid `index` is not a table or if the stack is empty.
	#[inline]
	pub fn raw_get(&mut self, index: i32) {
		assert!(self.is_table(index), "cannot raw get a value from the element at index {index} as it is not a table");
		assert!(self.get_top() > 0, "tried to raw get a value from the table at index {index} despite having an empty stack");

		// SAFETY: self.as_ptr() is always valid, the element specified is a table and the stack is not empty.
		unsafe { ffi::lua_rawget(self.as_ptr(), index); }
	}

	/// Accesses a table similarly to [`LuaState::set_table`], but without triggering any metamethod (see [`LuaState::set_table`]).
	///
	/// # Panics
	///
	/// Panics if the element at the given valid `index` is not a table or if the stack is empty.
	///
	/// # Lua errors
	///
	/// Not every value `k` can be used: if `k` is `nil` this function will error.
	#[inline]
	pub fn raw_set(&mut self, index: i32) {
		assert!(self.is_table(index), "cannot raw set a value for the element at index {index} as it is not a table");
		assert!(2 <= self.get_top(), "tried to raw set a value for the table at index {index} despite not having `v` and `k` in the stack");

		// SAFETY: self.as_ptr() is always valid, the element specified is a table and v and k are in the stack.
		unsafe { ffi::lua_rawset(self.as_ptr(), index); }
	}

	/// Does the equivalent of the Lua code `t[i]`, where `t` is the value located at the given valid `index`.
	///
	/// The result is pushed onto the stack.
	///
	/// # Panics
	///
	/// Panics if the element at the given valid `index` is not a table.
	#[inline]
	pub fn raw_geti(&mut self, index: i32, i: i32) {
		assert!(self.is_table(index), "cannot raw get an element from the element at index {index} as it is not a table");

		// SAFETY: self.as_ptr() is always valid and the element specified is a table.
		unsafe { ffi::lua_rawgeti(self.as_ptr(), index, i); }
	}

	/// Does the equivalent of the Lua code `t[i] = v`, where `t` is the value located at the given valid `index`
	/// and `v` is the value at the top of the stack (which is popped).
	///
	/// # Panics
	///
	/// Panics if the element at the given valid `index` is not a table or if the stack is empty.
	#[inline]
	pub fn raw_seti(&mut self, index: i32, i: i32) {
		assert!(self.is_table(index), "cannot raw set a value for the element at index {index} as it is not a table");
		assert!(self.get_top() > 0, "tried to raw set a value for the table at index {index} despite having an empty stack");

		// SAFETY: self.as_ptr() is always valid, the element specified is a table and the stack is not empty.
		unsafe { ffi::lua_rawseti(self.as_ptr(), index, i); }
	}

	/// Pops the last element of the stack, treating it as the key,
	/// and pushes onto the stack the *next* key-value pair from the table at the given `index`, returning `true`.
	///
	/// If there is no next key-value pair, then nothing is pushed and `false` is returned instead.
	///
	/// This function can be used to traverse all elements of a table, see the example below.
	///
	/// # Panics
	///
	/// Panics if the element at the given `index` is not a table.
	///
	/// # Lua Errors
	///
	/// If the given key is neither `nil` nor present in the table, this function will error.
	///
	/// # Examples
	///
	/// ```
	/// # use low_luajit::{LuaState, AmountReturned};
	/// # extern "C-unwind" fn example(mut lua: LuaState) -> AmountReturned {
	/// # let table_index = 1;
	/// // Assume the table is located in the stack at index t.
	/// lua.push_nil(); // Push the first (dummy) key.
	/// while lua.next(table_index) { // Get the next key-value pair.
	///     // (The key is located at index -2, while the value is located at index -1.)
	///     // Use the key-value pair in some way...
	///     lua.pop(1); // Pop the value, so that the key can be used to get the next key-value pair.
	/// }
	/// # AmountReturned::NONE }
	/// ```
	#[inline]
	pub fn next(&mut self, index: i32) -> bool {
		assert!(self.is_table(index), "cannot get the next key-value pair of the element at index {index} as it is not a table");

		// SAFETY: The element at the given index is a table.
		unsafe { self.next_unchecked(index) }
	}

	/// Behaves exactly as [`LuaState::next`], but without doing any checks.
	///
	/// # Safety
	///
	/// The element at the given `index` must be a table.
	#[inline]
	pub unsafe fn next_unchecked(&mut self, index: i32) -> bool {
		// SAFETY: The caller must ensure the element at the given index is a table.
		unsafe { ffi::lua_next(self.as_ptr(), index) != 0 }
	}

	/// Returns an iterator over all key-value pairs of the table at the given `index`,
	/// mimicking the behavior of Lua's `pairs` function.
	///
	/// For each iteration, the indices of the current key-value pair are returned.
	///
	/// *[See also the `Pairs` struct](struct.Pairs.html)*.
	///
	/// # Panics
	///
	/// Panics if the element at the given `index` is not a table.
	///
	/// # Examples
	///
	/// Printing the Lua standard library:
	///
	/// ```
	/// use low_luajit::LuaState;
	///
	/// # let mut lua = unsafe { low_luajit::__test!(luaL_openlibs) };
	/// for (k, v) in lua.pairs(LuaState::GLOBALS_INDEX) {
	///     println!("{}: {}", lua.display(k), lua.display(v));
	///
	///     if lua.is_table(v) {
	///         for (k, v) in lua.pairs(v) {
	///             println!("\t{}: {}", lua.display(k), lua.display(v));
	///         }
	///     }
	/// }
	/// ```
	#[inline]
	pub fn pairs(&self, index: i32) -> Pairs<'_> {
		Pairs::new(self, index)
	}

	/// Returns an iterator over the indexed key-value pairs of the table at the given `index`,
	/// mimicking the behavior of Lua's `ipairs` function.
	///
	/// For each iteration, the indices of the current key-value pair are returned.
	///
	/// *[See also the `IPairs` struct](struct.IPairs.html)*.
	///
	/// # Panics
	///
	/// Panics if the element at the given `index` is not a table.
	#[inline]
	pub fn ipairs(&self, index: i32) -> IPairs<'_> {
		IPairs::new(self, index)
	}

	/// Gets the metatable of the value at the given acceptable `index` and pushes it onto the stack, returning `true`.
	///
	/// If `index` is invalid or the value has no metatable, nothing is pushed and the function returns `false`.
	#[inline]
	pub fn get_metatable(&mut self, index: i32) -> bool {
		// SAFETY: self.as_ptr() is always valid.
		unsafe { ffi::lua_getmetatable(self.as_ptr(), index) != 0 }
	}

	/// Pops the top elemenet and sets it as the metatable of the value at the given acceptable `index`.
	///
	/// This function always returns `true`.
	///
	/// # Panics
	///
	/// Panics if the top element is not a table or if the stack is empty.
	#[inline]
	pub fn set_metatable(&mut self, index: i32) -> bool {
		assert!(self.get_top() > 0, "tried to set the metatable for the element at index {index} despite having an empty stack");
		assert!(self.is_table(-1), "cannot set the metatable for the element at index {index}, as the given metatable is not a table");

		// SAFETY: self.as_ptr() is always valid, the stack is not empty and the top element is a table.
		unsafe { ffi::lua_setmetatable(self.as_ptr(), index) != 0 }
	}

	/// Gets the environment table of the value at the given `index` and pushes it onto the stack.
	///
	/// If `index` is invalid or the value has no environment table, `nil` is pushed onto the stack.
	#[inline]
	pub fn get_fenv(&mut self, index: i32) {
		// SAFETY: self.as_ptr() is always valid.
		unsafe { ffi::lua_getfenv(self.as_ptr(), index); }
	}

	/// Pops a table from the stack and sets it as the environment table of the value at the given `index`, returning `true`.
	///
	/// Only functions, userdata and threads may have an environment table.
	///
	/// If `index` is invalid or the value's type is not one listed above, the stack is still popped but the function returns `false`.
	///
	/// # Panics
	///
	/// Panics if the top element is not a table or if the stack is empty.
	#[inline]
	pub fn set_fenv(&mut self, index: i32) -> bool {
		assert!(self.get_top() > 0, "tried to set the environment table for the element at index {index} despite having an empty stack");
		assert!(self.is_table(-1), "cannot set the environment table for the element at index {index}, as the given environment table is not a table");

		// SAFETY: self.as_ptr() is always valid, the stack is not empty and the top element is a table.
		unsafe { ffi::lua_setfenv(self.as_ptr(), index) != 0 }
	}

	/// Calls the metamethod `method` from the metatable of the element at the given `index`.
	///
	/// The metamethod will receive the element itself as its only argument and its result will be pushed onto the stack, returning `true`.
	///
	/// If there is no metatable or no metamethod, nothing is pushed and this function returns `false` instead.
	///
	/// # Lua Errors
	///
	/// The called metamethod may error.
	#[inline]
	pub fn call_meta(&mut self, index: i32, method: &CStr) -> bool {
		// SAFETY: self.as_ptr() is always valid and method's pointer is valid for the duration of this function.
		unsafe { ffi::luaL_callmeta(self.as_ptr(), index, method.as_ptr()) != 0 }
	}

	/// Searches for a table located at `path` within the table located at the given `index`, creating it if it doesn't exist.
	///
	/// `path` should be a list of keys separated by dots, these keys will be traversed starting from the given table.
	/// Any missing table (`nil`) will be created while traversing by calling [`LuaState::create_table`] using `size_hint` for `n_rec`
	/// *(only for the final table, any other table created will use `1` instead; `n_arr` is always `0`)*.
	///
	/// # Panics
	///
	/// Panics if the element at the given `index` is not a table.
	///
	/// # Errors
	///
	/// If any key in `path` leads to a value that is neither a table nor `nil`,
	/// this function will return `Err` with a substring of `path` containing the non-traversed keys.
	///
	/// # Examples
	///
	/// ```
	/// use low_luajit::LuaState;
	///
	/// # let mut lua = unsafe { low_luajit::__test!(luaL_openlibs) };
	/// assert_eq!(lua.find_table(LuaState::GLOBALS_INDEX, c"math", 0), Ok(()));
	/// assert_eq!(lua.find_table(LuaState::GLOBALS_INDEX, c"new.tables", 123), Ok(()));
	/// assert_eq!(lua.find_table(LuaState::GLOBALS_INDEX, c"math.pi", 0), Err(c"pi")); // Not a table!
	/// ```
	#[inline]
	pub fn find_table<'a>(&mut self, index: i32, path: &'a CStr, size_hint: i32) -> Result<(), &'a CStr> {
		assert!(self.is_table(index), "cannot find a table within the element at index {index} as it is not a table");

		// SAFETY: self.as_ptr() and path are always valid and luaL_findtable returns valid nul-terminated string pointers when not null.
		let fname = unsafe { ffi::luaL_findtable(self.as_ptr(), index, path.as_ptr(), size_hint) };
		map_ptr!(fname, unsafe { CStr::from_ptr(fname) }).map_or(Ok(()), Err)
	}
}