Refactor ewext a bit

ewext/src/addr_grabber.rs

@@ -1,7 +1,15 @@
-use std::{os::raw::c_void, ptr};
+use std::{mem, os::raw::c_void, ptr, sync::OnceLock};
 
 use iced_x86::{Decoder, DecoderOptions, Mnemonic};
 
+use crate::{
+    lua_bindings::{lua_State, LUA_GLOBALSINDEX},
+    noita::ntypes::{EntityManager, ThiscallFn},
+    LUA,
+};
+
+static GRABBED: OnceLock<Grabbed> = OnceLock::new();
+
 pub(crate) unsafe fn grab_addr_from_instruction(
     func: *const c_void,
     offset: isize,
@@ -25,3 +33,80 @@ pub(crate) unsafe fn grab_addr_from_instruction(
 
     instruction.memory_displacement32() as *mut c_void
 }
+
+struct Grabbed {
+    globals: GrabbedGlobals,
+    fns: GrabbedFns,
+}
+
+// This only stores pointers that are constant, so should be safe to share between threads.
+unsafe impl Sync for Grabbed {}
+unsafe impl Send for Grabbed {}
+
+pub(crate) struct GrabbedGlobals {
+    // These 3 actually point to a pointer.
+    pub(crate) game_global: *mut usize,
+    pub(crate) world_state_entity: *mut usize,
+    pub(crate) entity_manager: *const *mut EntityManager,
+}
+
+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;
+    let world_state_entity =
+        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.lua_getfield(lua, LUA_GLOBALSINDEX, c"GameGetFrameNum".as_ptr());
+    let base = LUA.lua_tocfunction(lua, -1).unwrap() as *const c_void;
+    let load_game_global =
+        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();
+    println!("Game global addr: 0x{:x}", game_global as usize);
+    // Pop the last element.
+    LUA.lua_settop(lua, -2);
+
+    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,
+    ));
+    println!("get_entity addr: 0x{:x}", get_entity as usize);
+    let entity_manager =
+        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);
+
+    GRABBED
+        .set(Grabbed {
+            globals: GrabbedGlobals {
+                game_global,
+                world_state_entity,
+                entity_manager,
+            },
+            fns: GrabbedFns { get_entity },
+        })
+        .ok();
+}
+
+pub(crate) fn grabbed_fns() -> &'static GrabbedFns {
+    &GRABBED.get().expect("to be initialized early").fns
+}
+
+pub(crate) fn grabbed_globals() -> &'static GrabbedGlobals {
+    &GRABBED.get().expect("to be initialized early").globals
+}

ewext/src/lib.rs

@@ -2,18 +2,16 @@ use std::{
     arch::asm,
     cell::{LazyCell, RefCell},
     ffi::{c_int, c_void},
-    mem,
     sync::LazyLock,
 };
 
-use iced_x86::Mnemonic;
+use addr_grabber::{grab_addrs, grabbed_fns, grabbed_globals};
-use lua_bindings::{lua_State, Lua51, LUA_GLOBALSINDEX};
+use lua_bindings::{lua_State, Lua51};
-use noita::{
+use lua_state::LuaState;
-    ntypes::{Entity, EntityManager, ThiscallFn},
+use noita::{ntypes::Entity, NoitaPixelRun, ParticleWorldState};
-    NoitaPixelRun, ParticleWorldState,
-};
 
 mod lua_bindings;
+mod lua_state;
 
 mod noita;
 
@@ -36,23 +34,10 @@ struct SavedWorldState {
     world_state_entity: usize,
 }
 
-struct GrabbedGlobals {
-    // These 3 actually point to a pointer.
-    game_global: *mut usize,
-    world_state_entity: *mut usize,
-    entity_manager: *const *mut EntityManager,
-}
-
-struct GrabbedFns {
-    get_entity: *const ThiscallFn, //unsafe extern "C" fn(*const EntityManager, u32) -> *mut Entity,
-}
-
 #[derive(Default)]
 struct ExtState {
     particle_world_state: Option<ParticleWorldState>,
-    globals: Option<GrabbedGlobals>,
     saved_world_state: Option<SavedWorldState>,
-    fns: Option<GrabbedFns>,
 }
 
 // const EWEXT: [(&'static str, Function); 1] = [("testfn", None)];
