3 files changed, 18 insertions, 15 deletions

diff --git a/src/lib_jit.c b/src/lib_jit.c
index 4da079c8..0728fc3a 100644
--- a/src/lib_jit.c
+++ b/src/lib_jit.c
@@ -437,9 +437,6 @@ static int jitopt_flag(jit_State *J, const char *str)
   return 0;  /* No match. */
 }
 
-/* Forward declaration. */
-static void jit_init_hotcount(jit_State *J);
-
 /* Parse optimization parameter. */
 static int jitopt_param(jit_State *J, const char *str)
 {
@@ -453,7 +450,7 @@ static int jitopt_param(jit_State *J, const char *str)
         lj_str_numconv(&str[len+1], &tv)) {
       J->param[i] = lj_num2int(tv.n);
       if (i == JIT_P_hotloop)
-        jit_init_hotcount(J);
+        lj_dispatch_init_hotcount(J2G(J));
       return 1;  /* Ok. */
     }
     lst += 1+len;
@@ -498,16 +495,6 @@ JIT_PARAMDEF(JIT_PARAMINIT)
 #undef JIT_PARAMINIT
   0
 };
-
-/* Initialize hotcount table. */
-static void jit_init_hotcount(jit_State *J)
-{
-  HotCount start = (HotCount)J->param[JIT_P_hotloop];
-  HotCount *hotcount = J2GG(J)->hotcount;
-  uint32_t i;
-  for (i = 0; i < HOTCOUNT_SIZE; i++)
-    hotcount[i] = start;
-}
 #endif
 
 /* Arch-dependent CPU detection. */
@@ -570,7 +557,6 @@ static void jit_init(lua_State *L)
   jit_State *J = L2J(L);
   J->flags = flags | JIT_F_ON | JIT_F_OPT_DEFAULT;
   memcpy(J->param, jit_param_default, sizeof(J->param));
-  jit_init_hotcount(J);
   lj_dispatch_update(G(L));
 #else
   UNUSED(flags);
diff --git a/src/lj_dispatch.c b/src/lj_dispatch.c
index 1e5b574c..a94afea9 100644
--- a/src/lj_dispatch.c
+++ b/src/lj_dispatch.c
@@ -34,6 +34,18 @@ void lj_dispatch_init(GG_State *GG)
   disp[BC_LOOP] = disp[BC_ILOOP];
 }
 
+#if LJ_HASJIT
+/* Initialize hotcount table. */
+void lj_dispatch_init_hotcount(global_State *g)
+{
+  HotCount start = (HotCount)G2J(g)->param[JIT_P_hotloop];
+  HotCount *hotcount = G2GG(g)->hotcount;
+  uint32_t i;
+  for (i = 0; i < HOTCOUNT_SIZE; i++)
+    hotcount[i] = start;
+}
+#endif
+
 /* Update dispatch table depending on various flags. */
 void lj_dispatch_update(global_State *g)
 {
@@ -77,6 +89,10 @@ void lj_dispatch_update(global_State *g)
       disp[BC_ITERL] = f_iterl;
       disp[BC_LOOP] = f_loop;
     }
+#if LJ_HASJIT
+    if ((mode & 1) && !(oldmode & 1))  /* JIT off to on transition. */
+      lj_dispatch_init_hotcount(g);
+#endif
   }
 }
 
diff --git a/src/lj_dispatch.h b/src/lj_dispatch.h
index ce9d9e11..cb5e5f64 100644
--- a/src/lj_dispatch.h
+++ b/src/lj_dispatch.h
@@ -56,6 +56,7 @@ typedef struct GG_State {
 
 /* Dispatch table management. */
 LJ_FUNC void lj_dispatch_init(GG_State *GG);
+LJ_FUNC void lj_dispatch_init_hotcount(global_State *g);
 LJ_FUNC void lj_dispatch_update(global_State *g);
 
 /* Instruction dispatch callback for instr/line hooks or when recording. */