Implementing setfenv in Lua 5.2, 5.3, and above

If you've programmed in Lua 5.1 you've probably come across the function setfenv (and the associated getfenv). Lua 5.2 removed those functions favor of the _ENV variable. The _ENV variable is an interesting change, and can accomplish many of the same things, but it is not a direct replacement.

View the completed implementation or read on to learn about how it’s created.

The _ENV variable

The _ENV variable provides a lexical way of setting a function’s environment. In the sense that functions have access to lexically scoped variables created a closure, a function’s environment is set by the closest lexically set value of _ENV.

Free variables referenced by closures are called upvalues. The _ENV variable is just an upvalue of a function that’s has the special purpose of being the place where global variable references are searched.

In order to re-implement setfenv, and set the function’s environment, that upvalue must be updated to point to a table to act as the environment.

Getting upvalues

Upvalues can be accessed and manipulated in Lua by using the debug functions debug.getupvalue, debug.setupvalue, and debug.upvaluejoin.

Lua upvalues are indexed by an integer, much like an array. The debug function debug.getupvalue(func, index) is used to access them.

Before implementing setfenv we'll examine a function’s upvalues. Here’s a basic function:

local number = 15
local function lucky()
  print("your lucky number: " .. number)
end

There is one local variable closed on the function, the variable number.

Here’s a chunk of code that will print out the upvalues. It works by enumerating positive integers until nil is returned by debug.getupvalue.

-- print out the upvalues of the function lucky
local idx = 1
while true do
  local name, val = debug.getupvalue(lucky, idx)
  if not name then break end
  print(name, val)
  idx = idx + 1
end

This will produce something like:

_ENV    table: 0x11e39f0
number  15

There are two upvalues, _ENV and number. What’s the default value of _ENV? It’s equal to the global variable _G.

As a quick example, how many upvalues does this function have? (note the absence of local on number)

number = 15

local function lucky2()
  print("your lucky number: " .. number)
end

Here’s the output:

_ENV    table: 0x11e39f0

number is no longer an upvalue because it is now a global variable. When number is referenced inside of that function it will look within the _ENV upvalue, effectively the value of _ENV.number.

One last example, what’s the upvalue of this (empty) function:

local function no_so_lucky()
end

You might be thinking “Well, just _ENV right?”. In reality this function has no upvalues. A function only has an upvalue on _ENV if it references a global variable. In the first example, the function referenced print globally, and the second example referenced both print and number. This is worth keeping in mind if you're ever trying to debug upvalue code.

Setting upvalues

The function debug.setupvalue(func, index, value) sets an upvalue for a function. As mentioned above, upvalues are referenced by their index. In order to set an upvalue by name, the index first must be found.

Here is the snippet above generalized to return a single upvalue’s index by its name:

local function get_upvalue(fn, search_name)
  local idx = 1
  while true do
    local name, val = debug.getupvalue(fn, idx)
    if not name then break end
    if name == search_name then
      return idx, val
    end
    idx = idx + 1
  end
end

Now it’s trivial to change an upvalue of the example function:

local number = 15
local function lucky()
  print("your lucky number: " .. number)
end

debug.setupvalue(get_upvalue(lucky, "number"), 22)

lucky() --> your lucky number: 22

Taking it a step further you might the upvalue _ENV to replace print with something else.

debug.setupvalue(lucky, get_upvalue(lucky, "_ENV"), {
  print = function()
    print("no lucky number for you!")
  end
})

lucky()
-- uhhhh....

This produces a stack overflow error, there’s an infinite loop within print. It kinda makes sense, we did just re-define print to call print with our new environment.

The problem is that upvalues can be shared across many functions. When we set the upvalue on lucky we actually changed every function environment. We changed the global environment to point to a function that just has print in it.

If print is called on the top level, instead of calling lucky, and the same infinite loop will happen. This is because the global scope also shared the same _ENV.

Additionally, calling any other regularly available global function, like type, will fail since the new environment only provides print.

Here’s another example:

local number = 15
local function lucky()
  print("your lucky number: " .. number)
end

local function lucky2()
  print("your other number: " .. number)
end

debug.setupvalue(lucky, get_upvalue(lucky, "number"), 22)
lucky2() --> your other number: 22

There are two functions that refer to the same upvalue number. Calling debug.setupvalue in this case has the same effect as writing number = 22, which changes what the results of both of the functions.

Since debug.setupvalue doesn’t work we'll need a new approach.

Replacing upvalues

The debug.upvaluejoin function replaces what an upvalue refers to. You must specify another function and upvalue index to do the replacement.

By creating a new anonymous function with and upvalue set to what _ENV should be in lucky, the upvalue can be replaced without affecting everything else.

local new_env = {
  print = function()
    print("no lucky number for you!")
  end
}

debug.upvaluejoin(lucky, get_upvalue(lucky, "_ENV"),
  function() return new_env end, 1)

lucky() --> prints "no lucky number for you"

The last two arguments reference upvalue 1 of the anonymous function. Since the anonymous function only has one upvalue it’s guaranteed to be at index 1. It’s not necessary to look it up by name.

Now we're ready for a completed setfenv implementation.

setfenv implementation

local function setfenv(fn, env)
  local i = 1
  while true do
    local name = debug.getupvalue(fn, i)
    if name == "_ENV" then
      debug.upvaluejoin(fn, i, (function()
        return env
      end), 1)
      break
    elseif not name then
      break
    end

    i = i + 1
  end

  return fn
end

getfenv implementation

local function getfenv(fn)
  local i = 1
  while true do
    local name, val = debug.getupvalue(fn, i)
    if name == "_ENV" then
      return val
    elseif not name then
      break
    end
    i = i + 1
  end
end