Vararg table

Not yet optimized nor documented.

9 files changed, 80 insertions, 31 deletions

diff --git a/lobject.h b/lobject.h
index cc3dd370..a805dcbf 100644
--- a/lobject.h
+++ b/lobject.h
@@ -583,8 +583,10 @@ typedef struct AbsLineInfo {
 /*
 ** Flags in Prototypes
 */
-#define PF_ISVARARG     1
-#define PF_FIXED        2  /* prototype has parts in fixed memory */
+#define PF_ISVARARG     1  /* function is vararg */
+#define PF_VATAB        2  /* function is vararg with table */
+#define PF_VAPTAB       4  /* function is vararg with pseudo-table */
+#define PF_FIXED        8  /* prototype has parts in fixed memory */
 
 
 /*
diff --git a/lopcodes.h b/lopcodes.h
index 9ad21021..c3f7f64d 100644
--- a/lopcodes.h
+++ b/lopcodes.h
@@ -338,7 +338,7 @@ OP_CLOSURE,/*	A Bx	R[A] := closure(KPROTO[Bx])			*/
 
 OP_VARARG,/*    A C     R[A], R[A+1], ..., R[A+C-2] = vararg            */
 
-OP_VARARGPREP,/*A       (adjust vararg parameters)                      */
+OP_VARARGPREP,/*        (adjust vararg parameters)                      */
 
 OP_EXTRAARG/*   Ax      extra (larger) argument for previous opcode     */
 } OpCode;
