A bit more ewext refactoring

ewext/noita_api_macro/src/lib.rs

@@ -1,3 +1,5 @@
+use std::ffi::CString;
+
 use heck::ToSnekCase;
 use proc_macro::TokenStream;
 use quote::{format_ident, quote};
@@ -94,3 +96,22 @@ fn generate_code_for_component(com: Component) -> proc_macro2::TokenStream {
         }
     }
 }
+
+#[proc_macro]
+pub fn add_lua_fn(item: TokenStream) -> TokenStream {
+    let mut tokens = item.into_iter();
+
+    let fn_name = tokens.next().unwrap().to_string();
+    let fn_name_ident = format_ident!("{fn_name}");
+    let bridge_fn_name = format_ident!("{fn_name}_lua_bridge");
+    let fn_name_c = proc_macro2::Literal::c_string(CString::new(fn_name).unwrap().as_c_str());
+    quote! {
+        unsafe extern "C" fn #bridge_fn_name(lua: *mut lua_State) -> c_int {
+            #fn_name_ident(LuaState::new(lua)) as c_int
+        }
+
+        LUA.lua_pushcclosure(lua, Some(#bridge_fn_name), 0);
+        LUA.lua_setfield(lua, -2, #fn_name_c.as_ptr());
+    }
+    .into()
+}

ewext/src/addr_grabber.rs

@@ -3,9 +3,8 @@ use std::{mem, os::raw::c_void, ptr, sync::OnceLock};
 use iced_x86::{Decoder, DecoderOptions, Mnemonic};
 
 use crate::{
-    lua_bindings::{lua_State, LUA_GLOBALSINDEX},
+    lua_state::LuaState,
     noita::ntypes::{EntityManager, ThiscallFn},
-    LUA,
 };
 
 static GRABBED: OnceLock<Grabbed> = OnceLock::new();
@@ -54,42 +53,42 @@ pub(crate) struct GrabbedFns {
     pub(crate) get_entity: *const ThiscallFn, //unsafe extern "C" fn(*const EntityManager, u32) -> *mut Entity,
 }
 
-pub(crate) unsafe fn grab_addrs(lua: *mut lua_State) {
-    LUA.lua_getfield(lua, LUA_GLOBALSINDEX, c"GameGetWorldStateEntity".as_ptr());
-    let base = LUA.lua_tocfunction(lua, -1).unwrap() as *const c_void;
+pub(crate) fn grab_addrs(lua: LuaState) {
+    lua.get_global(c"GameGetWorldStateEntity");
+    let base = lua.to_cfunction(-1).unwrap() as *const c_void;
     let world_state_entity =
-        grab_addr_from_instruction(base, 0x007aa7ce - 0x007aa540, Mnemonic::Mov).cast();
+        unsafe { grab_addr_from_instruction(base, 0x007aa7ce - 0x007aa540, Mnemonic::Mov).cast() };
     println!(
         "World state entity addr: 0x{:x}",
         world_state_entity as usize
     );
-    // Pop the last element.
-    LUA.lua_settop(lua, -2);
+    lua.pop_last();
 
-    LUA.lua_getfield(lua, LUA_GLOBALSINDEX, c"GameGetFrameNum".as_ptr());
-    let base = LUA.lua_tocfunction(lua, -1).unwrap() as *const c_void;
+    lua.get_global(c"GameGetFrameNum");
+    let base = lua.to_cfunction(-1).unwrap() as *const c_void;
     let load_game_global =
-        grab_addr_from_instruction(base, 0x007bf3c9 - 0x007bf140, Mnemonic::Call); // CALL load_game_global
+        unsafe { grab_addr_from_instruction(base, 0x007bf3c9 - 0x007bf140, Mnemonic::Call) }; // CALL load_game_global
     println!("Load game global addr: 0x{:x}", load_game_global as usize);
-    let game_global =
-        grab_addr_from_instruction(load_game_global, 0x00439c17 - 0x00439bb0, Mnemonic::Mov).cast();
+    let game_global = unsafe {
+        grab_addr_from_instruction(load_game_global, 0x00439c17 - 0x00439bb0, Mnemonic::Mov).cast()
+    };
     println!("Game global addr: 0x{:x}", game_global as usize);
-    // Pop the last element.
-    LUA.lua_settop(lua, -2);
+    lua.pop_last();
 
-    LUA.lua_getfield(lua, LUA_GLOBALSINDEX, c"EntityGetFilename".as_ptr());
-    let base = LUA.lua_tocfunction(lua, -1).unwrap() as *const c_void;
-    let get_entity = mem::transmute_copy(&grab_addr_from_instruction(
-        base,
-        0x0079782b - 0x00797570,
-        Mnemonic::Call,
-    ));
+    lua.get_global(c"EntityGetFilename");
+    let base = lua.to_cfunction(-1).unwrap() as *const c_void;
+    let get_entity = unsafe {
+        mem::transmute_copy(&grab_addr_from_instruction(
+            base,
+            0x0079782b - 0x00797570,
+            Mnemonic::Call,
+        ))
+    };
     println!("get_entity addr: 0x{:x}", get_entity as usize);
     let entity_manager =
-        grab_addr_from_instruction(base, 0x00797821 - 0x00797570, Mnemonic::Mov).cast();
+        unsafe { grab_addr_from_instruction(base, 0x00797821 - 0x00797570, Mnemonic::Mov).cast() };
     println!("entity_manager addr: 0x{:x}", entity_manager as usize);
-    // Pop the last element.
-    LUA.lua_settop(lua, -2);
+    lua.pop_last();
 
     GRABBED
         .set(Grabbed {

ewext/src/lib.rs

@@ -9,6 +9,7 @@ use addr_grabber::{grab_addrs, grabbed_fns, grabbed_globals};
 use lua_bindings::{lua_State, Lua51};
 use lua_state::LuaState;
 use noita::{ntypes::Entity, NoitaPixelRun, ParticleWorldState};
+use noita_api_macro::add_lua_fn;
 
 mod lua_bindings;
 mod lua_state;
@@ -40,13 +41,11 @@ struct ExtState {
     saved_world_state: Option<SavedWorldState>,
 }
 
-// const EWEXT: [(&'static str, Function); 1] = [("testfn", None)];
-
-unsafe extern "C" fn init_particle_world_state(lua: *mut lua_State) -> c_int {
+fn init_particle_world_state(lua: LuaState) -> c_int {
     println!("\nInitializing particle world state");
-    let world_pointer = unsafe { LUA.lua_tointeger(lua, 1) };
-    let chunk_map_pointer = unsafe { LUA.lua_tointeger(lua, 2) };
-    let material_list_pointer = unsafe { LUA.lua_tointeger(lua, 3) };
+    let world_pointer = lua.to_integer(1);
+    let chunk_map_pointer = lua.to_integer(2);
+    let material_list_pointer = lua.to_integer(3);
     println!("pws stuff: {world_pointer:?} {chunk_map_pointer:?}");
 
     STATE.with(|state| {
@@ -60,7 +59,8 @@ unsafe extern "C" fn init_particle_world_state(lua: *mut lua_State) -> c_int {
     0
 }
 
-unsafe extern "C" fn encode_area(lua: *mut lua_State) -> c_int {
+fn encode_area(lua: LuaState) -> c_int {
+    let lua = lua.raw();
     let start_x = unsafe { LUA.lua_tointeger(lua, 1) } as i32;
     let start_y = unsafe { LUA.lua_tointeger(lua, 2) } as i32;
     let end_x = unsafe { LUA.lua_tointeger(lua, 3) } as i32;
@@ -110,10 +110,9 @@ unsafe extern "C" fn load_world_state_lua(_lua: *mut lua_State) -> i32 {
     0
 }
 
-unsafe extern "C" fn make_ephemerial(lua: *mut lua_State) -> c_int {
-    let lua_state = LuaState::new(lua);
+fn make_ephemerial(lua: LuaState) -> c_int {
     unsafe {
-        let entity_id = lua_state.to_integer(1) as u32;
+        let entity_id = lua.to_integer(1) as u32;
 
         let entity_manager = grabbed_globals().entity_manager.read();
         let mut entity: *mut Entity;
@@ -128,13 +127,14 @@ unsafe extern "C" fn make_ephemerial(lua: *mut lua_State) -> c_int {
             out("ecx") _,
             out("eax") entity,
         );
-        // let entity = (state.fns.as_ref().unwrap().get_entity)(entity_manager, entity_id);
-        entity.cast::<c_void>().offset(0x8).cast::<u32>().write(0);
+        if !entity.is_null() {
+            entity.cast::<c_void>().offset(0x8).cast::<u32>().write(0);
+        }
     }
     0
 }
 
-unsafe extern "C" fn on_world_initialized(lua: *mut lua_State) -> c_int {
+fn on_world_initialized(lua: LuaState) -> i32 {
     grab_addrs(lua);
     0
 }
@@ -148,18 +148,14 @@ pub unsafe extern "C" fn luaopen_ewext0(lua: *mut lua_State) -> c_int {
     unsafe {
         LUA.lua_createtable(lua, 0, 0);
 
-        LUA.lua_pushcclosure(lua, Some(init_particle_world_state), 0);
-        LUA.lua_setfield(lua, -2, c"init_particle_world_state".as_ptr());
-        LUA.lua_pushcclosure(lua, Some(encode_area), 0);
-        LUA.lua_setfield(lua, -2, c"encode_area".as_ptr());
+        add_lua_fn!(init_particle_world_state);
+        add_lua_fn!(encode_area);
         LUA.lua_pushcclosure(lua, Some(load_world_state_lua), 0);
         LUA.lua_setfield(lua, -2, c"load_world_state".as_ptr());
         LUA.lua_pushcclosure(lua, Some(save_world_state_lua), 0);
         LUA.lua_setfield(lua, -2, c"save_world_state".as_ptr());
-        LUA.lua_pushcclosure(lua, Some(make_ephemerial), 0);
-        LUA.lua_setfield(lua, -2, c"make_ephemerial".as_ptr());
-        LUA.lua_pushcclosure(lua, Some(on_world_initialized), 0);
-        LUA.lua_setfield(lua, -2, c"on_world_initialized".as_ptr());
+        add_lua_fn!(make_ephemerial);
+        add_lua_fn!(on_world_initialized);
     }
     println!("Initializing ewext - Ok");
     1

ewext/src/lua_state.rs

@@ -1,14 +1,37 @@
-use crate::{lua_bindings::lua_State, LUA};
+use std::ffi::CStr;
+
+use crate::{
+    lua_bindings::{lua_CFunction, lua_State, LUA_GLOBALSINDEX},
+    LUA,
+};
 
 #[derive(Clone, Copy)]
-pub(crate) struct LuaState(*mut lua_State);
+pub(crate) struct LuaState {
+    lua: *mut lua_State,
+}
 
 impl LuaState {
     pub(crate) fn new(lua: *mut lua_State) -> Self {
-        Self(lua)
+        Self { lua }
+    }
+
+    pub(crate) fn raw(&self) -> *mut lua_State {
+        self.lua
     }
 
     pub(crate) fn to_integer(&self, index: i32) -> isize {
-        unsafe { LUA.lua_tointeger(self.0, index) }
+        unsafe { LUA.lua_tointeger(self.lua, index) }
+    }
+
+    pub(crate) fn to_cfunction(&self, index: i32) -> lua_CFunction {
+        unsafe { LUA.lua_tocfunction(self.lua, index) }
+    }
+
+    pub(crate) fn get_global(&self, name: &CStr) {
+        unsafe { LUA.lua_getfield(self.lua, LUA_GLOBALSINDEX, name.as_ptr()) };
+    }
+
+    pub(crate) fn pop_last(&self) {
+        unsafe { LUA.lua_settop(self.lua, -2) };
     }
 }