New year's gift - a patch for Lua 5.3.5 implementing 'defer' statement as in Ravi

[Dibyendu Majumdar <mobile@...>]

Hi,

I recently created a release for Ravi that has an implementation of
the 'defer' statement. Please see
https://github.com/dibyendumajumdar/ravi/releases for details.

I am pleased to attach a patch for Lua 5.3.5 that adds the 'defer'
statement. Minimally tested, but happy to receive any reports of
issues. Of course many thanks to Roberto for his work on Lua 5.4, as
this patch is based upon that work.

Happy new year all!

Regards
Dibyendu

Index: lfunc.c
IDEA additional info:
Subsystem: com.intellij.openapi.diff.impl.patch.CharsetEP
<+>UTF-8
===================================================================
--- lfunc.c	(revision 26d9fe86009d7d89b2decb3d669459ad4ec1747e)
+++ lfunc.c	(date 1577916887765)
@@ -14,6 +14,7 @@
 
 #include "lua.h"
 
+#include "ldo.h"
 #include "lfunc.h"
 #include "lgc.h"
 #include "lmem.h"
@@ -61,13 +62,15 @@
   lua_assert(isintwups(L) || L->openupval == NULL);
   while (*pp != NULL && (p = *pp)->v >= level) {
     lua_assert(upisopen(p));
-    if (p->v == level)  /* found a corresponding upvalue? */
-      return p;  /* return it */
+    if (p->v == level && !p->flags)  /* found a corresponding upvalue that is not a deferred value? */ {
+      return p; /* return it */
+    }
     pp = &p->u.open.next;
   }
   /* not found: create a new upvalue */
   uv = luaM_new(L, UpVal);
   uv->refcount = 0;
+  uv->flags = 0;
   uv->u.open.next = *pp;  /* link it to list of open upvalues */
   uv->u.open.touched = 1;
   *pp = uv;
@@ -79,20 +82,84 @@
   return uv;
 }
 
+static void calldeferred(lua_State *L, void *ud) {
+  UNUSED(ud);
+  luaD_callnoyield(L, L->top - 2, 0);
+}
+
+/*
+** Prepare deferred function plus its arguments for object 'obj' with
+** error message 'err'. (This function assumes EXTRA_STACK.)
+*/
+static int preparetocall(lua_State *L, TValue *func, TValue *err) {
+  StkId top = L->top;
+  setobj2s(L, top, func);  /* will call deferred function */
+  if (err) {
+    setobj2s(L, top + 1, err); /* and error msg. as 1st argument */
+  }
+  else {
+    setnilvalue(top + 1);
+  }
+  L->top = top + 2;  /* add function and arguments */
+  return 1;
+}
 
