Fix locs

src/lcode.c

@@ -380,24 +380,13 @@ static void removelastinstruction (FuncState *fs) {
 ** line information. Return 'i' position.
 */
 int luaK_code (FuncState *fs, Instruction i) {
-  int pc;
   Proto *f = fs->f;
   /* put new instruction in code array */
   luaM_growvector(fs->ls->L, f->code, fs->pc, f->sizecode, Instruction,
                   MAX_INT, "opcodes");
-  pc = fs->pc++;
-  f->code[pc] = i;
-#ifdef USE_YK
-  luaM_growvector(fs->ls->L, f->yklocs, pc, f->sizeyklocs, YkLocation,
-                  MAX_INT, "yklocs");
-  f->yklocs[pc] = yk_location_null();
-  if ((GET_OPCODE(i) == OP_JMP) && (GETARG_sJ(i) < 0))
-    f->yklocs[pc + GETARG_sJ(i)] = yk_location_new();
-  if (GET_OPCODE(i) == OP_FORLOOP)
-    f->yklocs[pc - GETARG_Bx(i) - 2] = yk_location_new();
-#endif
+  f->code[fs->pc++] = i;
   savelineinfo(fs, f, fs->ls->lastline);
-  return pc;  /* index of new instruction */
+  return fs->pc - 1;  /* index of new instruction */
 }
 
 

src/lparser.c

@@ -771,6 +771,35 @@ static void close_func (LexState *ls) {
   luaM_shrinkvector(L, f->upvalues, f->sizeupvalues, fs->nups, Upvaldesc);
   ls->fs = fs->prev;
   luaC_checkGC(L);
+
+#ifdef USE_YK
+  /*
+   * Identify loops in the Lua program and insert yk locations there.
+   *
+   * We do this here (and not in, e.g. `luaK_code()`), to ensure that bytecode
+   * operands (required to find loop starts) have been finalised.
+   */
+  f->yklocs = luaM_newvectorchecked(L, fs->pc, YkLocation);
+  for (int pc = 0; pc < fs->pc; pc++) {
+    Instruction i = f->code[pc];
+    f->yklocs[pc] = yk_location_null();
+    /*
+     * The computations for finding the start of loops is derived from
+     * `PrintCode()` in `luac.c`. Note that we have to deduct one because luac
+     * prints bytecode pcs starting from 1.
+     *
+     * The assertions below check that inserting a null location will never
+     * overwrite a non-null location in a later iteration of this loop.
+     */
+    if ((GET_OPCODE(i) == OP_JMP) && (GETARG_sJ(i) < 0)) {
+      lua_assert(GETARG_sJ(i) + pc + 2 - 1 < pc);
+      f->yklocs[GETARG_sJ(i) + pc + 2 - 1] = yk_location_new();
+    } else if (GET_OPCODE(i) == OP_FORLOOP) {
+      lua_assert(pc - GETARG_Bx(i) + 2 - 1 < pc);
+      f->yklocs[pc - GETARG_Bx(i) + 2 - 1] = yk_location_new();
+    }
+  }
+#endif
 }
 
 

src/luac.c

@@ -25,6 +25,10 @@
 #include "lstate.h"
 #include "lundump.h"
 
+#ifdef USE_YK
+#include <yk.h>
+#endif
+
 static void PrintFunction(const Proto* f, int full);
 #define luaU_print	PrintFunction
 
@@ -349,6 +353,13 @@ static void PrintCode(const Proto* f)
   int sbx=GETARG_sBx(i);
   int isk=GETARG_k(i);
   int line=luaG_getfuncline(f,pc);
+#ifdef USE_YK
+  if (yk_location_is_null(f->yklocs[pc])) {
+    printf("      ");
+  } else {
+    printf("ykloc:");
+  }
+#endif
   printf("\t%d\t",pc+1);
   if (line>0) printf("[%d]\t",line); else printf("[-]\t");
   printf("%-9s\t",opnames[o]);