@@ -92,11 +77,10 @@ unsafe extern "C" fn encode_area(lua: *mut lua_State) -> c_int {
 }
 
 unsafe fn save_world_state() {
+    let game_global = grabbed_globals().game_global.read();
+    let world_state_entity = grabbed_globals().world_state_entity.read();
     STATE.with(|state| {
         let mut state = state.borrow_mut();
-        let game_global = state.globals.as_ref().unwrap().game_global.read();
-        let world_state_entity = state.globals.as_ref().unwrap().world_state_entity.read();
-
         state.saved_world_state = Some(SavedWorldState {
             game_global,
             world_state_entity,
@@ -109,19 +93,14 @@ unsafe fn load_world_state() {
     STATE.with(|state| {
         let state = state.borrow_mut();
         let saved_ws = state.saved_world_state.as_ref().unwrap();
-        let globals = state.globals.as_ref().unwrap();
+        grabbed_globals().game_global.write(saved_ws.game_global);
-        globals.game_global.write(saved_ws.game_global);
+        grabbed_globals()
-        globals
             .world_state_entity
             .write(saved_ws.world_state_entity);
     });
 }
 
-unsafe extern "C" fn save_world_state_lua(lua: *mut lua_State) -> i32 {
+unsafe extern "C" fn save_world_state_lua(_lua: *mut lua_State) -> i32 {
-    if STATE.with(|state| state.borrow().globals.is_none()) {
-        grab_addrs(lua);
-    }
-
     save_world_state();
     0
 }
@@ -131,81 +110,32 @@ unsafe extern "C" fn load_world_state_lua(_lua: *mut lua_State) -> i32 {
     0
 }
 
-unsafe fn grab_addrs(lua: *mut lua_State) {
+unsafe extern "C" fn make_ephemerial(lua: *mut lua_State) -> c_int {
-    LUA.lua_getfield(lua, LUA_GLOBALSINDEX, c"GameGetWorldStateEntity".as_ptr());
+    let lua_state = LuaState::new(lua);
-    let base = LUA.lua_tocfunction(lua, -1).unwrap() as *const c_void;
+    unsafe {
-    let world_state_entity =
+        let entity_id = lua_state.to_integer(1) as u32;
-        addr_grabber::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.lua_getfield(lua, LUA_GLOBALSINDEX, c"GameGetFrameNum".as_ptr());
+        let entity_manager = grabbed_globals().entity_manager.read();
-    let base = LUA.lua_tocfunction(lua, -1).unwrap() as *const c_void;
+        let mut entity: *mut Entity;
-    let load_game_global =
+        asm!(
-        addr_grabber::grab_addr_from_instruction(base, 0x007bf3c9 - 0x007bf140, Mnemonic::Call); // CALL load_game_global
+            "mov ecx, {entity_manager}",
-    println!("Load game global addr: 0x{:x}", load_game_global as usize);
+            "push {entity_id:e}",
-    let game_global = addr_grabber::grab_addr_from_instruction(
+            "call {get_entity}",
-        load_game_global,
+            entity_manager = in(reg) entity_manager,
-        0x00439c17 - 0x00439bb0,
+            get_entity = in(reg) grabbed_fns().get_entity,
-        Mnemonic::Mov,
+            entity_id = in(reg) entity_id,
-    )
+            clobber_abi("C"),
-    .cast();
+            out("ecx") _,
-    println!("Game global addr: 0x{:x}", game_global as usize);
+            out("eax") entity,
-    // Pop the last element.
+        );
-    LUA.lua_settop(lua, -2);
+        // let entity = (state.fns.as_ref().unwrap().get_entity)(entity_manager, entity_id);
-
+        entity.cast::<c_void>().offset(0x8).cast::<u32>().write(0);
-    LUA.lua_getfield(lua, LUA_GLOBALSINDEX, c"EntityGetFilename".as_ptr());
+    }
-    let base = LUA.lua_tocfunction(lua, -1).unwrap() as *const c_void;
+    0
-    let get_entity = mem::transmute_copy(&addr_grabber::grab_addr_from_instruction(
-        base,
-        0x0079782b - 0x00797570,
-        Mnemonic::Call,
-    ));
-    println!("get_entity addr: 0x{:x}", get_entity as usize);
-    let entity_manager =
-        addr_grabber::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);
-
-    STATE.with(|state| {
-        state.borrow_mut().globals = Some(GrabbedGlobals {
-            game_global,
-            world_state_entity,
-            entity_manager,
-        });
-        state.borrow_mut().fns = Some(GrabbedFns { get_entity })
-    });
 }
 
-unsafe extern "C" fn make_ephemerial(lua: *mut lua_State) -> c_int {
+unsafe extern "C" fn on_world_initialized(lua: *mut lua_State) -> c_int {
-    unsafe {
+    grab_addrs(lua);
-        let entity_id = LUA.lua_tointeger(lua, 1) as u32;
-        STATE.with(|state| {
-            let state = state.borrow();
-            let entity_manager = state.globals.as_ref().unwrap().entity_manager.read();
-            let mut entity: *mut Entity;
-            asm!(
-                "mov ecx, {entity_manager}",
-                "push {entity_id:e}",
-                "call {get_entity}",
-                entity_manager = in(reg) entity_manager,
-                get_entity = in(reg) state.fns.as_ref().unwrap().get_entity,
-                entity_id = in(reg) entity_id,
-                clobber_abi("C"),
-                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);
-        })
-    }
     0
 }
 
@@ -228,6 +158,8 @@ pub unsafe extern "C" fn luaopen_ewext0(lua: *mut lua_State) -> c_int {
         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());
     }
     println!("Initializing ewext - Ok");
     1

ewext/src/lua_state.rs

@@ -0,0 +1,14 @@
+use crate::{lua_bindings::lua_State, LUA};
+
+#[derive(Clone, Copy)]
+pub(crate) struct LuaState(*mut lua_State);
+
+impl LuaState {
+    pub(crate) fn new(lua: *mut lua_State) -> Self {
+        Self(lua)
+    }
+
+    pub(crate) fn to_integer(&self, index: i32) -> isize {
+        unsafe { LUA.lua_tointeger(self.0, index) }
+    }
+}

quant.ew/files/system/ewext_init/ewext_init.lua

@@ -9,6 +9,7 @@ local module = {}
 
 function module.on_world_initialized()
     initial_world_state_entity = GameGetWorldStateEntity()
+    ewext.on_world_initialized()
     ewext.save_world_state()
     local grid_world = world_ffi.get_grid_world()
     local chunk_map = grid_world.vtable.get_chunk_map(grid_world)