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