issue #185

include/lopcodes.h

@@ -242,8 +242,11 @@ OP_RAVI_FARRAY_SETF, /*	A B C	R(A)[RK(B)] := RK(C) where RK(B) is an int, R(A) i
 
 OP_RAVI_FORLOOP_IP,  /* As FORLOOP, but with integer index and positive integer step */
 OP_RAVI_FORLOOP_I1,  /* As FORLOOP, but with integer index 1 and step 1 */
+OP_RAVI_FORLOOP_I,   /* As FORLOOP, but with integer index 1 and integer step */
+
 OP_RAVI_FORPREP_IP,  /* As FORPREP, but with integer index and positive integer step */
 OP_RAVI_FORPREP_I1,  /* As FORPREP, but with integer index 1 and step 1 */
+OP_RAVI_FORPREP_I,   /* As FORPREP, but with integer index 1 and integer step */
 
 OP_RAVI_SETUPVALI,  /*	A B	UpValue[B] := tointeger(R(A))			*/
 OP_RAVI_SETUPVALF,  /*	A B	UpValue[B] := tonumber(R(A))			*/

src/lopcodes.c

@@ -131,8 +131,10 @@ LUAI_DDEF const char *const luaP_opnames[NUM_OPCODES+1] = {
 
   "FORLOOP_IP",
   "FORLOOP_I1",
+  "FORLOOP_I",
   "FORPREP_IP",
   "FORPREP_I1",
+  "FORPREP_I",
 
   "SETUPVALI", /*	A B	UpValue[B] := tointeger(R(A))			*/
   "SETUPVALF", /*	A B	UpValue[B] := tonumber(R(A))			*/
@@ -275,8 +277,10 @@ LUAI_DDEF const lu_byte luaP_opmodes[NUM_OPCODES] = {
 
  ,opmode(0, 1, OpArgR, OpArgN, iAsBx)		/* OP_RAVI_FORLOOP_IP */
  ,opmode(0, 1, OpArgR, OpArgN, iAsBx)		/* OP_RAVI_FORLOOP_I1 */
+ ,opmode(0, 1, OpArgR, OpArgN, iAsBx)		/* OP_RAVI_FORLOOP_I */
  ,opmode(0, 1, OpArgR, OpArgN, iAsBx)		/* OP_RAVI_FORPREP_IP */
  ,opmode(0, 1, OpArgR, OpArgN, iAsBx)		/* OP_RAVI_FORPREP_I1 */
+ ,opmode(0, 1, OpArgR, OpArgN, iAsBx)		/* OP_RAVI_FORPREP_I */
 
  ,opmode(0, 0, OpArgU, OpArgN, iABC)		/* OP_RAVI_SETUPVALI */
  ,opmode(0, 0, OpArgU, OpArgN, iABC)		/* OP_RAVI_SETUPVALF */
@@ -542,9 +546,11 @@ static void PrintCode(membuff_t *mb, const Proto* f)
    case OP_FORLOOP:
    case OP_RAVI_FORLOOP_IP:
    case OP_RAVI_FORLOOP_I1:
+   case OP_RAVI_FORLOOP_I:
    case OP_FORPREP:
    case OP_RAVI_FORPREP_IP:
    case OP_RAVI_FORPREP_I1:
+   case OP_RAVI_FORPREP_I:
    case OP_TFORLOOP:
      membuff_add_fstring(mb, "\t; to %d",sbx+pc+2);
     break;

src/lparser.c

@@ -2051,6 +2051,10 @@ static void forbody (LexState *ls, int base, int line, int nvars, int isnum, For
       forprep_inst = OP_RAVI_FORPREP_I1;
       forloop_inst = OP_RAVI_FORLOOP_I1;
     }
+    else if (info && info->type == RAVI_TNUMINT) {
+      forprep_inst = OP_RAVI_FORPREP_I;
+      forloop_inst = OP_RAVI_FORLOOP_I;
+    }
   }
   prep = isnum ? luaK_codeAsBx(fs, forprep_inst, base, NO_JUMP) : luaK_jump(fs);
   enterblock(fs, &bl, 0);  /* scope for declared variables */
@@ -2107,7 +2111,7 @@ static void fornum (LexState *ls, TString *varname, int line) {
   if (tidx.type == tlimit.type && tlimit.type == tstep.type && 
      (tidx.type == RAVI_TNUMFLT || tidx.type == RAVI_TNUMINT)) {
     LocVar *vidx, *vlimit, *vstep, *vvar;
-    if (tidx.type == RAVI_TNUMINT && tstep.is_constant)
+    if (tidx.type == RAVI_TNUMINT)
       info = &tstep;
     /* Note that as locvars may be reallocated while creating variables
        therefore we access the variables here */

src/lvm.c

@@ -1253,8 +1253,10 @@ int luaV_execute (lua_State *L) {
     &&vmlabel(OP_RAVI_FARRAY_SETF),
     &&vmlabel(OP_RAVI_FORLOOP_IP),
     &&vmlabel(OP_RAVI_FORLOOP_I1),
+    &&vmlabel(OP_RAVI_FORLOOP_I),
     &&vmlabel(OP_RAVI_FORPREP_IP),
     &&vmlabel(OP_RAVI_FORPREP_I1),
+    &&vmlabel(OP_RAVI_FORPREP_I),
     &&vmlabel(OP_RAVI_SETUPVALI),
     &&vmlabel(OP_RAVI_SETUPVALF),
     &&vmlabel(OP_RAVI_SETUPVAL_IARRAY),
@@ -2096,6 +2098,17 @@ int luaV_execute (lua_State *L) {
         }
         vmbreak;
       }
+      vmcase(OP_RAVI_FORLOOP_I) {
+          lua_Integer step = ivalue(ra + 2);
+          lua_Integer idx = intop(+, ivalue(ra), step); /* increment index */
+          lua_Integer limit = ivalue(ra + 1);
+          if (RAVI_LIKELY((0 < step)) ? (idx <= limit) : (limit <= idx)) {
+              pc += GETARG_sBx(i);  /* jump back */
+              chgivalue(ra, idx);  /* update internal index... */
+              setivalue(ra + 3, idx);  /* ...and external index */
+          }
+          vmbreak;
+      }
       vmcase(OP_RAVI_FORPREP_IP)
       vmcase(OP_RAVI_FORPREP_I1) {
         TValue *pinit = ra;
@@ -2109,6 +2122,15 @@ int luaV_execute (lua_State *L) {
         pc += GETARG_sBx(i);
         vmbreak;
       }
+      vmcase(OP_RAVI_FORPREP_I) {
+          TValue* init = ra;
+          TValue* pstep = ra + 2;
+          /* all values are integer */
+          lua_Integer initv = ivalue(init);
+          setivalue(init, intop(-, initv, ivalue(pstep)));
+          pc += GETARG_sBx(i);
+          vmbreak;
+      }
       vmcase(OP_RAVI_NEW_IARRAY) {
         Table *t;
         savepc(L);  /* in case of allocation errors */

src/ravi_jitshared.c

@@ -649,8 +649,10 @@ static void scan_jump_targets(struct function *fn) {
       case OP_JMP:
       case OP_RAVI_FORPREP_IP:
       case OP_RAVI_FORPREP_I1:
+      case OP_RAVI_FORPREP_I:
       case OP_RAVI_FORLOOP_IP:
       case OP_RAVI_FORLOOP_I1:
+      case OP_RAVI_FORLOOP_I:
       case OP_FORLOOP:
       case OP_FORPREP:
       case OP_TFORLOOP: {
@@ -1678,6 +1680,7 @@ static void emit_op_iforprep(struct function *fn, int A, int pc, int step_one, i
     membuff_add_fstring(&fn->prologue, "lua_Number nlimit_%d = 0.0;\n", A);
     membuff_add_fstring(&fn->prologue, "lua_Number nstep_%d = 0.0;\n", A);
     membuff_add_fstring(&fn->prologue, "int intloop_%d = 0;\n", A);
+    membuff_add_fstring(&fn->prologue, "int is_positive_step_%d = 0;\n", A);
   }
   emit_reg(fn, "ra", A);
   membuff_add_fstring(&fn->body, "i_%d = ivalue(ra);\n", A);
@@ -1694,6 +1697,33 @@ static void emit_op_iforprep(struct function *fn, int A, int pc, int step_one, i
   membuff_add_fstring(&fn->body, "goto Lbc_%d;\n", pc);
 }
 
+static void emit_op_iforprep2(struct function* fn, int A, int pc, int pc1) {
+    (void)pc1;
+    if (!fn->locals[A]) {
+        fn->locals[A] = 1;  // Lua can reuse the same forloop vars if loop isn't nested
+        // Although in IFOR instructions we do not need all the vars below
+        // it can happen that the same slot is used by normal FOR loop
+        // The optimizer should get rid of unused vars anyway
+        membuff_add_fstring(&fn->prologue, "lua_Integer i_%d = 0;\n", A);
+        membuff_add_fstring(&fn->prologue, "lua_Integer limit_%d = 0;\n", A);
+        membuff_add_fstring(&fn->prologue, "lua_Integer step_%d = 0;\n", A);
+        membuff_add_fstring(&fn->prologue, "lua_Number ninit_%d = 0.0;\n", A);
+        membuff_add_fstring(&fn->prologue, "lua_Number nlimit_%d = 0.0;\n", A);
+        membuff_add_fstring(&fn->prologue, "lua_Number nstep_%d = 0.0;\n", A);
+        membuff_add_fstring(&fn->prologue, "int intloop_%d = 0;\n", A);
+        membuff_add_fstring(&fn->prologue, "int is_positive_step_%d = 0;\n", A);
+    }
+    emit_reg(fn, "ra", A);
+    membuff_add_fstring(&fn->body, "i_%d = ivalue(ra);\n", A);
+    membuff_add_fstring(&fn->body, "ra = R(%d);\n", A + 1);
+    membuff_add_fstring(&fn->body, "limit_%d = ivalue(ra);\n", A);
+    membuff_add_fstring(&fn->body, "ra = R(%d);\n", A + 2);
+    membuff_add_fstring(&fn->body, "step_%d = ivalue(ra);\n", A);
+    membuff_add_fstring(&fn->body, "is_positive_step_%d = 0 < step_%d;\n", A, A);
+    membuff_add_fstring(&fn->body, "i_%d -= step_%d;\n", A, A);
+    membuff_add_fstring(&fn->body, "goto Lbc_%d;\n", pc);
+}
+
 static void emit_op_forprep(struct function *fn, int A, int pc, int pc1) {
   (void)pc1;
   if (!fn->locals[A]) {
@@ -1705,6 +1735,7 @@ static void emit_op_forprep(struct function *fn, int A, int pc, int pc1) {
     membuff_add_fstring(&fn->prologue, "lua_Number nlimit_%d = 0.0;\n", A);
     membuff_add_fstring(&fn->prologue, "lua_Number nstep_%d = 0.0;\n", A);
     membuff_add_fstring(&fn->prologue, "int intloop_%d = 0;\n", A);
+    membuff_add_fstring(&fn->prologue, "int is_positive_step_%d = 0;\n", A);
   }
   emit_reg(fn, "ra", A);  // init
   membuff_add_string(&fn->body, "rb = ra+1; /*limit*/\n");
@@ -1765,6 +1796,19 @@ static void emit_op_iforloop(struct function *fn, int A, int pc, int step_one, i
                       A, A, A + 3, A, pc);
 }
 
+static void emit_op_iforloop2(struct function* fn, int A, int pc, int pc1) {
+    (void)pc1;
+    membuff_add_fstring(&fn->body, "i_%d += step_%d;\n", A, A);
+    membuff_add_fstring(&fn->body, "if (is_positive_step_%d) goto Lpos%d; else goto Lneg%d;\n", A, pc, pc);
+    membuff_add_fstring(&fn->body, "Lpos%d:\n", pc);
+    membuff_add_fstring(&fn->body, "if (i_%d > limit_%d) goto Lend%d; else goto Lbody%d;\n", A, A, pc, pc);
+    membuff_add_fstring(&fn->body, "Lneg%d:\n", pc);
+    membuff_add_fstring(&fn->body, "if (i_%d < limit_%d) goto Lend%d; else goto Lbody%d;\n", A, A, pc, pc);
+    membuff_add_fstring(&fn->body, "Lbody%d:\n", pc);
+    membuff_add_fstring(&fn->body, "ra = R(%d);\n   setivalue(ra, i_%d);\n   goto Lbc_%d;\n", A + 3, A, pc);
+    membuff_add_fstring(&fn->body,  "Lend%d:\n", pc);
+}
+
 static void emit_op_forloop(struct function *fn, int A, int pc, int pc1) {
   (void)pc1;
   membuff_add_fstring(&fn->body, "if (intloop_%d) {\n", A);
@@ -1867,12 +1911,22 @@ bool raviJ_codegen(struct lua_State *L, struct Proto *p, struct ravi_compile_opt
         int B = GETARG_B(i);
         emit_op_return(&fn, A, B, pc);
       } break;
+      case OP_RAVI_FORPREP_I: {
+        int sbx = GETARG_sBx(i);
+        int j = sbx + pc + 1;
+        emit_op_iforprep2(&fn, A, j, pc);
+      } break;
       case OP_RAVI_FORPREP_I1:
       case OP_RAVI_FORPREP_IP: {
         int sbx = GETARG_sBx(i);
         int j = sbx + pc + 1;
         emit_op_iforprep(&fn, A, j, op == OP_RAVI_FORPREP_I1, pc);
       } break;
+      case OP_RAVI_FORLOOP_I: {
+        int sbx = GETARG_sBx(i);
+        int j = sbx + pc + 1;
+        emit_op_iforloop2(&fn, A, j, pc);
+      } break;
       case OP_RAVI_FORLOOP_I1:
       case OP_RAVI_FORLOOP_IP: {
         int sbx = GETARG_sBx(i);

src/ravi_llvmcodegen.cpp

@@ -1623,6 +1623,7 @@ bool RaviCodeGenerator::compile(lua_State *L, Proto *p,
         int j = sbx + pc + 1;
         emit_iFORPREP(def, A, j, op == OP_RAVI_FORPREP_I1, pc);
       } break;
+      case OP_RAVI_FORPREP_I:
       case OP_FORPREP: {
         int sbx = GETARG_sBx(i);
         int j = sbx + pc + 1;
@@ -1639,6 +1640,7 @@ bool RaviCodeGenerator::compile(lua_State *L, Proto *p,
         emit_iFORLOOP(def, A, j, def->jmp_targets[pc], op == OP_RAVI_FORLOOP_I1,
                       pc);
       } break;
+      case OP_RAVI_FORLOOP_I: 
       case OP_FORLOOP: {
         int sbx = GETARG_sBx(i);
         int j = sbx + pc + 1;
@@ -2013,8 +2015,10 @@ void RaviCodeGenerator::scan_jump_targets(RaviFunctionDef *def, Proto *p) {
       case OP_JMP:
       case OP_RAVI_FORPREP_IP:
       case OP_RAVI_FORPREP_I1:
+      case OP_RAVI_FORPREP_I:
       case OP_RAVI_FORLOOP_IP:
       case OP_RAVI_FORLOOP_I1:
+      case OP_RAVI_FORLOOP_I:
       case OP_FORLOOP:
       case OP_FORPREP:
       case OP_TFORLOOP: {
@@ -2023,10 +2027,10 @@ void RaviCodeGenerator::scan_jump_targets(RaviFunctionDef *def, Proto *p) {
         if (op == OP_JMP)
           targetname = "jmp";
         else if (op == OP_FORLOOP || op == OP_RAVI_FORLOOP_IP ||
-                 op == OP_RAVI_FORLOOP_I1)
+                 op == OP_RAVI_FORLOOP_I1 || op == OP_RAVI_FORLOOP_I)
           targetname = "forbody";
         else if (op == OP_FORPREP || op == OP_RAVI_FORPREP_IP ||
-                 op == OP_RAVI_FORPREP_I1)
+                 op == OP_RAVI_FORPREP_I1 || op == OP_RAVI_FORPREP_I)
 #if RAVI_CODEGEN_FORPREP2
           targetname = "forloop_ilt";
 #else