From c90cac5189943563de7bfe1d55e85b1da2ef0dad Mon Sep 17 00:00:00 2001
From: Alex Orlenko <zxteam@protonmail.com>
Date: Sun, 6 Jul 2025 10:35:07 +0100
Subject: [PATCH] Add `Lua::set_globals` method to replace global environment.
 Closes #611

---
 src/state.rs   | 33 +++++++++++++++++++++++++++++++++
 tests/tests.rs | 25 +++++++++++++++++++++++++
 2 files changed, 58 insertions(+)

diff --git a/src/state.rs b/src/state.rs
index fd34ee9..1070112 100644
--- a/src/state.rs
+++ b/src/state.rs
@@ -1550,6 +1550,39 @@ impl Lua {
         }
     }
 
+    /// Sets the global environment.
+    ///
+    /// This will replace the current global environment with the provided `globals` table.
+    ///
+    /// For Lua 5.2+ the globals table is stored in the registry and shared between all threads.
+    /// For Lua 5.1 and Luau the globals table is stored in each thread.
+    ///
+    /// Please note that any existing Lua functions have cached global environment and will not
+    /// see the changes made by this method.
+    /// To update the environment for existing Lua functions, use [`Function::set_environment`].
+    pub fn set_globals(&self, globals: Table) -> Result<()> {
+        let lua = self.lock();
+        let state = lua.state();
+        unsafe {
+            #[cfg(feature = "luau")]
+            if (*lua.extra.get()).sandboxed {
+                return Err(Error::runtime("cannot change globals in a sandboxed Lua state"));
+            }
+
+            let _sg = StackGuard::new(state);
+            check_stack(state, 1)?;
+
+            lua.push_ref(&globals.0);
+
+            #[cfg(any(feature = "lua54", feature = "lua53", feature = "lua52"))]
+            ffi::lua_rawseti(state, ffi::LUA_REGISTRYINDEX, ffi::LUA_RIDX_GLOBALS);
+            #[cfg(any(feature = "lua51", feature = "luajit", feature = "luau"))]
+            ffi::lua_replace(state, ffi::LUA_GLOBALSINDEX);
+        }
+
+        Ok(())
+    }
+
     /// Returns a handle to the active `Thread`.
     ///
     /// For calls to `Lua` this will be the main Lua thread, for parameters given to a callback,
diff --git a/tests/tests.rs b/tests/tests.rs
index f0cad94..25dfb49 100644
--- a/tests/tests.rs
+++ b/tests/tests.rs
@@ -147,6 +147,31 @@ fn test_eval() -> Result<()> {
     Ok(())
 }
 
+#[test]
+fn test_replace_globals() -> Result<()> {
+    let lua = Lua::new();
+
+    let globals = lua.create_table()?;
+    globals.set("foo", "bar")?;
+
+    lua.set_globals(globals.clone())?;
+    let val = lua.load("return foo").eval::<StdString>()?;
+    assert_eq!(val, "bar");
+
+    // Updating globals in sandboxed Lua state is not allowed
+    #[cfg(feature = "luau")]
+    {
+        lua.sandbox(true)?;
+        match lua.set_globals(globals) {
+            Err(Error::RuntimeError(msg))
+                if msg.contains("cannot change globals in a sandboxed Lua state") => {}
+            r => panic!("expected RuntimeError(...) with a specific error message, got {r:?}"),
+        }
+    }
+
+    Ok(())
+}
+
 #[test]
 fn test_load_mode() -> Result<()> {
     let lua = unsafe { Lua::unsafe_new() };