From c764ca71a639f5585b5f466bea25dc42b855a4b0 Mon Sep 17 00:00:00 2001
From: Roberto Ierusalimschy <roberto@inf.puc-rio.br>
Date: Mon, 25 Apr 2022 14:42:51 -0300
Subject: [PATCH] Bug: Wrong code generation in bitwise operations

---
 lcode.c               | 16 ++++++++++------
 ltests.h              |  7 +++++++
 testes/constructs.lua | 25 +++++++++++++++++++++++++
 3 files changed, 42 insertions(+), 6 deletions(-)

diff --git a/lcode.c b/lcode.c
index 06425a1db..cb724a090 100644
--- a/src/lcode.c
+++ b/src/lcode.c
@@ -1391,7 +1391,10 @@ static void finishbinexpval (FuncState *fs, expdesc *e1, expdesc *e2,
 */
 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 v2 = luaK_exp2anyreg(fs, e2);  /* make sure 'e2' is in a register */
+  /* 'e1' must be already in a register or it is a constant */
+  lua_assert((VNIL <= e1->k && e1->k <= VKSTR) ||
+             e1->k == VNONRELOC || e1->k == VRELOC);
   lua_assert(OP_ADD <= op && op <= OP_SHR);
   finishbinexpval(fs, e1, e2, op, v2, 0, line, OP_MMBIN,
                   cast(TMS, (op - OP_ADD) + TM_ADD));
@@ -1478,7 +1481,7 @@ static void codecommutative (FuncState *fs, BinOpr op,
 
 
 /*
-** Code bitwise operations; they are all associative, so the function
+** Code bitwise operations; they are all commutative, so the function
 ** tries to put an integer constant as the 2nd operand (a K operand).
 */
 static void codebitwise (FuncState *fs, BinOpr opr,
@@ -1486,11 +1489,11 @@ static void codebitwise (FuncState *fs, BinOpr opr,
   int flip = 0;
   int v2;
   OpCode op;
-  if (e1->k == VKINT && luaK_exp2RK(fs, e1)) {
+  if (e1->k == VKINT && luaK_exp2K(fs, e1)) {
     swapexps(e1, e2);  /* 'e2' will be the constant operand */
     flip = 1;
   }
-  else if (!(e2->k == VKINT && luaK_exp2RK(fs, e2))) {  /* no constants? */
+  else if (!(e2->k == VKINT && luaK_exp2K(fs, e2))) {  /* no constants? */
     op = cast(OpCode, opr + OP_ADD);
     codebinexpval(fs, op, e1, e2, line);  /* all-register opcodes */
     return;
@@ -1551,7 +1554,7 @@ static void codeeq (FuncState *fs, BinOpr opr, expdesc *e1, expdesc *e2) {
     op = OP_EQI;
     r2 = im;  /* immediate operand */
   }
-  else if (luaK_exp2RK(fs, e2)) {  /* 1st expression is constant? */
+  else if (luaK_exp2RK(fs, e2)) {  /* 2nd expression is constant? */
     op = OP_EQK;
     r2 = e2->u.info;  /* constant index */
   }
@@ -1611,7 +1614,8 @@ void luaK_infix (FuncState *fs, BinOpr op, expdesc *v) {
     case OPR_SHL: case OPR_SHR: {
       if (!tonumeral(v, NULL))
         luaK_exp2anyreg(fs, v);
-      /* else keep numeral, which may be folded with 2nd operand */
+      /* else keep numeral, which may be folded or used as an immediate
+         operand */
       break;
     }
     case OPR_EQ: case OPR_NE: {
diff --git a/ltests.h b/ltests.h
index cb3a0b480..ec520498b 100644
--- a/testes/ltests/ltests.h
+++ b/testes/ltests/ltests.h
@@ -125,6 +125,13 @@ LUA_API void *debug_realloc (void *ud, void *block,
 #define LUAI_USER_ALIGNMENT_T   union { char b[sizeof(void*) * 8]; }
 
 
+/*
+** This one is not compatible with tests for opcode optimizations,
+** as it blocks some optimizations
+#define MAXINDEXRK	0
+*/
+
+
 /* make stack-overflow tests run faster */
 #undef LUAI_MAXSTACK
 #define LUAI_MAXSTACK   50000
diff --git a/testes/constructs.lua b/testes/constructs.lua
index a74a8c041..0d9ec92d7 100644
--- a/testes/constructs.lua
+++ b/testes/constructs.lua
@@ -103,6 +103,31 @@ do  -- test old bug (first name could not be an `upvalue')
  local a; function f(x) x={a=1}; x={x=1}; x={G=1} end
 end
 
+
+do   -- bug since 5.4.0
+  -- create code with a table using more than 256 constants
+  local code = {"local x = {"}
+  for i = 1, 257 do
+    code[#code + 1] = i .. ".1,"
+  end
+  code[#code + 1] = "};"
+  code = table.concat(code)
+
+  -- add "ret" to the end of that code and checks that
+  -- it produces the expected value "val"
+  local function check (ret, val)
+    local code = code .. ret
+    code = load(code)
+    assert(code() == val)
+  end
+
+  check("return (1 ~ (2 or 3))", 1 ~ 2)
+  check("return (1 | (2 or 3))", 1 | 2)
+  check("return (1 + (2 or 3))", 1 + 2)
+  check("return (1 << (2 or 3))", 1 << 2)
+end
+
+
 function f (i)
   if type(i) ~= 'number' then return i,'jojo'; end;
   if i > 0 then return i, f(i-1); end;