-void luaF_close (lua_State *L, StkId level) {
+/*
+** Prepare and call a deferred function. If status is OK, code is still
+** inside the original protected call, and so any error will be handled
+** there. Otherwise, a previous error already activated the original
+** protected call, and so the call to the deferred method must be
+** protected here. (A status == -1 behaves like a previous
+** error, to also run the closing method in protected mode).
+** If status is OK, the call to the deferred method will be pushed
+** at the top of the stack. Otherwise, values are pushed after
+** the 'level' of the upvalue containing deferred function, as everything after
+** that won't be used again.
+*/
+static int calldeferredfunction(lua_State *L, StkId level, int status) {
+  TValue *uv = level; /* value being closed */
+  if (status == LUA_OK) {
+    preparetocall(L, uv, NULL); /* something to call? */
+    calldeferred(L, NULL);      /* call closing method */
+  }
+  else { /* must close the object in protected mode */
+    ptrdiff_t oldtop;
+    level++;                            /* space for error message */
+    oldtop = savestack(L, level + 1);   /* top will be after that */
+    luaD_seterrorobj(L, status, level); /* set error message */
+    preparetocall(L, uv, level);
+    int newstatus = luaD_pcall(L, calldeferred, NULL, oldtop, 0);
+    if (newstatus != LUA_OK && status == -1) /* first error? */
+      status = newstatus;                    /* this will be the new error */
+    else {
+      /* leave original error (or nil) on top */
+      L->top = restorestack(L, oldtop);
+    }
+  }
+  return status;
+}
+
+int luaF_close (lua_State *L, StkId level, int status) {
   UpVal *uv;
   while (L->openupval != NULL && (uv = L->openupval)->v >= level) {
     lua_assert(upisopen(uv));
     L->openupval = uv->u.open.next;  /* remove from 'open' list */
-    if (uv->refcount == 0)  /* no references? */
-      luaM_free(L, uv);  /* free upvalue */
+    if (uv->refcount == 0) {        /* no references? */
+      UpVal uv1 = *uv;              /* copy the upvalue as we will free it below */
+      luaM_free(L, uv);             /* free upvalue before invoking any deferred functions */
+      if (uv1.flags && ttisfunction(uv1.v)) {
+        ptrdiff_t levelrel = savestack(L, level);
+        status = calldeferredfunction(L, uv1.v, status);
+        level = restorestack(L, levelrel);
+      }
+    }
     else {
       setobj(L, &uv->u.value, uv->v);  /* move value to upvalue slot */
       uv->v = &uv->u.value;  /* now current value lives here */
       luaC_upvalbarrier(L, uv);
     }
   }
+  return status;
 }
 
 
Index: lparser.c
IDEA additional info:
Subsystem: com.intellij.openapi.diff.impl.patch.CharsetEP
<+>UTF-8
===================================================================
--- lparser.c	(revision 26d9fe86009d7d89b2decb3d669459ad4ec1747e)
+++ lparser.c	(date 1577916906509)
@@ -519,10 +519,17 @@
 ** so that, if it invokes the GC, the GC knows which registers
 ** are in use at that time.
 */
-static void codeclosure (LexState *ls, expdesc *v) {
+static void codeclosure (LexState *ls, expdesc *v, int deferred) {
   FuncState *fs = ls->fs->prev;
+  int pc = -1;
+  if (deferred) {
+    pc = luaK_codeABC(fs, OP_DEFER, 0, 0, 0);
+  }
   init_exp(v, VRELOCABLE, luaK_codeABx(fs, OP_CLOSURE, 0, fs->np - 1));
   luaK_exp2nextreg(fs, v);  /* fix it at the last register */
+  if (deferred) {
+    SETARG_A(fs->f->code[pc], v->u.info);
+  }
 }
 
 
@@ -779,24 +786,26 @@
 }
 
 
-static void body (LexState *ls, expdesc *e, int ismethod, int line) {
+static void body (LexState *ls, expdesc *e, int ismethod, int line, int deferred) {
   /* body ->  '(' parlist ')' block END */
   FuncState new_fs;
   BlockCnt bl;
   new_fs.f = addprototype(ls);
   new_fs.f->linedefined = line;
   open_func(ls, &new_fs, &bl);
-  checknext(ls, '(');
-  if (ismethod) {
-    new_localvarliteral(ls, "self");  /* create 'self' parameter */
-    adjustlocalvars(ls, 1);
-  }
-  parlist(ls);
-  checknext(ls, ')');
+  if (!deferred) {
+    checknext(ls, '(');
+    if (ismethod) {
+      new_localvarliteral(ls, "self");  /* create 'self' parameter */
+      adjustlocalvars(ls, 1);
+    }
+    parlist(ls);
+    checknext(ls, ')');
+  }
   statlist(ls);
   new_fs.f->lastlinedefined = ls->linenumber;
   check_match(ls, TK_END, TK_FUNCTION, line);
-  codeclosure(ls, e);
+  codeclosure(ls, e, deferred);
   close_func(ls);
 }
 
@@ -971,7 +980,7 @@
     }
     case TK_FUNCTION: {
       luaX_next(ls);
-      body(ls, v, 0, ls->linenumber);
+      body(ls, v, 0, ls->linenumber, 0);
       return;
     }
     default: {
@@ -1428,12 +1437,18 @@
 }
 
 
-static void localfunc (LexState *ls) {
+static void localfunc (LexState *ls, int defer) {
   expdesc b;
   FuncState *fs = ls->fs;
-  new_localvar(ls, str_checkname(ls));  /* new local variable */
+  if (defer) {
+    static const char funcname[] = "(deferred function)";
+    new_localvar(ls, luaX_newstring(ls, funcname, sizeof funcname-1));  /* new local variable */
+    markupval(fs, fs->nactvar);
+  } else {
+    new_localvar(ls, str_checkname(ls));  /* new local variable */
+  }
   adjustlocalvars(ls, 1);  /* enter its scope */
-  body(ls, &b, 0, ls->linenumber);  /* function created in next register */
+  body(ls, &b, 0, ls->linenumber, defer);  /* function created in next register */
   /* debug information will only see the variable after this point! */
   getlocvar(fs, b.u.info)->startpc = fs->pc;
 }
@@ -1479,7 +1494,7 @@
   expdesc v, b;
   luaX_next(ls);  /* skip FUNCTION */
   ismethod = funcname(ls, &v);
-  body(ls, &b, ismethod, line);
+  body(ls, &b, ismethod, line, 0);
   luaK_storevar(ls->fs, &v, &b);
   luaK_fixline(ls->fs, line);  /* definition "happens" in the first line */
 }
@@ -1571,10 +1586,15 @@
     case TK_LOCAL: {  /* stat -> localstat */
       luaX_next(ls);  /* skip LOCAL */
       if (testnext(ls, TK_FUNCTION))  /* local function? */
-        localfunc(ls);
+        localfunc(ls, 0);
       else
         localstat(ls);
       break;
+    }
+    case TK_DEFER: {  /* stat -> deferstat */
+      luaX_next(ls);  /* skip DEFER */
+      localfunc(ls, 1);
+      break;
     }
     case TK_DBCOLON: {  /* stat -> label */
       luaX_next(ls);  /* skip double colon */
Index: ldo.h
IDEA additional info:
Subsystem: com.intellij.openapi.diff.impl.patch.CharsetEP
<+>UTF-8
===================================================================
--- ldo.h	(revision 26d9fe86009d7d89b2decb3d669459ad4ec1747e)
+++ ldo.h	(date 1577913499995)
@@ -53,6 +53,7 @@
 
 LUAI_FUNC l_noret luaD_throw (lua_State *L, int errcode);
 LUAI_FUNC int luaD_rawrunprotected (lua_State *L, Pfunc f, void *ud);
+LUAI_FUNC void luaD_seterrorobj (lua_State *L, int errcode, StkId oldtop);
 
 #endif
 
Index: llex.c
IDEA additional info:
Subsystem: com.intellij.openapi.diff.impl.patch.CharsetEP
<+>UTF-8
===================================================================
--- llex.c	(revision 26d9fe86009d7d89b2decb3d669459ad4ec1747e)
+++ llex.c	(date 1577914263669)
@@ -40,7 +40,7 @@
 static const char *const luaX_tokens [] = {
     "and", "break", "do", "else", "elseif",
     "end", "false", "for", "function", "goto", "if",
-    "in", "local", "nil", "not", "or", "repeat",
+    "in", "local", "defer", "nil", "not", "or", "repeat",
     "return", "then", "true", "until", "while",
     "//", "..", "...", "==", ">=", "<=", "~=",
     "<<", ">>", "::", "<eof>",
Index: lopcodes.c
IDEA additional info:
Subsystem: com.intellij.openapi.diff.impl.patch.CharsetEP
<+>UTF-8
===================================================================
--- lopcodes.c	(revision 26d9fe86009d7d89b2decb3d669459ad4ec1747e)
+++ lopcodes.c	(date 1577914348993)
@@ -65,6 +65,7 @@
   "CLOSURE",
   "VARARG",
   "EXTRAARG",
+  "DEFER",
   NULL
 };
 
@@ -119,6 +120,7 @@
  ,opmode(0, 0, OpArgU, OpArgU, iABC)		/* OP_SETLIST */
  ,opmode(0, 1, OpArgU, OpArgN, iABx)		/* OP_CLOSURE */
  ,opmode(0, 1, OpArgU, OpArgN, iABC)		/* OP_VARARG */
- ,opmode(0, 0, OpArgU, OpArgU, iAx)		/* OP_EXTRAARG */
+ ,opmode(0, 0, OpArgU, OpArgU, iAx)		  /* OP_EXTRAARG */
+ ,opmode(0, 1, OpArgN, OpArgN, iABC)		/* OP_DEFER */
 };
 
Index: lstate.c
IDEA additional info:
Subsystem: com.intellij.openapi.diff.impl.patch.CharsetEP
<+>UTF-8
===================================================================
--- lstate.c	(revision 26d9fe86009d7d89b2decb3d669459ad4ec1747e)
+++ lstate.c	(date 1577915913519)
@@ -241,7 +241,7 @@
 
 static void close_state (lua_State *L) {
   global_State *g = G(L);
-  luaF_close(L, L->stack);  /* close all upvalues for this thread */
+  luaF_close(L, L->stack, -1);  /* close all upvalues for this thread */
   luaC_freeallobjects(L);  /* collect all objects */
   if (g->version)  /* closing a fully built state? */
     luai_userstateclose(L);
@@ -284,7 +284,7 @@
 
 void luaE_freethread (lua_State *L, lua_State *L1) {
   LX *l = fromstate(L1);
-  luaF_close(L1, L1->stack);  /* close all upvalues for this thread */
+  luaF_close(L1, L1->stack, -1);  /* close all upvalues for this thread */
   lua_assert(L1->openupval == NULL);
   luai_userstatefree(L, L1);
   freestack(L1);
Index: llex.h
IDEA additional info:
Subsystem: com.intellij.openapi.diff.impl.patch.CharsetEP
<+>UTF-8
===================================================================
--- llex.h	(revision 26d9fe86009d7d89b2decb3d669459ad4ec1747e)
+++ llex.h	(date 1577913638648)
@@ -27,7 +27,7 @@
   /* terminal symbols denoted by reserved words */
   TK_AND = FIRST_RESERVED, TK_BREAK,
   TK_DO, TK_ELSE, TK_ELSEIF, TK_END, TK_FALSE, TK_FOR, TK_FUNCTION,
-  TK_GOTO, TK_IF, TK_IN, TK_LOCAL, TK_NIL, TK_NOT, TK_OR, TK_REPEAT,
+  TK_GOTO, TK_IF, TK_IN, TK_LOCAL, TK_DEFER, TK_NIL, TK_NOT, TK_OR, TK_REPEAT,
   TK_RETURN, TK_THEN, TK_TRUE, TK_UNTIL, TK_WHILE,
   /* other terminal symbols */
   TK_IDIV, TK_CONCAT, TK_DOTS, TK_EQ, TK_GE, TK_LE, TK_NE,
Index: lvm.c
IDEA additional info:
Subsystem: com.intellij.openapi.diff.impl.patch.CharsetEP
<+>UTF-8
===================================================================
--- lvm.c	(revision 26d9fe86009d7d89b2decb3d669459ad4ec1747e)
+++ lvm.c	(date 1577916990044)
@@ -737,7 +737,7 @@
 /* execute a jump instruction */
 #define dojump(ci,i,e) \
   { int a = GETARG_A(i); \
-    if (a != 0) luaF_close(L, ci->u.l.base + a - 1); \
+    if (a != 0) Protect(luaF_close(L, ci->u.l.base + a - 1, LUA_OK)); \
     ci->u.l.savedpc += GETARG_sBx(i) + e; }
 
 /* for test instructions, execute the jump instruction that follows it */
@@ -1159,7 +1159,7 @@
           StkId lim = nci->u.l.base + getproto(nfunc)->numparams;
           int aux;
           /* close all upvalues from previous call */
-          if (cl->p->sizep > 0) luaF_close(L, oci->u.l.base);
+          if (cl->p->sizep > 0) Protect(luaF_close(L, oci->u.l.base, LUA_OK));
           /* move new frame into old one */
           for (aux = 0; nfunc + aux < lim; aux++)
             setobjs2s(L, ofunc + aux, nfunc + aux);
@@ -1175,7 +1175,7 @@
       }
       vmcase(OP_RETURN) {
         int b = GETARG_B(i);
-        if (cl->p->sizep > 0) luaF_close(L, base);
+        if (cl->p->sizep > 0) Protect(luaF_close(L, base, LUA_OK));
         b = luaD_poscall(L, ci, ra, (b != 0 ? b - 1 : cast_int(L->top - ra)));
         if (ci->callstatus & CIST_FRESH)  /* local 'ci' still from callee */
           return;  /* external invocation: return */
@@ -1313,6 +1313,12 @@
       vmcase(OP_EXTRAARG) {
         lua_assert(0);
         vmbreak;
+      }
+      vmcase(OP_DEFER) {
+        UpVal *up = luaF_findupval(L, ra); /* create new upvalue */
+        up->flags = 1;  /* mark it as deferred */
+        setnilvalue(ra);  /* initialize it with nil */
+        vmbreak;
       }
     }
   }
Index: lopcodes.h
IDEA additional info:
Subsystem: com.intellij.openapi.diff.impl.patch.CharsetEP
<+>UTF-8
===================================================================
--- lopcodes.h	(revision 26d9fe86009d7d89b2decb3d669459ad4ec1747e)
+++ lopcodes.h	(date 1577913734350)
@@ -230,11 +230,13 @@
 
 OP_VARARG,/*	A B	R(A), R(A+1), ..., R(A+B-2) = vararg		*/
 
-OP_EXTRAARG/*	Ax	extra (larger) argument for previous opcode	*/
+OP_EXTRAARG,/*	Ax	extra (larger) argument for previous opcode	*/
+OP_DEFER    /*  A   mark variable A "deferred"	  */
+
 } OpCode;
 
 
