Fix undumped locs

src/lparser.c

@@ -752,6 +752,36 @@ static void open_func (LexState *ls, FuncState *fs, BlockCnt *bl) {
   enterblock(fs, bl, 0);
 }
 
+#ifdef USE_YK
+/*
+ * Identify loops in the function and insert yk locations there.
+ *
+ * This should only be called at a time when the bytecode instructions have
+ * been properly finalised.
+ */
+void assign_yklocs(lua_State *L, Proto *f, int num_insts) {
+  f->yklocs = luaM_newvectorchecked(L, num_insts, YkLocation);
+  for (int pc = 0; pc < num_insts; 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
 
 static void close_func (LexState *ls) {
   lua_State *L = ls->L;
@@ -773,32 +803,7 @@ static void close_func (LexState *ls) {
   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();
-    }
-  }
+  assign_yklocs(L, f, fs->pc);
 #endif
 }
 

src/lparser.h

@@ -166,6 +166,9 @@ typedef struct FuncState {
 LUAI_FUNC int luaY_nvarstack (FuncState *fs);
 LUAI_FUNC LClosure *luaY_parser (lua_State *L, ZIO *z, Mbuffer *buff,
                                  Dyndata *dyd, const char *name, int firstchar);
+#ifdef USE_YK
+LUAI_FUNC void assign_yklocs(lua_State *L, Proto *f, int num_insts);
+#endif
 
 
 #endif

src/lundump.c

@@ -23,6 +23,9 @@
 #include "lfunc.h"
 #include "lmem.h"
 #include "lobject.h"
+#ifdef USE_YK
+#include "lparser.h"
+#endif
 #include "lstring.h"
 #include "lundump.h"
 #include "lzio.h"
@@ -148,13 +151,11 @@ static void loadCode (LoadState *S, Proto *f) {
   int n = loadInt(S);
   f->code = luaM_newvectorchecked(S->L, n, Instruction);
   f->sizecode = n;
+  loadVector(S, f->code, n);
 #ifdef USE_YK
-  f->yklocs = luaM_newvectorchecked(S->L, n, YkLocation);
-  for (int i=0; i<n; i++)
-    f->yklocs[i] = yk_location_new();
-  f->sizeyklocs = n;
+   /* FIXME: Ideally we'd persist the locations across a dump+undump too */
+  assign_yklocs(S->L, f, n);
 #endif
-  loadVector(S, f->code, n);
 }