feat: add get/setFnEnvironment functions

natecraddock · natecraddock · commit 458aced082ed · 2023-02-15T21:49:24.000-07:00
diff --git a/src/ziglua-5.1/lib.zig b/src/ziglua-5.1/lib.zig
@@ -461,6 +461,12 @@ pub const Lua = struct {
         return c.lua_getallocf(lua.state, @ptrCast([*c]?*anyopaque, data)).?;
     }
 
+    /// Pushes onto the stack the environment table of the value at the given index.
+    /// See https://www.lua.org/manual/5.1/manual.html#lua_getfenv
+    pub fn getFnEnvironment(lua: *Lua, index: i32) void {
+        c.lua_getfenv(lua.state, index);
+    }
+
     /// Pushes onto the stack the value t[key] where t is the value at the given index
     /// See https://www.lua.org/manual/5.4/manual.html#lua_getfield
     pub fn getField(lua: *Lua, index: i32, key: [:0]const u8) void {
@@ -807,6 +813,13 @@ pub const Lua = struct {
         c.lua_setallocf(lua.state, alloc_fn, data);
     }
 
+    /// Pops a table from the stack and sets it as the new environment for the value at the
+    /// given index. Returns an error if the value at that index is not a function or thread or userdata.
+    /// See https://www.lua.org/manual/5.1/manual.html#setfenv
+    pub fn setFnEnvironment(lua: *Lua, index: i32) !void {
+        if (c.lua_setfenv(lua.state, index) == 0) return error.Fail;
+    }
+
     /// Does the equivalent to t[`k`] = v where t is the value at the given `index`
     /// and v is the value on the top of the stack
     pub fn setField(lua: *Lua, index: i32, k: [:0]const u8) void {
diff --git a/src/ziglua-5.1/tests.zig b/src/ziglua-5.1/tests.zig
@@ -1116,3 +1116,35 @@ test "refs" {
     _ = Lua.doFile;
     _ = Lua.loadFile;
 }
+
+test "function environments" {
+    var lua = try Lua.init(testing.allocator);
+    defer lua.deinit();
+
+    try lua.doString("function test() return x end");
+
+    // set the global _G.x to be 10
+    lua.pushInteger(10);
+    lua.setGlobal("x");
+
+    lua.getGlobal("test");
+    try lua.protectedCall(0, 1, 0);
+    try testing.expectEqual(@as(Integer, 10), lua.toInteger(1));
+    lua.pop(1);
+
+    // now set the functions table to have a different value of x
+    lua.getGlobal("test");
+    lua.newTable();
+    lua.pushInteger(20);
+    lua.setField(2, "x");
+    try lua.setFnEnvironment(1);
+
+    try lua.protectedCall(0, 1, 0);
+    try testing.expectEqual(@as(Integer, 20), lua.toInteger(1));
+    lua.pop(1);
+
+    lua.getGlobal("test");
+    lua.getFnEnvironment(1);
+    lua.getField(2, "x");
+    try testing.expectEqual(@as(Integer, 20), lua.toInteger(3));
+}
