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