Add library constants support for Luau compiler.

Add ability to specify compile-time constants for known library members.

Author: khvzak

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<HashMap<(String, String), CompileConstant>>;
+
 /// Luau compiler
 #[cfg(any(feature = "luau", doc))]
 #[cfg_attr(docsrs, doc(cfg(feature = "luau")))]
@@ -144,6 +166,9 @@ pub struct Compiler {
     vector_type: Option<String>,
     mutable_globals: Vec<String>,
     userdata_types: Vec<String>,
+    libraries_with_known_members: Vec<String>,
+    library_constants: Option<LibraryMemberConstantMap>,
+    disabled_builtins: Vec<String>,
 }
 
 #[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,23 +271,56 @@ impl Compiler {
     ///
     /// It disables the import optimization for fields accessed through these.
     #[must_use]
-    pub fn set_mutable_globals(mut self, globals: Vec<String>) -> Self {
-        self.mutable_globals = globals;
+    pub fn set_mutable_globals<S: Into<String>>(mut self, globals: Vec<S>) -> 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<String>) -> Self {
-        self.userdata_types = types;
+    pub fn set_userdata_types<S: Into<String>>(mut self, types: Vec<S>) -> 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<L, M>(mut self, constants: Vec<(L, M, CompileConstant)>) -> Self
+    where
+        L: Into<String>,
+        M: Into<String>,
+    {
+        let map = constants
+            .into_iter()
+            .map(|(lib, member, cons)| ((lib.into(), member.into()), cons))
+            .collect::<HashMap<_, _>>();
+        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<S: Into<String>>(mut self, builtins: Vec<S>) -> Self {
+        self.disabled_builtins = builtins.into_iter().map(|s| s.into()).collect();
         self
     }
 
     /// Compiles the `source` into bytecode.
     ///
     /// Returns [`Error::SyntaxError`] if the source code is invalid.
     pub fn compile(&self, source: impl AsRef<[u8]>) -> Result<Vec<u8>> {
-        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<LibraryMemberConstantMap> = 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)
         };
 

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")))]

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)]

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::<bool>().unwrap();
+    assert_eq!(const_bool, true);
+    let const_num = lua.load("return mylib.const_num").eval::<f64>().unwrap();
+    assert_eq!(const_num, 123.0);
+    let const_vec = lua.load("return mylib.const_vec").eval::<Vector>().unwrap();
+    assert_eq!(const_vec, Vector::zero());
+    let const_str = lua.load("return mylib.const_str").eval::<String>();
+    assert_eq!(const_str.unwrap(), "value1");
+}
+
 #[test]
 fn test_chunk_wrap() -> Result<()> {
     let lua = Lua::new();