NONE - LuaJIT frames and other improvements

NONE - LuaJIT frames and other improvements

Author: yogwoggf

packages/autorun-luajit/src/helpers.rs

@@ -60,3 +60,8 @@ pub fn push_tvalue(state: &mut LJState, tvalue: &TValue) {
 		state.top = state.top.add(1);
 	}
 }
+
+pub fn push_frame_func(state: &mut LJState, frame: &Frame) -> anyhow::Result<()> {
+	push_tvalue(state, unsafe { &*frame.get_func_tv() });
+	Ok(())
+}

packages/autorun-luajit/src/types/common.rs

@@ -2,13 +2,16 @@
 
 use anyhow::Context;
 use std::ffi::{c_int, c_void};
+use std::fmt::Debug;
 
 // IMPORTANT: GMod's LUA_IDSIZE was randomly changed to 128 instead of 60 like in vanilla LuaJIT
 #[cfg(feature = "gmod")]
 pub const LUA_IDSIZE: i32 = 128;
 #[cfg(not(feature = "gmod"))]
 pub const LUA_IDSIZE: i32 = 60;
 
+pub const LJ_FR2: u32 = 1;
+
 pub const LJ_TNIL: u32 = !0u32;
 pub const LJ_TFALSE: u32 = !1u32;
 pub const LJ_TTRUE: u32 = !2u32;
