Add environment option to luau load built-in

CHANGELOG.md

@@ -52,6 +52,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   print(now:toUniversalTime())
   ```
 
+- Added support for setting a custom environment in load options for `luau.load`, not subject to `getfenv` / `setfenv` deoptimizations
 - Added [Terrain:GetMaterialColor](https://create.roblox.com/docs/reference/engine/classes/Terrain#GetMaterialColor) and [Terrain:SetMaterialColor](https://create.roblox.com/docs/reference/engine/classes/Terrain#SetMaterialColor) ([#93])
 - Added support for a variable number of arguments for CFrame methods ([#85])
 

src/lune/builtins/luau/mod.rs

@@ -33,7 +33,27 @@ fn load_source<'lua>(
     lua: &'lua Lua,
     (source, options): (LuaString<'lua>, LuauLoadOptions),
 ) -> LuaResult<LuaFunction<'lua>> {
-    lua.load(source.as_bytes())
-        .set_name(options.debug_name)
-        .into_function()
+    let mut chunk = lua.load(source.as_bytes()).set_name(options.debug_name);
+
+    if let Some(environment) = options.environment {
+        let environment_with_globals = lua.create_table()?;
+
+        if let Some(meta) = environment.get_metatable() {
+            environment_with_globals.set_metatable(Some(meta));
+        }
+
+        for pair in lua.globals().pairs() {
+            let (key, value): (LuaValue, LuaValue) = pair?;
+            environment_with_globals.set(key, value)?;
+        }
+
+        for pair in environment.pairs() {
+            let (key, value): (LuaValue, LuaValue) = pair?;
+            environment_with_globals.set(key, value)?;
+        }
+
+        chunk = chunk.set_environment(environment_with_globals);
+    }
+
+    chunk.into_function()
 }

src/lune/builtins/luau/options.rs

@@ -73,19 +73,21 @@ impl<'lua> FromLua<'lua> for LuauCompileOptions {
     }
 }
 
-pub struct LuauLoadOptions {
+pub struct LuauLoadOptions<'lua> {
     pub(crate) debug_name: String,
+    pub(crate) environment: Option<LuaTable<'lua>>,
 }
 
-impl Default for LuauLoadOptions {
+impl Default for LuauLoadOptions<'_> {
     fn default() -> Self {
         Self {
             debug_name: DEFAULT_DEBUG_NAME.to_string(),
+            environment: None,
         }
     }
 }
 
-impl<'lua> FromLua<'lua> for LuauLoadOptions {
+impl<'lua> FromLua<'lua> for LuauLoadOptions<'lua> {
     fn from_lua(value: LuaValue<'lua>, _: &'lua Lua) -> LuaResult<Self> {
         Ok(match value {
             LuaValue::Nil => Self::default(),
@@ -96,10 +98,15 @@ impl<'lua> FromLua<'lua> for LuauLoadOptions {
                     options.debug_name = debug_name;
                 }
 
+                if let Some(environment) = t.get("environment")? {
+                    options.environment = Some(environment);
+                }
+
                 options
             }
             LuaValue::String(s) => Self {
                 debug_name: s.to_string_lossy().to_string(),
+                environment: None,
             },
             _ => {
                 return Err(LuaError::FromLuaConversionError {

tests/luau/load.luau

@@ -31,3 +31,51 @@ local success = pcall(function()
 end)
 
 assert(success, "expected `luau.load` to be able to process the result of `luau.compile`")
+
+local CUSTOM_SOURCE_WITH_FOO_FN = "return foo()"
+
+-- NOTE: We use newproxy here to make a userdata to ensure
+-- we get the *exact* same value sent back, not some copy
+local fooValue = newproxy(false)
+local fooFn = luau.load(CUSTOM_SOURCE_WITH_FOO_FN, {
+	environment = {
+		foo = function()
+			return fooValue
+		end,
+	},
+})
+
+local fooFnRet = fooFn()
+assert(fooFnRet == fooValue, "expected `luau.load` with custom environment to return proper values")
+
+local CUSTOM_SOURCE_WITH_PRINT_FN = "return print()"
+
+-- NOTE: Same as what we did above, new userdata to guarantee unique-ness
+local overriddenValue = newproxy(false)
+local overriddenFn = luau.load(CUSTOM_SOURCE_WITH_PRINT_FN, {
+	environment = {
+		print = function()
+			return overriddenValue
+		end,
+	},
+})
+
+local overriddenFnRet = overriddenFn()
+assert(
+	overriddenFnRet == overriddenValue,
+	"expected `luau.load` with overridden environment to return proper values"
+)
+
+local CUSTOM_SOURCE_WITH_DEFAULT_FN = "return string.lower(...)"
+
+local overriddenFn2 = luau.load(CUSTOM_SOURCE_WITH_DEFAULT_FN, {
+	environment = {
+		hello = "world",
+	},
+})
+
+local overriddenFn2Ret = overriddenFn2("LOWERCASE")
+assert(
+	overriddenFn2Ret == "lowercase",
+	"expected `luau.load` with overridden environment to contain default globals"
+)

types/luau.luau

@@ -27,10 +27,12 @@ export type CompileOptions = {
 
 	This is a dictionary that may contain one or more of the following values:
 
-	* `debugName` - The debug name of the closure. Defaults to `string ["..."]`
+	* `debugName` - The debug name of the closure. Defaults to `luau.load(...)`
+	* `environment` - Environment values to set and/or override. Includes default globals unless overwritten.
 ]=]
 export type LoadOptions = {
-	debugName: string,
+	debugName: string?,
+	environment: { [string]: any }?,
 }
 
 --[=[