diff --git a/lparser.c b/lparser.c
index 5abcd407..f7e78793 100644
--- a/lparser.c
+++ b/lparser.c
@@ -1041,9 +1041,10 @@ static void constructor (LexState *ls, expdesc *t) {
 /* }====================================================================== */
 
 
-static void setvararg (FuncState *fs, int nparams) {
-  fs->f->flag |= PF_ISVARARG;
-  luaK_codeABC(fs, OP_VARARGPREP, nparams, 0, 0);
+static void setvararg (FuncState *fs, int kind) {
+  lua_assert(kind & PF_ISVARARG);
+  fs->f->flag |= cast_byte(kind);
+  luaK_codeABC(fs, OP_VARARGPREP, 0, 0, 0);
 }
 
 
@@ -1052,7 +1053,7 @@ static void parlist (LexState *ls) {
   FuncState *fs = ls->fs;
   Proto *f = fs->f;
   int nparams = 0;
-  int isvararg = 0;
+  int varargk = 0;
   if (ls->t.token != ')') {  /* is 'parlist' not empty? */
     do {
       switch (ls->t.token) {
@@ -1062,19 +1063,27 @@ static void parlist (LexState *ls) {
           break;
         }
         case TK_DOTS: {
+          varargk |= PF_ISVARARG;
           luaX_next(ls);
-          isvararg = 1;
+          if (testnext(ls, '=')) {
+            new_varkind(ls, str_checkname(ls), RDKVATAB);
+            varargk |= PF_VATAB;
+          }
           break;
         }
         default: luaX_syntaxerror(ls, "<name> or '...' expected");
       }
-    } while (!isvararg && testnext(ls, ','));
+    } while (!varargk && testnext(ls, ','));
   }
   adjustlocalvars(ls, nparams);
   f->numparams = cast_byte(fs->nactvar);
-  if (isvararg)
-    setvararg(fs, f->numparams);  /* declared vararg */
-  luaK_reserveregs(fs, fs->nactvar);  /* reserve registers for parameters */
+  if (varargk != 0) {
+    setvararg(fs, varargk);  /* declared vararg */
+    if (varargk & PF_VATAB)
+      adjustlocalvars(ls, 1);  /* vararg table */
+  }
+  /* reserve registers for parameters (and vararg variable, if present) */
+  luaK_reserveregs(fs, fs->nactvar);
 }
 
 
@@ -2099,7 +2108,7 @@ static void mainfunc (LexState *ls, FuncState *fs) {
   BlockCnt bl;
   Upvaldesc *env;
   open_func(ls, fs, &bl);
-  setvararg(fs, 0);  /* main function is always declared vararg */
+  setvararg(fs, PF_ISVARARG);  /* main function is always vararg */
   env = allocupvalue(fs);  /* ...set environment upvalue */
   env->instack = 1;
   env->idx = 0;
diff --git a/lparser.h b/lparser.h
index fdbb9b8a..e479905e 100644
--- a/lparser.h
+++ b/lparser.h
@@ -97,10 +97,11 @@ typedef struct expdesc {
 /* kinds of variables */
 #define VDKREG          0   /* regular local */
 #define RDKCONST        1   /* local constant */
-#define RDKTOCLOSE      2   /* to-be-closed */
-#define RDKCTC          3   /* local compile-time constant */
-#define GDKREG          4   /* regular global */
-#define GDKCONST        5   /* global constant */
+#define RDKVATAB        2   /* vararg table */
+#define RDKTOCLOSE      3   /* to-be-closed */
+#define RDKCTC          4   /* local compile-time constant */
+#define GDKREG          5   /* regular global */
+#define GDKCONST        6   /* global constant */
 
 /* variables that live in registers */
 #define varinreg(v)     ((v)->vd.kind <= RDKTOCLOSE)
diff --git a/ltm.c b/ltm.c
index d1a61a62..619be59e 100644
--- a/ltm.c
+++ b/ltm.c
@@ -224,11 +224,38 @@ int luaT_callorderiTM (lua_State *L, const TValue *p1, int v2,
 }
 
 
-void luaT_adjustvarargs (lua_State *L, int nfixparams, CallInfo *ci,
-                         const Proto *p) {
+/*
+** Create a vararg table at the top of the stack, with 'n' elements
+** starting at 'f'.
+*/
+static void createvarargtab (lua_State *L, StkId f, int n) {
   int i;
-  int actual = cast_int(L->top.p - ci->func.p) - 1;  /* number of arguments */
-  int nextra = actual - nfixparams;  /* number of extra arguments */
+  TValue key, value;
+  Table *t = luaH_new(L);
+  sethvalue(L, s2v(L->top.p), t);
+  L->top.p++;
+  luaH_resize(L, t, cast_uint(n), 1);
+  setsvalue(L, &key, luaS_new(L, "n"));  /* key is "n" */
+  setivalue(&value, n);  /* value is n */
+  /* No need to anchor the key: Due to the resize, the next operation
+     cannot trigger a garbage collection */
+  luaH_set(L, t, &key, &value);  /* t.n = n */
+  for (i = 0; i < n; i++)
+    luaH_setint(L, t, i + 1, s2v(f + i));
+}
+
+
+/*
+** initial stack:  func arg1 ... argn extra1 ...
+**                 ^ ci->func                    ^ L->top
+** final stack: func nil ... nil extra1 ... func arg1 ... argn
+**                                          ^ ci->func         ^ L->top
+*/
+void luaT_adjustvarargs (lua_State *L, CallInfo *ci, const Proto *p) {
+  int i;
+  int totalargs = cast_int(L->top.p - ci->func.p) - 1;
+  int nfixparams = p->numparams;
+  int nextra = totalargs - nfixparams;  /* number of extra arguments */
   ci->u.l.nextraargs = nextra;
   luaD_checkstack(L, p->maxstacksize + 1);
   /* copy function to the top of the stack */
@@ -238,8 +265,14 @@ void luaT_adjustvarargs (lua_State *L, int nfixparams, CallInfo *ci,
     setobjs2s(L, L->top.p++, ci->func.p + i);
     setnilvalue(s2v(ci->func.p + i));  /* erase original parameter (for GC) */
   }
-  ci->func.p += actual + 1;
-  ci->top.p += actual + 1;
+  if (p->flag & (PF_VAPTAB | PF_VATAB)) {  /* is there a vararg table? */
+    if (p->flag & PF_VAPTAB)  /* is vararg table fake? */
+      setnilvalue(s2v(L->top.p));  /* initialize it */
+    else
+      createvarargtab(L, ci->func.p + nfixparams + 1, nextra);
+  }
+  ci->func.p += totalargs + 1;
+  ci->top.p += totalargs + 1;
   lua_assert(L->top.p <= ci->top.p && ci->top.p <= L->stack_last.p);
 }
 
diff --git a/ltm.h b/ltm.h
index ba2e4760..ed479bb4 100644
--- a/ltm.h
+++ b/ltm.h
@@ -95,8 +95,8 @@ LUAI_FUNC int luaT_callorderTM (lua_State *L, const TValue *p1,
 LUAI_FUNC int luaT_callorderiTM (lua_State *L, const TValue *p1, int v2,
                                  int inv, int isfloat, TMS event);
 
-LUAI_FUNC void luaT_adjustvarargs (lua_State *L, int nfixparams,
-                                   struct CallInfo *ci, const Proto *p);
+LUAI_FUNC void luaT_adjustvarargs (lua_State *L, struct CallInfo *ci,
+                                                 const Proto *p);
 LUAI_FUNC void luaT_getvarargs (lua_State *L, struct CallInfo *ci,
                                               StkId where, int wanted);
 
diff --git a/lundump.c b/lundump.c
index 76f0ddc1..74839af8 100644
--- a/lundump.c
+++ b/lundump.c
@@ -327,7 +327,8 @@ static void loadFunction (LoadState *S, Proto *f) {
   f->linedefined = loadInt(S);
   f->lastlinedefined = loadInt(S);
   f->numparams = loadByte(S);
-  f->flag = loadByte(S) & PF_ISVARARG;  /* get only the meaningful flags */
+  /* get only the meaningful flags */
+  f->flag = cast_byte(loadByte(S) & ~PF_FIXED);
   if (S->fixed)
     f->flag |= PF_FIXED;  /* signal that code is fixed */
   f->maxstacksize = loadByte(S);
diff --git a/lvm.c b/lvm.c
index d0a1c05d..d88a80d1 100644
--- a/lvm.c
+++ b/lvm.c
@@ -1927,7 +1927,7 @@ void luaV_execute (lua_State *L, CallInfo *ci) {
         vmbreak;
       }
       vmcase(OP_VARARGPREP) {
-        ProtectNT(luaT_adjustvarargs(L, GETARG_A(i), ci, cl->p));
+        ProtectNT(luaT_adjustvarargs(L, ci, cl->p));
         if (l_unlikely(trap)) {  /* previous "Protect" updated trap */
           luaD_hookcall(L, ci);
           L->oldpc = 1;  /* next opcode will be seen as a "new" line */
diff --git a/testes/vararg.lua b/testes/vararg.lua
index 10553de2..4320684e 100644
--- a/testes/vararg.lua
+++ b/testes/vararg.lua
@@ -3,9 +3,12 @@
 
 print('testing vararg')
 
-local function f (a, ...)
+local function f (a, ...=t)
   local x = {n = select('#', ...), ...}
-  for i = 1, x.n do assert(a[i] == x[i]) end
+  assert(x.n == t.n)
+  for i = 1, x.n do
+    assert(a[i] == x[i] and x[i] == t[i])
+  end
   return x.n
 end
 
@@ -17,7 +20,7 @@ local function c12 (...)
   return res, 2
 end
 
-local function vararg (...) return {n = select('#', ...), ...} end
+local function vararg (...=t) return t end
 
 local call = function (f, args) return f(table.unpack(args, 1, args.n)) end
 
@@ -99,7 +102,7 @@ assert(a==nil and b==nil and c==nil and d==nil and e==nil)
 
 
 -- varargs for main chunks
-local f = load[[ return {...} ]]
+local f = assert(load[[ return {...} ]])
 local x = f(2,3)
 assert(x[1] == 2 and x[2] == 3 and x[3] == undef)