test: support upvalues in exec_lua

1 files changed, 158 insertions, 6 deletions

diff --git a/test/functional/testnvim.lua b/test/functional/testnvim.lua
index ac2ec89271..36e6b4d7bd 100644
--- a/test/functional/testnvim.lua
+++ b/test/functional/testnvim.lua
@@ -837,19 +837,171 @@ function M.exec_capture(code)
   return M.api.nvim_exec2(code, { output = true }).output
 end
 
+--- @param f function
+--- @return table<string,any>
+local function get_upvalues(f)
+  local i = 1
+  local upvalues = {} --- @type table<string,any>
+  while true do
+    local n, v = debug.getupvalue(f, i)
+    if not n then
+      break
+    end
+    upvalues[n] = v
+    i = i + 1
+  end
+  return upvalues
+end
+
+--- @param f function
+--- @param upvalues table<string,any>
+local function set_upvalues(f, upvalues)
+  local i = 1
+  while true do
+    local n = debug.getupvalue(f, i)
+    if not n then
+      break
+    end
+    if upvalues[n] then
+      debug.setupvalue(f, i, upvalues[n])
+    end
+    i = i + 1
+  end
+end
+
+--- @type fun(f: function): table<string,any>
+_G.__get_upvalues = nil
+
+--- @type fun(f: function, upvalues: table<string,any>)
+_G.__set_upvalues = nil
+
+--- @param self table<string,function>
+--- @param bytecode string
+--- @param upvalues table<string,any>
+--- @param ... any[]
+--- @return any[] result
+--- @return table<string,any> upvalues
+local function exec_lua_handler(self, bytecode, upvalues, ...)
+  local f = assert(loadstring(bytecode))
+  self.set_upvalues(f, upvalues)
+  local ret = { f(...) } --- @type any[]
+  --- @type table<string,any>
+  local new_upvalues = self.get_upvalues(f)
+
+  do -- Check return value types for better error messages
+    local invalid_types = {
+      ['thread'] = true,
+      ['function'] = true,
+      ['userdata'] = true,
+    }
+
+    for k, v in pairs(ret) do
+      if invalid_types[type(v)] then
+        error(
+          string.format(
+            "Return index %d with value '%s' of type '%s' cannot be serialized over RPC",
+            k,
+            tostring(v),
+            type(v)
+          )
+        )
+      end
+    end
+  end
+
+  return ret, new_upvalues
+end
+
+--- Execute Lua code in the wrapped Nvim session.
+---
+--- When `code` is passed as a function, it is converted into Lua byte code.
+---
+--- Direct upvalues are copied over, however upvalues contained
+--- within nested functions are not. Upvalues are also copied back when `code`
+--- finishes executing. See `:help lua-upvalue`.
+---
+--- Only types which can be serialized can be transferred over, e.g:
+--- `table`, `number`, `boolean`, `string`.
+---
+--- `code` runs with a different environment and thus will have a different global
+--- environment. See `:help lua-environments`.
+---
+--- Example:
+--- ```lua
+--- local upvalue1 = 'upvalue1'
+--- exec_lua(function(a, b, c)
+---   print(upvalue1, a, b, c)
+---   (function()
+---     print(upvalue2)
+---   end)()
+--- end, 'a', 'b', 'c'
+--- ```
+--- Prints:
+--- ```
+--- upvalue1 a b c
+--- nil
+--- ```
+---
+--- Not supported:
+--- ```lua
+--- local a = vim.uv.new_timer()
+--- exec_lua(function()
+---   print(a) -- Error: a is of type 'userdata' which cannot be serialized.
+--- end)
+--- ```
 --- @param code string|function
+--- @param ... any
 --- @return any
 function M.exec_lua(code, ...)
-  if type(code) == 'function' then
-    return M.api.nvim_exec_lua(
+  if type(code) == 'string' then
+    return M.api.nvim_exec_lua(code, { ... })
+  end
+
+  assert(session)
+
+  if not session.exec_lua_setup then
+    M.api.nvim_exec_lua(
       [[
-      local code = ...
-      return loadstring(code)(select(2, ...))
+      _G.__test_exec_lua = {
+        get_upvalues = loadstring((select(1,...))),
+        set_upvalues = loadstring((select(2,...))),
+        handler = loadstring((select(3,...)))
+      }
+      setmetatable(_G.__test_exec_lua, { __index = _G.__test_exec_lua })
     ]],
-      { string.dump(code), ... }
+      { string.dump(get_upvalues), string.dump(set_upvalues), string.dump(exec_lua_handler) }
     )
+    session.exec_lua_setup = true
   end
-  return M.api.nvim_exec_lua(code, { ... })
+
+  --- @type any[], table<string,any>
+  local ret, upvalues = unpack(M.api.nvim_exec_lua(
+    [[
+      return {
+        _G.__test_exec_lua:handler(...)
+      }
+    ]],
+    {
+      string.dump(code),
+      get_upvalues(code),
+      ...,
+    }
+  ))
+
+  -- Update upvalues
+  if next(upvalues) then
+    local caller = debug.getinfo(2)
+    local f = caller.func
+    -- On PUC-Lua, if the function is a tail call, then func will be nil.
+    -- In this case we need to use the current function.
+    if not f then
+      assert(caller.source == '=(tail call)')
+      f = debug.getinfo(1).func
+    end
+    set_upvalues(f, upvalues)
+  end
+
+  return unpack(ret, 1, table.maxn(ret))
 end
 
 function M.get_pathsep()