From 6e9b719694bffb8de711f182d405ec37d32ae0b1 Mon Sep 17 00:00:00 2001
From: Roberto Ierusalimschy <roberto@inf.puc-rio.br>
Date: Fri, 26 Oct 2018 10:38:50 -0300
Subject: [PATCH] More uniformity in code generation for 'for' loops

Added new instruction 'OP_TFORPREP' to prepare a generic for loop.
Currently it is equivalent to a jump (but with a format 'iABx',
similar to other for-loop preparing instructions), but soon it will
be the place to create upvalues for closing loop states.
---
 ljumptab.h |  1 +
 lopcodes.c |  1 +
 lopcodes.h |  1 +
 lopnames.h |  1 +
 lparser.c  | 16 ++++++----------
 lvm.c      |  4 ++++
 6 files changed, 14 insertions(+), 10 deletions(-)

diff --git a/ljumptab.h b/ljumptab.h
index da4cf7b7..6767e95b 100644
--- a/ljumptab.h
+++ b/ljumptab.h
@@ -96,6 +96,7 @@ static void *disptab[] = {
 &&L_OP_FORPREP1,
 &&L_OP_FORLOOP,
 &&L_OP_FORPREP,
+&&L_OP_TFORPREP,
 &&L_OP_TFORCALL,
 &&L_OP_TFORLOOP,
 &&L_OP_SETLIST,
diff --git a/lopcodes.c b/lopcodes.c
index f6915beb..11a73c29 100644
--- a/lopcodes.c
+++ b/lopcodes.c
@@ -90,6 +90,7 @@ LUAI_DDEF const lu_byte luaP_opmodes[NUM_OPCODES] = {
  ,opmode(0, 0, 0, 1, iABx)		/* OP_FORPREP1 */
  ,opmode(0, 0, 0, 1, iABx)		/* OP_FORLOOP */
  ,opmode(0, 0, 0, 1, iABx)		/* OP_FORPREP */
+ ,opmode(0, 0, 0, 0, iABx)		/* OP_TFORPREP */
  ,opmode(0, 0, 0, 0, iABC)		/* OP_TFORCALL */
  ,opmode(0, 0, 0, 1, iABx)		/* OP_TFORLOOP */
  ,opmode(0, 1, 0, 0, iABC)		/* OP_SETLIST */
diff --git a/lopcodes.h b/lopcodes.h
index 4d144000..4797d7c3 100644
--- a/lopcodes.h
+++ b/lopcodes.h
@@ -282,6 +282,7 @@ OP_FORLOOP,/*	A Bx	R(A)+=R(A+2);
 			if R(A) <?= R(A+1) then { pc-=Bx; R(A+3)=R(A) }	*/
 OP_FORPREP,/*	A Bx	R(A)-=R(A+2); pc+=Bx				*/
 
+OP_TFORPREP,/*	A Bx	create upvalue A; pc+=Bx			*/
 OP_TFORCALL,/*	A C	R(A+3), ... ,R(A+2+C) := R(A)(R(A+1), R(A+2));	*/
 OP_TFORLOOP,/*	A Bx	if R(A+1) ~= nil then { R(A)=R(A+1); pc -= Bx }	*/
 
diff --git a/lopnames.h b/lopnames.h
index 304d3b6c..b2c4fe21 100644
--- a/lopnames.h
+++ b/lopnames.h
@@ -81,6 +81,7 @@ static const char *const opnames[] = {
   "FORPREP1",
   "FORLOOP",
   "FORPREP",
+  "TFORPREP",
   "TFORCALL",
   "TFORLOOP",
   "SETLIST",
diff --git a/lparser.c b/lparser.c
index c0c40eae..c3dc6362 100644
--- a/lparser.c
+++ b/lparser.c
@@ -1350,30 +1350,26 @@ static void fixforjump (FuncState *fs, int pc, int dest, int back) {
 */
 static void forbody (LexState *ls, int base, int line, int nvars, int kind) {
   /* forbody -> DO block */
+  static OpCode forprep[3] = {OP_FORPREP, OP_FORPREP1, OP_TFORPREP};
+  static OpCode forloop[3] = {OP_FORLOOP, OP_FORLOOP1, OP_TFORLOOP};
   BlockCnt bl;
   FuncState *fs = ls->fs;
   int prep, endfor;
   adjustlocalvars(ls, 3);  /* control variables */
   checknext(ls, TK_DO);
-  prep = (kind == 0) ? luaK_codeABx(fs, OP_FORPREP, base, 0)
-       : (kind == 1) ? luaK_codeABx(fs, OP_FORPREP1, base, 0)
-       : luaK_jump(fs);
+  prep = luaK_codeABx(fs, forprep[kind], base, 0);
   enterblock(fs, &bl, 0);  /* scope for declared variables */
   adjustlocalvars(ls, nvars);
   luaK_reserveregs(fs, nvars);
   block(ls);
   leaveblock(fs);  /* end of scope for declared variables */
+  fixforjump(fs, prep, luaK_getlabel(fs), 0);
   if (kind == 2) {  /* generic for? */
-    luaK_patchtohere(fs, prep);
     luaK_codeABC(fs, OP_TFORCALL, base, 0, nvars);
     luaK_fixline(fs, line);
-    endfor = luaK_codeABx(fs, OP_TFORLOOP, base + 2, 0);
-  }
-  else {
-    fixforjump(fs, prep, luaK_getlabel(fs), 0);
-    endfor = (kind == 0) ? luaK_codeABx(fs, OP_FORLOOP, base, 0)
-                         : luaK_codeABx(fs, OP_FORLOOP1, base, 0);
+    base += 2;  /* base for 'OP_TFORLOOP' (skips function and state) */
   }
+  endfor = luaK_codeABx(fs, forloop[kind], base, 0);
   fixforjump(fs, endfor, prep + 1, 1);
   luaK_fixline(fs, line);
 }
diff --git a/lvm.c b/lvm.c
index 2a1ee175..35a58089 100644
--- a/lvm.c
+++ b/lvm.c
@@ -1744,6 +1744,10 @@ void luaV_execute (lua_State *L, CallInfo *ci) {
         pc += GETARG_Bx(i);
         vmbreak;
       }
+      vmcase(OP_TFORPREP) {
+        pc += GETARG_Bx(i);
+        vmbreak;
+      }
       vmcase(OP_TFORCALL) {
         StkId cb = ra + 3;  /* call base */
         setobjs2s(L, cb+2, ra+2);