-#define NUM_OPCODES	(cast(int, OP_EXTRAARG) + 1)
+#define NUM_OPCODES	(cast(int, OP_DEFER) + 1)
 
 
 
Index: ldo.c
IDEA additional info:
Subsystem: com.intellij.openapi.diff.impl.patch.CharsetEP
<+>UTF-8
===================================================================
--- ldo.c	(revision 26d9fe86009d7d89b2decb3d669459ad4ec1747e)
+++ ldo.c	(date 1577916852909)
@@ -88,7 +88,7 @@
 };
 
 
-static void seterrorobj (lua_State *L, int errcode, StkId oldtop) {
+void luaD_seterrorobj (lua_State *L, int errcode, StkId oldtop) {
   switch (errcode) {
     case LUA_ERRMEM: {  /* memory error? */
       setsvalue2s(L, oldtop, G(L)->memerrmsg); /* reuse preregistered msg. */
@@ -121,7 +121,7 @@
     }
     else {  /* no handler at all; abort */
       if (g->panic) {  /* panic function? */
-        seterrorobj(L, errcode, L->top);  /* assume EXTRA_STACK */
+        luaD_seterrorobj(L, errcode, L->top);  /* assume EXTRA_STACK */
         if (L->ci->top < L->top)
           L->ci->top = L->top;  /* pushing msg. can break this invariant */
         lua_unlock(L);
@@ -584,8 +584,8 @@
   if (ci == NULL) return 0;  /* no recovery point */
   /* "finish" luaD_pcall */
   oldtop = restorestack(L, ci->extra);
-  luaF_close(L, oldtop);
-  seterrorobj(L, status, oldtop);
+  luaF_close(L, oldtop, status);
+  luaD_seterrorobj(L, status, oldtop);
   L->ci = ci;
   L->allowhook = getoah(ci->callstatus);  /* restore original 'allowhook' */
   L->nny = 0;  /* should be zero to be yieldable */
@@ -671,7 +671,7 @@
     }
     if (errorstatus(status)) {  /* unrecoverable error? */
       L->status = cast_byte(status);  /* mark thread as 'dead' */
-      seterrorobj(L, status, L->top);  /* push error message */
+      luaD_seterrorobj(L, status, L->top);  /* push error message */
       L->ci->top = L->top;
     }
     else lua_assert(status == L->status);  /* normal end or yield */
@@ -729,11 +729,12 @@
   status = luaD_rawrunprotected(L, func, u);
   if (status != LUA_OK) {  /* an error occurred? */
     StkId oldtop = restorestack(L, old_top);
-    luaF_close(L, oldtop);  /* close possible pending closures */
-    seterrorobj(L, status, oldtop);
     L->ci = old_ci;
     L->allowhook = old_allowhooks;
     L->nny = old_nny;
+    status = luaF_close(L, oldtop, status);  /* close possible pending closures */
+    oldtop = restorestack(L, old_top);
+    luaD_seterrorobj(L, status, oldtop);
     luaD_shrinkstack(L);
   }
   L->errfunc = old_errfunc;
Index: lfunc.h
IDEA additional info:
Subsystem: com.intellij.openapi.diff.impl.patch.CharsetEP
<+>UTF-8
===================================================================
--- lfunc.h	(revision 26d9fe86009d7d89b2decb3d669459ad4ec1747e)
+++ lfunc.h	(date 1577913564617)
@@ -34,7 +34,8 @@
 */
 struct UpVal {
   TValue *v;  /* points to stack or to its own value */
-  lu_mem refcount;  /* reference counter */
+  unsigned int refcount;  /* reference counter */
+  unsigned int flags; /* Used to mark deferred values */
   union {
     struct {  /* (when open) */
       UpVal *next;  /* linked list */
@@ -52,7 +53,7 @@
 LUAI_FUNC LClosure *luaF_newLclosure (lua_State *L, int nelems);
 LUAI_FUNC void luaF_initupvals (lua_State *L, LClosure *cl);
 LUAI_FUNC UpVal *luaF_findupval (lua_State *L, StkId level);
-LUAI_FUNC void luaF_close (lua_State *L, StkId level);
+LUAI_FUNC int luaF_close (lua_State *L, StkId level, int status);
 LUAI_FUNC void luaF_freeproto (lua_State *L, Proto *f);
 LUAI_FUNC const char *luaF_getlocalname (const Proto *func, int local_number,
                                          int pc);