Skip to main content

mlua/luau/
mod.rs

1//! Luau-specific extensions and types.
2//!
3//! This module provides Luau-specific functionality including custom [`require`] implementations,
4//! heap memory analysis, and Luau VM integration utilities.
5//!
6//! [`require`]: crate::Lua::create_require_function
7
8use std::ffi::{CStr, CString};
9use std::os::raw::c_int;
10use std::ptr;
11
12use crate::chunk::ChunkMode;
13use crate::error::{Error, Result};
14use crate::function::Function;
15use crate::state::{ExtraData, Lua, callback_error_ext};
16use crate::traits::{FromLuaMulti, IntoLua};
17use crate::types::MaybeSend;
18
19pub use heap_dump::HeapDump;
20pub use require::{FsRequirer, NavigateError, Require};
21
22// Since Luau has some missing standard functions, we re-implement them here
23
24impl Lua {
25    /// Create a custom Luau `require` function using provided [`Require`] implementation to find
26    /// and load modules.
27    #[cfg(any(feature = "luau", doc))]
28    #[cfg_attr(docsrs, doc(cfg(feature = "luau")))]
29    pub fn create_require_function<R: Require + MaybeSend + 'static>(&self, require: R) -> Result<Function> {
30        require::create_require_function(self, require)
31    }
32
33    /// Set the memory category for subsequent allocations from this Lua state.
34    ///
35    /// The category "main" is reserved for the default memory category.
36    /// Maximum of 255 categories can be registered.
37    /// The category is set per Lua thread (state) and affects all allocations made from that
38    /// thread.
39    ///
40    /// Return error if too many categories are registered or if the category name is invalid.
41    ///
42    /// See [`Lua::heap_dump`] for tracking memory usage by category.
43    #[cfg(any(feature = "luau", doc))]
44    #[cfg_attr(docsrs, doc(cfg(feature = "luau")))]
45    pub fn set_memory_category(&self, category: &str) -> Result<()> {
46        let lua = self.lock();
47
48        if category.contains(|c| !matches!(c, 'a'..='z' | 'A'..='Z' | '0'..='9' | '-' | '_')) {
49            return Err(Error::runtime("invalid memory category name"));
50        }
51        let cat_id = unsafe {
52            let extra = ExtraData::get(lua.state());
53            match ((*extra).mem_categories.iter().enumerate())
54                .find(|&(_, name)| name.as_bytes() == category.as_bytes())
55            {
56                Some((id, _)) => id as u8,
57                None => {
58                    let new_id = (*extra).mem_categories.len() as u8;
59                    if new_id == 255 {
60                        return Err(Error::runtime("too many memory categories registered"));
61                    }
62                    (*extra).mem_categories.push(CString::new(category).unwrap());
63                    new_id
64                }
65            }
66        };
67        unsafe { ffi::lua_setmemcat(lua.state(), cat_id as i32) };
68
69        Ok(())
70    }
71
72    /// Dumps the current Lua VM heap state.
73    ///
74    /// The returned `HeapDump` can be used to analyze memory usage.
75    /// It's recommended to call [`Lua::gc_collect`] before dumping the heap.
76    #[cfg(any(feature = "luau", doc))]
77    #[cfg_attr(docsrs, doc(cfg(feature = "luau")))]
78    pub fn heap_dump(&self) -> Result<HeapDump> {
79        let lua = self.lock();
80        unsafe { heap_dump::HeapDump::new(lua.state()).ok_or_else(|| Error::runtime("unable to dump heap")) }
81    }
82
83    pub(crate) unsafe fn configure_luau(&self) -> Result<()> {
84        let globals = self.globals();
85
86        globals.raw_set("collectgarbage", self.create_c_function(lua_collectgarbage)?)?;
87        globals.raw_set("loadstring", self.create_c_function(lua_loadstring)?)?;
88
89        // Set `_VERSION` global to include version number
90        // The environment variable `LUAU_VERSION` set by the build script
91        if let Some(version) = ffi::luau_version() {
92            globals.raw_set("_VERSION", format!("Luau {version}"))?;
93        }
94
95        // Enable default `require` implementation
96        let require = self.create_require_function(FsRequirer::new())?;
97        self.globals().raw_set("require", require)?;
98
99        Ok(())
100    }
101}
102
103unsafe extern "C-unwind" fn lua_collectgarbage(state: *mut ffi::lua_State) -> c_int {
104    let option = ffi::luaL_optstring(state, 1, cstr!("collect"));
105    let option = CStr::from_ptr(option);
106    let arg = ffi::luaL_optinteger(state, 2, 0);
107    let is_sandboxed = (*ExtraData::get(state)).sandboxed;
108    match option.to_str() {
109        Ok("collect") if !is_sandboxed => {
110            ffi::lua_gc(state, ffi::LUA_GCCOLLECT, 0);
111            0
112        }
113        Ok("stop") if !is_sandboxed => {
114            ffi::lua_gc(state, ffi::LUA_GCSTOP, 0);
115            0
116        }
117        Ok("restart") if !is_sandboxed => {
118            ffi::lua_gc(state, ffi::LUA_GCRESTART, 0);
119            0
120        }
121        Ok("count") => {
122            let kbytes = ffi::lua_gc(state, ffi::LUA_GCCOUNT, 0) as ffi::lua_Number;
123            let kbytes_rem = ffi::lua_gc(state, ffi::LUA_GCCOUNTB, 0) as ffi::lua_Number;
124            ffi::lua_pushnumber(state, kbytes + kbytes_rem / 1024.0);
125            1
126        }
127        Ok("step") if !is_sandboxed => {
128            let res = ffi::lua_gc(state, ffi::LUA_GCSTEP, arg as _);
129            ffi::lua_pushboolean(state, res);
130            1
131        }
132        Ok("isrunning") if !is_sandboxed => {
133            let res = ffi::lua_gc(state, ffi::LUA_GCISRUNNING, 0);
134            ffi::lua_pushboolean(state, res);
135            1
136        }
137        _ => ffi::luaL_error(state, cstr!("collectgarbage called with invalid option")),
138    }
139}
140
141unsafe extern "C-unwind" fn lua_loadstring(state: *mut ffi::lua_State) -> c_int {
142    callback_error_ext(state, ptr::null_mut(), false, move |extra, nargs| {
143        let rawlua = (*extra).raw_lua();
144        let (chunk, chunk_name) =
145            <(String, Option<String>)>::from_stack_args(nargs, 1, Some("loadstring"), rawlua)?;
146        let chunk_name = chunk_name.as_deref().unwrap_or("=(loadstring)");
147        (rawlua.lua())
148            .load(chunk)
149            .set_name(chunk_name)
150            .set_mode(ChunkMode::Text)
151            .into_function()?
152            .push_into_stack(rawlua)?;
153        Ok(1)
154    })
155}
156
157mod heap_dump;
158mod json;
159mod require;