From 681297187ec45268e872b26753c441586c12bdd8 Mon Sep 17 00:00:00 2001
From: Roberto Ierusalimschy <roberto@inf.puc-rio.br>
Date: Fri, 16 Apr 2021 15:41:44 -0300
Subject: [PATCH] Bug: yielding in '__close' mess up number of returns

Yielding in a __close metamethod called when returning vararg results
changes the top and so messes up the number of returned values.
---
 lstate.h          |  2 +-
 lvm.c             | 12 +++++++++-
 testes/locals.lua | 59 +++++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 71 insertions(+), 2 deletions(-)

diff --git a/lstate.h b/lstate.h
index c1283bb6b..44cf939cb 100644
--- a/lstate.h
+++ b/lstate.h
@@ -165,7 +165,7 @@ typedef struct stringtable {
 ** - field 'nyield' is used only while a function is "doing" an
 ** yield (from the yield until the next resume);
 ** - field 'nres' is used only while closing tbc variables when
-** returning from a C function;
+** returning from a function;
 ** - field 'transferinfo' is used only during call/returnhooks,
 ** before the function starts or after it ends.
 */
diff --git a/lvm.c b/lvm.c
index 16e01d683..e4b1903e7 100644
--- a/lvm.c
+++ b/lvm.c
@@ -847,10 +847,19 @@ void luaV_finishOp (lua_State *L) {
       luaV_concat(L, total);  /* concat them (may yield again) */
       break;
     }
-    case OP_CLOSE:  case OP_RETURN: {  /* yielded closing variables */
+    case OP_CLOSE: {  /* yielded closing variables */
       ci->u.l.savedpc--;  /* repeat instruction to close other vars. */
       break;
     }
+    case OP_RETURN: {  /* yielded closing variables */
+      StkId ra = base + GETARG_A(inst);
+      /* adjust top to signal correct number of returns, in case the
+         return is "up to top" ('isIT') */
+      L->top = ra + ci->u2.nres;
+      /* repeat instruction to close other vars. and complete the return */
+      ci->u.l.savedpc--;
+      break;
+    }
     default: {
       /* only these other opcodes can yield */
       lua_assert(op == OP_TFORCALL || op == OP_CALL ||
@@ -1672,6 +1681,7 @@ void luaV_execute (lua_State *L, CallInfo *ci) {
           n = cast_int(L->top - ra);  /* get what is available */
         savepc(ci);
         if (TESTARG_k(i)) {  /* may there be open upvalues? */
+          ci->u2.nres = n;  /* save number of returns */
           if (L->top < ci->top)
             L->top = ci->top;
           luaF_close(L, base, CLOSEKTOP, 1);
diff --git a/testes/locals.lua b/testes/locals.lua
index 6aad5d253..6151f64d0 100644
--- a/testes/locals.lua
+++ b/testes/locals.lua
@@ -813,6 +813,65 @@ do
 end
 
 
+do
+  -- yielding inside closing metamethods while returning
+  -- (bug in 5.4.3)
+
+  local extrares    -- result from extra yield (if any)
+
+  local function check (body, extra, ...)
+    local t = table.pack(...)   -- expected returns
+    local co = coroutine.wrap(body)
+    if extra then
+      extrares = co()    -- runs until first (extra) yield
+    end
+    local res = table.pack(co())   -- runs until yield inside '__close'
+    assert(res.n == 2 and res[2] == nil)
+    local res2 = table.pack(co())   -- runs until end of function
+    assert(res2.n == t.n)
+    for i = 1, #t do
+      if t[i] == "x" then
+        assert(res2[i] == res[1])    -- value that was closed
+      else
+        assert(res2[i] == t[i])
+      end
+    end
+  end
+
+  local function foo ()
+    local x <close> = func2close(coroutine.yield)
+    local extra <close> = func2close(function (self)
+      assert(self == extrares)
+      coroutine.yield(100)
+    end)
+    extrares = extra
+    return table.unpack{10, x, 30}
+  end
+  check(foo, true, 10, "x", 30)
+  assert(extrares == 100)
+
+  local function foo ()
+    local x <close> = func2close(coroutine.yield)
+    return
+  end
+  check(foo, false)
+
+  local function foo ()
+    local x <close> = func2close(coroutine.yield)
+    local y, z = 20, 30
+    return x
+  end
+  check(foo, false, "x")
+
+  local function foo ()
+    local x <close> = func2close(coroutine.yield)
+    local extra <close> = func2close(coroutine.yield)
+    return table.unpack({}, 1, 100)   -- 100 nils
+  end
+  check(foo, true, table.unpack({}, 1, 100))
+
+end
+
 do
   -- yielding inside closing metamethods after an error