From 18497f15289221a514e650982286e38ba5704e02 Mon Sep 17 00:00:00 2001 From: Alex Orlenko Date: Tue, 11 Feb 2025 00:06:03 +0000 Subject: [PATCH] Add library constants support for Luau compiler. Add ability to specify compile-time constants for known library members. --- src/chunk.rs | 118 ++++++++++++++++++++++++++++++++++++++++++++++--- src/lib.rs | 7 ++- src/prelude.rs | 4 +- tests/chunk.rs | 35 ++++++++++++--- 4 files changed, 152 insertions(+), 12 deletions(-) diff --git a/src/chunk.rs b/src/chunk.rs index 41a096ef..02dd35a9 100644 --- a/src/chunk.rs +++ b/src/chunk.rs @@ -130,6 +130,28 @@ pub enum ChunkMode { Binary, } +/// Represents a constant value that can be used by Luau compiler. +#[cfg(any(feature = "luau", doc))] +#[cfg_attr(docsrs, doc(cfg(feature = "luau")))] +#[derive(Clone, Debug)] +pub enum CompileConstant { + Nil, + Boolean(bool), + Number(crate::Number), + Vector(crate::Vector), + String(String), +} + +#[cfg(feature = "luau")] +impl From<&'static str> for CompileConstant { + fn from(s: &'static str) -> Self { + CompileConstant::String(s.to_string()) + } +} + +#[cfg(any(feature = "luau", doc))] +type LibraryMemberConstantMap = std::sync::Arc>; + /// Luau compiler #[cfg(any(feature = "luau", doc))] #[cfg_attr(docsrs, doc(cfg(feature = "luau")))] @@ -144,6 +166,9 @@ pub struct Compiler { vector_type: Option, mutable_globals: Vec, userdata_types: Vec, + libraries_with_known_members: Vec, + library_constants: Option, + disabled_builtins: Vec, } #[cfg(any(feature = "luau", doc))] @@ -168,6 +193,9 @@ impl Compiler { vector_type: None, mutable_globals: Vec::new(), userdata_types: Vec::new(), + libraries_with_known_members: Vec::new(), + library_constants: None, + disabled_builtins: Vec::new(), } } @@ -200,6 +228,7 @@ impl Compiler { /// Possible values: /// * 0 - generate for native modules (default) /// * 1 - generate for all modules + #[must_use] pub const fn set_type_info_level(mut self, level: u8) -> Self { self.type_info_level = level; self @@ -242,15 +271,46 @@ impl Compiler { /// /// It disables the import optimization for fields accessed through these. #[must_use] - pub fn set_mutable_globals(mut self, globals: Vec) -> Self { - self.mutable_globals = globals; + pub fn set_mutable_globals>(mut self, globals: Vec) -> Self { + self.mutable_globals = globals.into_iter().map(|s| s.into()).collect(); self } /// Sets a list of userdata types that will be included in the type information. #[must_use] - pub fn set_userdata_types(mut self, types: Vec) -> Self { - self.userdata_types = types; + pub fn set_userdata_types>(mut self, types: Vec) -> Self { + self.userdata_types = types.into_iter().map(|s| s.into()).collect(); + self + } + + /// Sets constants for known library members. + /// + /// The constants are used by the compiler to optimize the generated bytecode. + /// Optimization level must be at least 2 for this to have any effect. + /// + /// The first element of the tuple is the library name,the second is the member name, and the + /// third is the constant value. + #[must_use] + pub fn set_library_constants(mut self, constants: Vec<(L, M, CompileConstant)>) -> Self + where + L: Into, + M: Into, + { + let map = constants + .into_iter() + .map(|(lib, member, cons)| ((lib.into(), member.into()), cons)) + .collect::>(); + self.library_constants = Some(std::sync::Arc::new(map)); + self.libraries_with_known_members = (self.library_constants.clone()) + .map(|map| map.keys().map(|(lib, _)| lib.clone()).collect()) + .unwrap_or_default(); + self + } + + /// Sets a list of builtins that should be disabled. + #[must_use] + pub fn set_disabled_builtins>(mut self, builtins: Vec) -> Self { + self.disabled_builtins = builtins.into_iter().map(|s| s.into()).collect(); self } @@ -258,7 +318,9 @@ impl Compiler { /// /// Returns [`Error::SyntaxError`] if the source code is invalid. pub fn compile(&self, source: impl AsRef<[u8]>) -> Result> { - use std::os::raw::c_int; + use std::cell::RefCell; + use std::ffi::CStr; + use std::os::raw::{c_char, c_int}; use std::ptr; let vector_lib = self.vector_lib.clone(); @@ -290,6 +352,44 @@ impl Compiler { vec2cstring_ptr!(mutable_globals, mutable_globals_ptr); vec2cstring_ptr!(userdata_types, userdata_types_ptr); + vec2cstring_ptr!(libraries_with_known_members, libraries_with_known_members_ptr); + vec2cstring_ptr!(disabled_builtins, disabled_builtins_ptr); + + thread_local! { + static LIBRARY_MEMBER_CONSTANT_MAP: RefCell = Default::default(); + } + + #[cfg(feature = "luau")] + unsafe extern "C-unwind" fn library_member_constant_callback( + library: *const c_char, + member: *const c_char, + constant: *mut ffi::lua_CompileConstant, + ) { + let library = CStr::from_ptr(library).to_string_lossy(); + let member = CStr::from_ptr(member).to_string_lossy(); + LIBRARY_MEMBER_CONSTANT_MAP.with_borrow(|map| { + if let Some(cons) = map.get(&(library.to_string(), member.to_string())) { + match cons { + CompileConstant::Nil => ffi::luau_set_compile_constant_nil(constant), + CompileConstant::Boolean(b) => { + ffi::luau_set_compile_constant_boolean(constant, *b as c_int) + } + CompileConstant::Number(n) => ffi::luau_set_compile_constant_number(constant, *n), + CompileConstant::Vector(v) => { + #[cfg(not(feature = "luau-vector4"))] + ffi::luau_set_compile_constant_vector(constant, v.x(), v.y(), v.z(), 0.0); + #[cfg(feature = "luau-vector4")] + ffi::luau_set_compile_constant_vector(constant, v.x(), v.y(), v.z(), v.w()); + } + CompileConstant::String(s) => ffi::luau_set_compile_constant_string( + constant, + s.as_ptr() as *const c_char, + s.len(), + ), + } + } + }) + } let bytecode = unsafe { let mut options = ffi::lua_CompileOptions::default(); @@ -302,6 +402,14 @@ impl Compiler { options.vectorType = vector_type.map_or(ptr::null(), |s| s.as_ptr()); options.mutableGlobals = mutable_globals_ptr; options.userdataTypes = userdata_types_ptr; + options.librariesWithKnownMembers = libraries_with_known_members_ptr; + if let Some(map) = self.library_constants.as_ref() { + if !self.libraries_with_known_members.is_empty() { + LIBRARY_MEMBER_CONSTANT_MAP.with_borrow_mut(|gmap| *gmap = map.clone()); + options.libraryMemberConstantCallback = Some(library_member_constant_callback); + } + } + options.disabledBuiltins = disabled_builtins_ptr; ffi::luau_compile(source.as_ref(), options) }; diff --git a/src/lib.rs b/src/lib.rs index a404594c..913bf589 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -126,7 +126,12 @@ pub use crate::hook::HookTriggers; #[cfg(any(feature = "luau", doc))] #[cfg_attr(docsrs, doc(cfg(feature = "luau")))] -pub use crate::{buffer::Buffer, chunk::Compiler, function::CoverageInfo, vector::Vector}; +pub use crate::{ + buffer::Buffer, + chunk::{CompileConstant, Compiler}, + function::CoverageInfo, + vector::Vector, +}; #[cfg(feature = "async")] #[cfg_attr(docsrs, doc(cfg(feature = "async")))] diff --git a/src/prelude.rs b/src/prelude.rs index 68ba8f2f..b84adf49 100644 --- a/src/prelude.rs +++ b/src/prelude.rs @@ -22,7 +22,9 @@ pub use crate::HookTriggers as LuaHookTriggers; #[cfg(feature = "luau")] #[doc(no_inline)] -pub use crate::{CoverageInfo as LuaCoverageInfo, Vector as LuaVector}; +pub use crate::{ + CompileConstant as LuaCompileConstant, CoverageInfo as LuaCoverageInfo, Vector as LuaVector, +}; #[cfg(feature = "async")] #[doc(no_inline)] diff --git a/tests/chunk.rs b/tests/chunk.rs index 16df553b..17becbe2 100644 --- a/tests/chunk.rs +++ b/tests/chunk.rs @@ -96,8 +96,6 @@ fn test_chunk_macro() -> Result<()> { #[cfg(feature = "luau")] #[test] fn test_compiler() -> Result<()> { - use std::vec; - let compiler = mlua::Compiler::new() .set_optimization_level(2) .set_debug_level(2) @@ -106,10 +104,11 @@ fn test_compiler() -> Result<()> { .set_vector_lib("vector") .set_vector_ctor("new") .set_vector_type("vector") - .set_mutable_globals(vec!["mutable_global".into()]) - .set_userdata_types(vec!["MyUserdata".into()]); + .set_mutable_globals(vec!["mutable_global"]) + .set_userdata_types(vec!["MyUserdata"]) + .set_disabled_builtins(vec!["tostring"]); - assert!(compiler.compile("return vector.new(1, 2, 3)").is_ok()); + assert!(compiler.compile("return tostring(vector.new(1, 2, 3))").is_ok()); // Error match compiler.compile("%") { @@ -122,6 +121,32 @@ fn test_compiler() -> Result<()> { Ok(()) } +#[cfg(feature = "luau")] +#[test] +fn test_compiler_library_constants() { + use mlua::{CompileConstant, Compiler, Vector}; + + let compiler = Compiler::new() + .set_optimization_level(2) + .set_library_constants(vec![ + ("mylib", "const_bool", CompileConstant::Boolean(true)), + ("mylib", "const_num", CompileConstant::Number(123.0)), + ("mylib", "const_vec", CompileConstant::Vector(Vector::zero())), + ("mylib", "const_str", "value1".into()), + ]); + + let lua = Lua::new(); + lua.set_compiler(compiler); + let const_bool = lua.load("return mylib.const_bool").eval::().unwrap(); + assert_eq!(const_bool, true); + let const_num = lua.load("return mylib.const_num").eval::().unwrap(); + assert_eq!(const_num, 123.0); + let const_vec = lua.load("return mylib.const_vec").eval::().unwrap(); + assert_eq!(const_vec, Vector::zero()); + let const_str = lua.load("return mylib.const_str").eval::(); + assert_eq!(const_str.unwrap(), "value1"); +} + #[test] fn test_chunk_wrap() -> Result<()> { let lua = Lua::new();