@@ -79,6 +82,12 @@ pub struct GCRef {
 }
 
 impl GCRef {
+	pub fn from_ptr<T>(ptr: *mut T) -> Self {
+		Self {
+			gcptr64: (ptr as u64) & LJ_GCVMASK,
+		}
+	}
+
 	// equivalent to the gcref macro in LuaJIT
 	pub fn as_ptr<T>(&self) -> *mut T {
 		self.gcptr64 as *mut T
@@ -112,6 +121,22 @@ pub union TValue {
 	pub ftsz: u64,
 }
 
+impl Debug for TValue {
+	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+		write!(f, "TValue {{ itype: {}, gcr: {:x} }}", self.itype(), unsafe {
+			self.gcr.gcptr64
+		})
+	}
+}
+
+macro_rules! impl_tvalue_type_check {
+	($function_name:ident, $lj_type_const:expr) => {
+		pub fn $function_name(&self) -> bool {
+			self.itype() == $lj_type_const
+		}
+	};
+}
+
 impl TValue {
 	pub fn as_ptr<T: IntoLJType>(&self) -> anyhow::Result<*mut T> {
 		if self.itype() != T::LJ_TYPE {
@@ -132,6 +157,21 @@ impl TValue {
 	pub fn itype(&self) -> u32 {
 		unsafe { ((self.it64 >> 47) & 0xFFFFFFFF) as u32 }
 	}
+
+	impl_tvalue_type_check!(is_nil, LJ_TNIL);
+	impl_tvalue_type_check!(is_false, LJ_TFALSE);
+	impl_tvalue_type_check!(is_true, LJ_TTRUE);
+	impl_tvalue_type_check!(is_lightud, LJ_TLIGHTUD);
+	impl_tvalue_type_check!(is_str, LJ_TSTR);
+	impl_tvalue_type_check!(is_upval, LJ_TUPVAL);
+	impl_tvalue_type_check!(is_thread, LJ_TTHREAD);
+	impl_tvalue_type_check!(is_proto, LJ_TPROTO);
+	impl_tvalue_type_check!(is_func, LJ_TFUNC);
+	impl_tvalue_type_check!(is_trace, LJ_TTRACE);
+	impl_tvalue_type_check!(is_cdata, LJ_TCDATA);
+	impl_tvalue_type_check!(is_tab, LJ_TTAB);
+	impl_tvalue_type_check!(is_udata, LJ_TUDATA);
+	impl_tvalue_type_check!(is_numx, LJ_TNUMX);
 }
 
 #[repr(C, packed)]
@@ -160,6 +200,17 @@ pub struct GCfuncL {
 	pub uvptr: [GCRef; 1],
 }
 
+impl GCfuncL {
+	pub fn get_proto(&self) -> anyhow::Result<&GCProto> {
+		let pc_ref = self.header.pc;
+		// proto starts immediately before the pc pointer
+		let proto = unsafe { pc_ref.as_ptr::<GCProto>().offset(-1) };
+		let proto_ref = unsafe { proto.as_ref().context("Failed to dereference GCProto from GCfuncL")? };
+
+		Ok(proto_ref)
+	}
+}
+
 #[repr(C)]
 #[derive(Clone, Copy)]
 pub union GCfunc {
@@ -349,7 +400,15 @@ impl IntoLJType for GCUpval {
 	const LJ_TYPE: u32 = LJ_TUPVAL;
 }
 
-pub type BCIns = u32;
+#[derive(Debug, Clone, Copy)]
+pub struct BCIns(u32);
+
+impl BCIns {
+	pub fn a(&self) -> u8 {
+		//#define bc_a(i)		((BCReg)(((i)>>8)&0xff))
+		((self.0 >> 8) & 0xff) as u8
+	}
+}
 
 pub type BCLine = u32;
 

packages/autorun-luajit/src/types/frame.rs

@@ -1,9 +1,11 @@
 // Helper module for managing the frame stack in LuaJIT, which is highly complicated and fragile.
-use crate::{GCfunc, LJState, TValue};
+use crate::{BCIns, GCRef, GCfunc, LJ_FR2, LJState, TValue};
 
 pub const FRAME_TYPE: u8 = 3;
 pub const FRAME_P: u8 = 4;
 pub const FRAME_TYPEP: u8 = FRAME_TYPE | FRAME_P;
+
+#[derive(Debug, Copy, Clone, PartialEq, Eq)]
 pub enum FrameType {
 	Lua = 0,
 	C = 1,
@@ -14,26 +16,65 @@ pub enum FrameType {
 	FfPcallWithHook = 7,
 }
 
-#[derive(Debug)]
+impl std::fmt::Display for FrameType {
+	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+		let type_str = match self {
+			FrameType::Lua => "Lua",
+			FrameType::C => "C",
+			FrameType::Continuation => "Continuation",
+			FrameType::LuaVararg => "LuaVararg",
+			FrameType::Cpcall => "Cpcall",
+			FrameType::FfPcall => "FfPcall",
+			FrameType::FfPcallWithHook => "FfPcallWithHook",
+		};
+
+		write!(f, "{}", type_str)
+	}
+}
+
+#[derive(Debug, Copy, Clone)]
 /// Frames are stored in the LJ state as TValues with their payloads below them in the stack.
 /// This struct helps manage and interpret these frames.
 pub struct Frame {
 	pub tvalue: *mut TValue,
-	pub size: u32,
+	payload_copy: TValue,
 }
 
 impl Frame {
-	pub fn new(tvalue: *mut TValue, size: u32) -> Self {
-		Self { tvalue, size }
+	pub fn new(tvalue: *mut TValue) -> Self {
+		Self {
+			tvalue,
+			payload_copy: unsafe { *tvalue.offset(-1) },
+		}
 	}
 
 	pub fn from_debug_ci(state: &mut LJState, i_ci: i32) -> Self {
 		let offset = (i_ci as u32) & 0xffff;
-		let size = (i_ci as u32) >> 16;
 		let tvstack = state.stack.as_ptr::<TValue>();
 		let frametv = unsafe { tvstack.add(offset as usize) };
 
-		Frame::new(frametv, size)
+		frametv.into()
+	}
+
+	pub fn walk_stack(state: *mut LJState) -> Vec<Frame> {
+		let mut frames = Vec::new();
+
+		// traversal is in backwards order
+		let start = unsafe { (*state).stack.as_ptr::<TValue>().add(LJ_FR2 as usize) };
+		let mut frame = unsafe { (*state).base.offset(-1) };
+
+		while frame > start {
+			let current_frame: Frame = frame.into();
+			frames.push(current_frame);
+
+			if current_frame.is_lua_frame() {
+				frame = current_frame.get_previous_lua_frame().tvalue;
+			} else {
+				frame = current_frame.get_previous_delta_frame().tvalue;
+			}
+		}
+
+		frames
 	}
 
 	pub fn get_type(&self) -> FrameType {
@@ -51,14 +92,14 @@ impl Frame {
 	}
 
 	pub fn is_lua_frame(&self) -> bool {
-		matches!(self.get_type(), FrameType::Lua | FrameType::LuaVararg)
+		self.get_type() == FrameType::Lua
 	}
 
 	pub fn is_c_frame(&self) -> bool {
-		matches!(self.get_type(), FrameType::C | FrameType::Cpcall)
+		self.get_type() == FrameType::C
 	}
 
-	pub fn get_gc_func(&mut self) -> anyhow::Result<&mut GCfunc> {
+	pub fn get_gc_func(&self) -> anyhow::Result<&mut GCfunc> {
 		unsafe {
 			let func_tv = self.tvalue.offset(-1);
 			let gcfunc = (*func_tv).as_mut::<GCfunc>()?;
@@ -69,4 +110,54 @@ impl Frame {
 	pub fn get_func_tv(&self) -> *mut TValue {
 		unsafe { self.tvalue.offset(-1) }
 	}
+
+	fn overwrite_payload_gcr(&mut self, new_gcr: GCRef) {
+		unsafe {
+			(*self.tvalue.offset(-1)).gcr = new_gcr;
+		}
+	}
+
+	fn restore_payload(&mut self) {
+		unsafe {
+			*self.tvalue.offset(-1) = self.payload_copy;
+		}
+	}
+
+	pub fn mark_as_dummy_frame(&mut self, state: *mut LJState) {
+		// We can overwrite the payload to point to the Lua state,
+		// which is a special GCRef that indicates a dummy frame.
+		self.overwrite_payload_gcr(GCRef::from_ptr(state));
+	}
+
+	pub fn restore_from_dummy_frame(&mut self) {
+		self.restore_payload();
+	}
+
+	pub fn get_pc(&self) -> *const BCIns {
+		unsafe { (*self.tvalue).ftsz as *const BCIns }
+	}
+
+	pub fn get_previous_lua_frame(&self) -> Self {
+		let bc_ins_ptr = self.get_pc();
+		let bc_ins = unsafe { bc_ins_ptr.offset(-1).read_unaligned() };
+		let bc_a = bc_ins.a();
+		let offset = (1 + LJ_FR2 + (bc_a as u32)) as isize;
+
+		unsafe { self.tvalue.offset(-offset) }.into()
+	}
+
+	pub fn get_delta_size(&self) -> u64 {
+		unsafe { (*self.tvalue).ftsz & !(FRAME_TYPEP as u64) }
+	}
+
+	pub fn get_previous_delta_frame(&self) -> Self {
+		let size = self.get_delta_size() as usize;
+		unsafe { self.tvalue.byte_sub(size).into() }
+	}
+}
+
+impl Into<Frame> for *mut TValue {
+	fn into(self) -> Frame {
+		Frame::new(self)
+	}
 }

services/autorun-lib/src/lib.rs

@@ -25,6 +25,8 @@ pub fn main() -> anyhow::Result<()> {
 #[cfg(target_os = "windows")]
 #[unsafe(no_mangle)]
 extern "C" fn autorun_entrypoint() {
+	// Yield to allow gmod to initialize properly
+	std::thread::sleep(std::time::Duration::from_secs(1));
 	// Redirect stdout to stderr.
 	// This is a hack because for some reason stdout isn't intercepted on windows?
 	// Might be gmod's fault. I don't care.