1 files changed, 73 insertions, 53 deletions

diff --git a/lcode.c b/lcode.c
index 054b28fd..1872ede2 100644
--- a/lcode.c
+++ b/lcode.c
@@ -371,10 +371,6 @@ int luaK_codeABCk (FuncState *fs, OpCode o, int a, int b, int c, int k) {
 }
 
 
-#define codeABsC(fs,o,a,b,c,k)  luaK_codeABCk(fs,o,a,b,((c) + OFFSET_sC),k)
-
-
-
 /*
 ** Format and emit an 'iABx' instruction.
 */
@@ -885,30 +881,45 @@ void luaK_exp2val (FuncState *fs, expdesc *e) {
 
 
 /*
+** Try to make 'e' a K expression with an index in the range of R/K
+** indices. Return true iff succeeded.
+*/
+static int luaK_exp2K (FuncState *fs, expdesc *e) {
+  if (!hasjumps(e)) {
+    int info;
+    switch (e->k) {  /* move constants to 'k' */
+      case VTRUE: info = boolK(fs, 1); break;
+      case VFALSE: info = boolK(fs, 0); break;
+      case VNIL: info = nilK(fs); break;
+      case VKINT: info = luaK_intK(fs, e->u.ival); break;
+      case VKFLT: info = luaK_numberK(fs, e->u.nval); break;
+      case VK: info = e->u.info; break;
+      default: return 0;  /* not a constant */
+    }
+    if (info <= MAXINDEXRK) {  /* does constant fit in 'argC'? */
+      e->k = VK;  /* make expression a 'K' expression */
+      e->u.info = info;
+      return 1;
+    }
+  }
+  /* else, expression doesn't fit; leave it unchanged */
+  return 0;
+}
+
+
+/*
 ** Ensures final expression result is in a valid R/K index
 ** (that is, it is either in a register or in 'k' with an index
 ** in the range of R/K indices).
-** Returns 1 if expression is K, 0 otherwise.
+** Returns 1 iff expression is K.
 */
 int luaK_exp2RK (FuncState *fs, expdesc *e) {
-  luaK_exp2val(fs, e);
-  switch (e->k) {  /* move constants to 'k' */
-    case VTRUE: e->u.info = boolK(fs, 1); goto vk;
-    case VFALSE: e->u.info = boolK(fs, 0); goto vk;
-    case VNIL: e->u.info = nilK(fs); goto vk;
-    case VKINT: e->u.info = luaK_intK(fs, e->u.ival); goto vk;
-    case VKFLT: e->u.info = luaK_numberK(fs, e->u.nval); goto vk;
-    case VK:
-     vk:
-      e->k = VK;
-      if (e->u.info <= MAXINDEXRK)  /* constant fits in 'argC'? */
-        return 1;
-      else break;
-    default: break;
+  if (luaK_exp2K(fs, e))
+    return 1;
+  else {  /* not a constant in the right range: put it in a register */
+    luaK_exp2anyreg(fs, e);
+    return 0;
   }
-  /* not a constant in the right range: put it in a register */
-  luaK_exp2anyreg(fs, e);
-  return 0;
 }
 
 
@@ -1232,8 +1243,19 @@ static void codeunexpval (FuncState *fs, OpCode op, expdesc *e, int line) {
 }
 
 
+/*
+** Emit code for binary expressions that "produce values"
+** (everything but logical operators 'and'/'or' and comparison
+** operators).
+** Expression to produce final result will be encoded in 'e1'.
+** Because 'luaK_exp2anyreg' can free registers, its calls must be
+** in "stack order" (that is, first on 'e2', which may have more
+** recent registers to be released).
+*/
 static void finishbinexpval (FuncState *fs, expdesc *e1, expdesc *e2,
-                             int pc, int line) {
+                             OpCode op, int v2, int k, int line) {
+  int v1 = luaK_exp2anyreg(fs, e1);
+  int pc = luaK_codeABCk(fs, op, 0, v1, v2, k);
   freeexps(fs, e1, e2);
   e1->u.info = pc;
   e1->k = VRELOC;  /* all those operations are relocatable */
@@ -1242,20 +1264,13 @@ static void finishbinexpval (FuncState *fs, expdesc *e1, expdesc *e2,
 
 
 /*
-** Emit code for binary expressions that "produce values"
-** (everything but logical operators 'and'/'or' and comparison
-** operators).
-** Expression to produce final result will be encoded in 'e1'.
-** Because 'luaK_exp2anyreg' can free registers, its calls must be
-** in "stack order" (that is, first on 'e2', which may have more
-** recent registers to be released).
+** Emit code for binary expressions that "produce values" over
+** two registers.
 */
 static void codebinexpval (FuncState *fs, OpCode op,
                            expdesc *e1, expdesc *e2, int line) {
   int v2 = luaK_exp2anyreg(fs, e2);  /* both operands are in registers */
-  int v1 = luaK_exp2anyreg(fs, e1);
-  int pc = luaK_codeABC(fs, op, 0, v1, v2);  /* generate opcode */
-  finishbinexpval(fs, e1, e2, pc, line);
+  finishbinexpval(fs, e1, e2, op, v2, 0, line);
 }
 
 
@@ -1264,41 +1279,48 @@ static void codebinexpval (FuncState *fs, OpCode op,
 */
 static void codebini (FuncState *fs, OpCode op,
                        expdesc *e1, expdesc *e2, int k, int line) {
-  int v2 = cast_int(e2->u.ival);  /* immediate operand */
-  int v1 = luaK_exp2anyreg(fs, e1);
-  int pc = codeABsC(fs, op, 0, v1, v2, k);  /* generate opcode */
-  finishbinexpval(fs, e1, e2, pc, line);
+  int v2 = cast_int(e2->u.ival) + OFFSET_sC;  /* immediate operand */
+  finishbinexpval(fs, e1, e2, op, v2, k, line);
+}
+
+
+static void swapexps (expdesc *e1, expdesc *e2) {
+  expdesc temp = *e1; *e1 = *e2; *e2 = temp;  /* swap 'e1' and 'e2' */
 }
 
 
 /*
 ** Code arithmetic operators ('+', '-', ...). If second operand is a
 ** constant in the proper range, use variant opcodes with immediate
-** operands.
+** operands or K operands.
 */
 static void codearith (FuncState *fs, OpCode op,
                        expdesc *e1, expdesc *e2, int flip, int line) {
-  if (!isSCint(e2))
-    codebinexpval(fs, op, e1, e2, line);  /* use standard operators */
-  else  /* use immediate operators */
+  if (isSCint(e2))  /* immediate operand? */
     codebini(fs, cast(OpCode, op - OP_ADD + OP_ADDI), e1, e2, flip, line);
-}
-
-
-static void swapexps (expdesc *e1, expdesc *e2) {
-  expdesc temp = *e1; *e1 = *e2; *e2 = temp;  /* swap 'e1' and 'e2' */
+  else if (tonumeral(e2, NULL) && luaK_exp2K(fs, e2)) {  /* K operand? */
+    int v2 = e2->u.info;  /* K index */
+    op = cast(OpCode, op - OP_ADD + OP_ADDK);
+    finishbinexpval(fs, e1, e2, op, v2, flip, line);
+  }
+  else {  /* 'e2' is neither an immediate nor a K operand */
+    if (flip)
+      swapexps(e1, e2);  /* back to original order */
+    codebinexpval(fs, op, e1, e2, line);  /* use standard operators */
+  }
 }
 
 
 /*
 ** Code commutative operators ('+', '*'). If first operand is a
-** constant, change order of operands to use immediate operator.
+** numeric constant, change order of operands to try to use an
+** immediate or K operator.
 */
 static void codecommutative (FuncState *fs, OpCode op,
                              expdesc *e1, expdesc *e2, int line) {
   int flip = 0;
-  if (isSCint(e1)) {
-    swapexps(e1, e2);
+  if (tonumeral(e1, NULL)) {  /* is first operand a numeric constant? */
+    swapexps(e1, e2);  /* change order */
     flip = 1;
   }
   codearith(fs, op, e1, e2, flip, line);
@@ -1312,7 +1334,7 @@ static void codecommutative (FuncState *fs, OpCode op,
 static void codebitwise (FuncState *fs, BinOpr opr,
                          expdesc *e1, expdesc *e2, int line) {
   int inv = 0;
-  int v1, v2, pc;
+  int v2;
   OpCode op;
   if (e1->k == VKINT && luaK_exp2RK(fs, e1)) {
     swapexps(e1, e2);  /* 'e2' will be the constant operand */
@@ -1323,12 +1345,10 @@ static void codebitwise (FuncState *fs, BinOpr opr,
     codebinexpval(fs, op, e1, e2, line);  /* all-register opcodes */
     return;
   }
-  v1 = luaK_exp2anyreg(fs, e1);
   v2 = e2->u.info;  /* index in K array */
   op = cast(OpCode, opr - OPR_BAND + OP_BANDK);
   lua_assert(ttisinteger(&fs->f->k[v2]));
-  pc = luaK_codeABCk(fs, op, 0, v1, v2, inv);
-  finishbinexpval(fs, e1, e2, pc, line);
+  finishbinexpval(fs, e1, e2, op, v2, inv, line);
 }