From d3b48cf2f314f780dfa1d2eb91645cc3c78db4ee Mon Sep 17 00:00:00 2001
From: Alex Orlenko <zxteam@protonmail.com>
Date: Mon, 18 Jul 2022 00:59:17 +0100
Subject: [PATCH] Use Luau tags to mark userdata objects as destructed

---
 Cargo.toml          |  2 +-
 src/ffi/luau/lua.rs |  8 +++++++-
 src/lua.rs          | 10 +++++-----
 src/types.rs        |  2 +-
 src/util.rs         | 14 ++++++--------
 5 files changed, 20 insertions(+), 16 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index 3c72242..98bbcb2 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -58,7 +58,7 @@ cc = { version = "1.0" }
 pkg-config = { version = "0.3.17" }
 lua-src = { version = ">= 544.0.0, < 550.0.0", optional = true }
 luajit-src = { version = ">= 210.4.0, < 220.0.0", optional = true }
-luau0-src = { version = "0.3.2", optional = true }
+luau0-src = { version = "0.3.6", optional = true }
 
 [dev-dependencies]
 rustyline = "9.0"
diff --git a/src/ffi/luau/lua.rs b/src/ffi/luau/lua.rs
index 8aa83fc..cbc2c40 100644
--- a/src/ffi/luau/lua.rs
+++ b/src/ffi/luau/lua.rs
@@ -260,6 +260,7 @@ extern "C" {
     pub fn lua_concat(L: *mut lua_State, n: c_int);
     // TODO: lua_encodepointer
     pub fn lua_clock() -> c_double;
+    pub fn lua_setuserdatatag(L: *mut lua_State, idx: c_int, tag: c_int);
     pub fn lua_setuserdatadtor(
         L: *mut lua_State,
         tag: c_int,
@@ -437,7 +438,12 @@ extern "C" {
     pub fn lua_setupvalue(L: *mut lua_State, funcindex: c_int, n: c_int) -> *const c_char;
 
     pub fn lua_singlestep(L: *mut lua_State, enabled: c_int);
-    pub fn lua_breakpoint(L: *mut lua_State, funcindex: c_int, line: c_int, enabled: c_int);
+    pub fn lua_breakpoint(
+        L: *mut lua_State,
+        funcindex: c_int,
+        line: c_int,
+        enabled: c_int,
+    ) -> c_int;
 
     pub fn lua_getcoverage(
         L: *mut lua_State,
diff --git a/src/lua.rs b/src/lua.rs
index 14a91be..60f8bc9 100644
--- a/src/lua.rs
+++ b/src/lua.rs
@@ -25,7 +25,7 @@ use crate::string::String;
 use crate::table::Table;
 use crate::thread::Thread;
 use crate::types::{
-    Callback, CallbackUpvalue, DestructedUserdataMT, Integer, LightUserData, LuaRef, MaybeSend,
+    Callback, CallbackUpvalue, DestructedUserdata, Integer, LightUserData, LuaRef, MaybeSend,
     Number, RegistryKey,
 };
 use crate::userdata::{AnyUserData, UserData, UserDataCell};
@@ -602,13 +602,13 @@ impl Lua {
             "Error while storing extra data",
         );
 
-        // Register `DestructedUserdataMT` type
+        // Register `DestructedUserdata` type
         get_destructed_userdata_metatable(main_state);
         let destructed_mt_ptr = ffi::lua_topointer(main_state, -1);
-        let destructed_mt_typeid = Some(TypeId::of::<DestructedUserdataMT>());
+        let destructed_ud_typeid = TypeId::of::<DestructedUserdata>();
         (*extra.get())
             .registered_userdata_mt
-            .insert(destructed_mt_ptr, destructed_mt_typeid);
+            .insert(destructed_mt_ptr, Some(destructed_ud_typeid));
         ffi::lua_pop(main_state, 1);
 
         mlua_debug_assert!(
@@ -2587,7 +2587,7 @@ impl Lua {
 
         let extra = &*self.extra.get();
         match extra.registered_userdata_mt.get(&mt_ptr) {
-            Some(&type_id) if type_id == Some(TypeId::of::<DestructedUserdataMT>()) => {
+            Some(&type_id) if type_id == Some(TypeId::of::<DestructedUserdata>()) => {
                 Err(Error::UserDataDestructed)
             }
             Some(&type_id) => Ok(type_id),
diff --git a/src/types.rs b/src/types.rs
index 68a7032..7b1f653 100644
--- a/src/types.rs
+++ b/src/types.rs
@@ -83,7 +83,7 @@ pub trait MaybeSend {}
 #[cfg(not(feature = "send"))]
 impl<T> MaybeSend for T {}
 
-pub(crate) struct DestructedUserdataMT;
+pub(crate) struct DestructedUserdata;
 
 /// An auto generated key into the Lua registry.
 ///
diff --git a/src/util.rs b/src/util.rs
index e326c94..66b26e7 100644
--- a/src/util.rs
+++ b/src/util.rs
@@ -310,10 +310,7 @@ pub unsafe fn push_userdata<T>(state: *mut ffi::lua_State, t: T, protect: bool)
 #[inline]
 pub unsafe fn push_userdata<T>(state: *mut ffi::lua_State, t: T, protect: bool) -> Result<()> {
     unsafe extern "C" fn destructor<T>(ud: *mut c_void) {
-        let ud = ud as *mut T;
-        if *(ud.offset(1) as *mut u8) == 0 {
-            ptr::drop_in_place(ud);
-        }
+        ptr::drop_in_place(ud as *mut T);
     }
 
     let size = mem::size_of::<T>() + 1;
@@ -325,7 +322,6 @@ pub unsafe fn push_userdata<T>(state: *mut ffi::lua_State, t: T, protect: bool)
         ffi::lua_newuserdatadtor(state, size, destructor::<T>) as *mut T
     };
     ptr::write(ud, t);
-    *(ud.offset(1) as *mut u8) = 0; // Mark as not destructed
 
     Ok(())
 }
@@ -369,10 +365,12 @@ pub unsafe fn take_userdata<T>(state: *mut ffi::lua_State) -> T {
     get_destructed_userdata_metatable(state);
     ffi::lua_setmetatable(state, -2);
     let ud = get_userdata::<T>(state, -1);
+
+    // Update userdata tag to disable destructor and mark as destructed
+    #[cfg(feature = "luau")]
+    ffi::lua_setuserdatatag(state, -1, 1);
+
     ffi::lua_pop(state, 1);
-    if cfg!(feature = "luau") {
-        *(ud.offset(1) as *mut u8) = 1; // Mark as destructed
-    }
     ptr::read(ud)
 }