mlua_codemp_patch/state/
extra.rs

1use std::any::TypeId;
2use std::cell::UnsafeCell;
3use std::mem::MaybeUninit;
4use std::os::raw::{c_int, c_void};
5use std::ptr;
6use std::rc::Rc;
7use std::sync::Arc;
8
9use parking_lot::Mutex;
10use rustc_hash::FxHashMap;
11
12use crate::error::Result;
13use crate::state::RawLua;
14use crate::stdlib::StdLib;
15use crate::types::{AppData, ReentrantMutex, XRc};
16use crate::util::{get_internal_metatable, push_internal_userdata, TypeKey, WrappedFailure};
17
18#[cfg(any(feature = "luau", doc))]
19use crate::chunk::Compiler;
20
21#[cfg(feature = "async")]
22use {futures_util::task::noop_waker_ref, std::ptr::NonNull, std::task::Waker};
23
24use super::{Lua, WeakLua};
25
26// Unique key to store `ExtraData` in the registry
27static EXTRA_REGISTRY_KEY: u8 = 0;
28
29const WRAPPED_FAILURE_POOL_SIZE: usize = 64;
30const REF_STACK_RESERVE: c_int = 1;
31
32/// Data associated with the Lua state.
33pub(crate) struct ExtraData {
34    pub(super) lua: MaybeUninit<Lua>,
35    pub(super) weak: MaybeUninit<WeakLua>,
36    pub(super) owned: bool,
37
38    pub(super) registered_userdata: FxHashMap<TypeId, c_int>,
39    pub(super) registered_userdata_mt: FxHashMap<*const c_void, Option<TypeId>>,
40    pub(super) last_checked_userdata_mt: (*const c_void, Option<TypeId>),
41
42    // When Lua instance dropped, setting `None` would prevent collecting `RegistryKey`s
43    pub(super) registry_unref_list: Arc<Mutex<Option<Vec<c_int>>>>,
44
45    // Container to store arbitrary data (extensions)
46    pub(super) app_data: AppData,
47
48    pub(super) safe: bool,
49    pub(super) libs: StdLib,
50    // Used in module mode
51    pub(super) skip_memory_check: bool,
52
53    // Auxiliary thread to store references
54    pub(super) ref_thread: *mut ffi::lua_State,
55    pub(super) ref_stack_size: c_int,
56    pub(super) ref_stack_top: c_int,
57    pub(super) ref_free: Vec<c_int>,
58
59    // Pool of `WrappedFailure` enums in the ref thread (as userdata)
60    pub(super) wrapped_failure_pool: Vec<c_int>,
61    // Pool of `Thread`s (coroutines) for async execution
62    #[cfg(feature = "async")]
63    pub(super) thread_pool: Vec<c_int>,
64
65    // Address of `WrappedFailure` metatable
66    pub(super) wrapped_failure_mt_ptr: *const c_void,
67
68    // Waker for polling futures
69    #[cfg(feature = "async")]
70    pub(super) waker: NonNull<Waker>,
71
72    #[cfg(not(feature = "luau"))]
73    pub(super) hook_callback: Option<crate::types::HookCallback>,
74    #[cfg(not(feature = "luau"))]
75    pub(super) hook_thread: *mut ffi::lua_State,
76    #[cfg(feature = "lua54")]
77    pub(super) warn_callback: Option<crate::types::WarnCallback>,
78    #[cfg(feature = "luau")]
79    pub(super) interrupt_callback: Option<crate::types::InterruptCallback>,
80
81    #[cfg(feature = "luau")]
82    pub(super) sandboxed: bool,
83    #[cfg(feature = "luau")]
84    pub(super) compiler: Option<Compiler>,
85    #[cfg(feature = "luau-jit")]
86    pub(super) enable_jit: bool,
87}
88
89impl Drop for ExtraData {
90    fn drop(&mut self) {
91        unsafe {
92            if !self.owned {
93                self.lua.assume_init_drop();
94            }
95
96            self.weak.assume_init_drop();
97        }
98        *self.registry_unref_list.lock() = None;
99    }
100}
101
102static EXTRA_TYPE_KEY: u8 = 0;
103
104impl TypeKey for XRc<UnsafeCell<ExtraData>> {
105    #[inline(always)]
106    fn type_key() -> *const c_void {
107        &EXTRA_TYPE_KEY as *const u8 as *const c_void
108    }
109}
110
111impl ExtraData {
112    // Index of `error_traceback` function in auxiliary thread stack
113    #[cfg(any(feature = "lua51", feature = "luajit", feature = "luau"))]
114    pub(super) const ERROR_TRACEBACK_IDX: c_int = 1;
115
116    pub(super) unsafe fn init(state: *mut ffi::lua_State, owned: bool) -> XRc<UnsafeCell<Self>> {
117        // Create ref stack thread and place it in the registry to prevent it
118        // from being garbage collected.
119        let ref_thread = mlua_expect!(
120            protect_lua!(state, 0, 0, |state| {
121                let thread = ffi::lua_newthread(state);
122                ffi::luaL_ref(state, ffi::LUA_REGISTRYINDEX);
123                thread
124            }),
125            "Error while creating ref thread",
126        );
127
128        let wrapped_failure_mt_ptr = {
129            get_internal_metatable::<WrappedFailure>(state);
130            let ptr = ffi::lua_topointer(state, -1);
131            ffi::lua_pop(state, 1);
132            ptr
133        };
134
135        // Store `error_traceback` function on the ref stack
136        #[cfg(any(feature = "lua51", feature = "luajit", feature = "luau"))]
137        {
138            ffi::lua_pushcfunction(ref_thread, crate::util::error_traceback);
139            assert_eq!(ffi::lua_gettop(ref_thread), Self::ERROR_TRACEBACK_IDX);
140        }
141
142        #[allow(clippy::arc_with_non_send_sync)]
143        let extra = XRc::new(UnsafeCell::new(ExtraData {
144            lua: MaybeUninit::uninit(),
145            weak: MaybeUninit::uninit(),
146            owned,
147            registered_userdata: FxHashMap::default(),
148            registered_userdata_mt: FxHashMap::default(),
149            last_checked_userdata_mt: (ptr::null(), None),
150            registry_unref_list: Arc::new(Mutex::new(Some(Vec::new()))),
151            app_data: AppData::default(),
152            safe: false,
153            libs: StdLib::NONE,
154            skip_memory_check: false,
155            ref_thread,
156            // We need some reserved stack space to move values in and out of the ref stack.
157            ref_stack_size: ffi::LUA_MINSTACK - REF_STACK_RESERVE,
158            ref_stack_top: ffi::lua_gettop(ref_thread),
159            ref_free: Vec::new(),
160            wrapped_failure_pool: Vec::with_capacity(WRAPPED_FAILURE_POOL_SIZE),
161            #[cfg(feature = "async")]
162            thread_pool: Vec::new(),
163            wrapped_failure_mt_ptr,
164            #[cfg(feature = "async")]
165            waker: NonNull::from(noop_waker_ref()),
166            #[cfg(not(feature = "luau"))]
167            hook_callback: None,
168            #[cfg(not(feature = "luau"))]
169            hook_thread: ptr::null_mut(),
170            #[cfg(feature = "lua54")]
171            warn_callback: None,
172            #[cfg(feature = "luau")]
173            interrupt_callback: None,
174            #[cfg(feature = "luau")]
175            sandboxed: false,
176            #[cfg(feature = "luau")]
177            compiler: None,
178            #[cfg(feature = "luau-jit")]
179            enable_jit: true,
180        }));
181
182        // Store it in the registry
183        mlua_expect!(Self::store(&extra, state), "Error while storing extra data");
184
185        extra
186    }
187
188    pub(super) unsafe fn set_lua(&mut self, raw: &XRc<ReentrantMutex<RawLua>>) {
189        self.lua.write(Lua {
190            raw: XRc::clone(raw),
191            collect_garbage: false,
192        });
193        if self.owned {
194            XRc::decrement_strong_count(XRc::as_ptr(raw));
195        }
196        self.weak.write(WeakLua(XRc::downgrade(raw)));
197    }
198
199    pub(super) unsafe fn get(state: *mut ffi::lua_State) -> *mut Self {
200        #[cfg(feature = "luau")]
201        if cfg!(not(feature = "module")) {
202            // In the main app we can use `lua_callbacks` to access ExtraData
203            return (*ffi::lua_callbacks(state)).userdata as *mut _;
204        }
205
206        let extra_key = &EXTRA_REGISTRY_KEY as *const u8 as *const c_void;
207        if ffi::lua_rawgetp(state, ffi::LUA_REGISTRYINDEX, extra_key) != ffi::LUA_TUSERDATA {
208            // `ExtraData` can be null only when Lua state is foreign.
209            // This case in used in `Lua::try_from_ptr()`.
210            ffi::lua_pop(state, 1);
211            return ptr::null_mut();
212        }
213        let extra_ptr = ffi::lua_touserdata(state, -1) as *mut Rc<UnsafeCell<ExtraData>>;
214        ffi::lua_pop(state, 1);
215        (*extra_ptr).get()
216    }
217
218    unsafe fn store(extra: &XRc<UnsafeCell<Self>>, state: *mut ffi::lua_State) -> Result<()> {
219        #[cfg(feature = "luau")]
220        if cfg!(not(feature = "module")) {
221            (*ffi::lua_callbacks(state)).userdata = extra.get() as *mut _;
222            return Ok(());
223        }
224
225        push_internal_userdata(state, XRc::clone(extra), true)?;
226        protect_lua!(state, 1, 0, fn(state) {
227            let extra_key = &EXTRA_REGISTRY_KEY as *const u8 as *const c_void;
228            ffi::lua_rawsetp(state, ffi::LUA_REGISTRYINDEX, extra_key);
229        })
230    }
231
232    #[inline(always)]
233    pub(super) unsafe fn lua(&self) -> &Lua {
234        self.lua.assume_init_ref()
235    }
236
237    #[inline(always)]
238    pub(super) unsafe fn raw_lua(&self) -> &RawLua {
239        &*self.lua.assume_init_ref().raw.data_ptr()
240    }
241
242    #[inline(always)]
243    pub(super) unsafe fn weak(&self) -> &WeakLua {
244        self.weak.assume_init_ref()
245    }
246}
mlua_codemp_patch/state/extra.rs

mlua_codemp_patch/state/
extra.rs
