diff --git a/cpu-exec.c b/cpu-exec.c
index a1b8b96..a984b7c 100644
--- a/cpu-exec.c
+++ b/cpu-exec.c
@@ -85,7 +85,7 @@ static void cpu_exec_nocache(CPUArchState *env, int max_cycles,
     if (max_cycles > CF_COUNT_MASK)
         max_cycles = CF_COUNT_MASK;
 
-    tb = tb_gen_code(env, orig_tb->pc, orig_tb->cs_base, orig_tb->flags,
+    tb = tb_gen_code(cpu, orig_tb->pc, orig_tb->cs_base, orig_tb->flags,
                      max_cycles);
     cpu->current_tb = tb;
     /* execute the generated code */
@@ -138,7 +138,7 @@ static TranslationBlock *tb_find_slow(CPUArchState *env,
     }
  not_found:
    /* if no translated code available, then translate it now */
-    tb = tb_gen_code(env, pc, cs_base, flags, 0);
+    tb = tb_gen_code(cpu, pc, cs_base, flags, 0);
 
  found:
     /* Move the last found TB to the head of the list */
diff --git a/exec.c b/exec.c
index 0a81967..1b5a0cd 100644
--- a/exec.c
+++ b/exec.c
@@ -1486,7 +1486,7 @@ static void check_watchpoint(int offset, int len_mask, int flags)
                     cpu_loop_exit(cpu);
                 } else {
                     cc->get_tb_cpu_state(cpu, &pc, &cs_base, &cpu_flags);
-                    tb_gen_code(env, pc, cs_base, cpu_flags, 1);
+                    tb_gen_code(cpu, pc, cs_base, cpu_flags, 1);
                     cpu_resume_from_signal(env, NULL);
                 }
             }
diff --git a/hw/i386/kvmvapic.c b/hw/i386/kvmvapic.c
index 5757a1b..bcf8d9d 100644
--- a/hw/i386/kvmvapic.c
+++ b/hw/i386/kvmvapic.c
@@ -450,7 +450,7 @@ static void patch_instruction(VAPICROMState *s, X86CPU *cpu, target_ulong ip)
 
     if (!kvm_enabled()) {
         cs->current_tb = NULL;
-        tb_gen_code(env, current_pc, current_cs_base, current_flags, 1);
+        tb_gen_code(cs, current_pc, current_cs_base, current_flags, 1);
         cpu_resume_from_signal(env, NULL);
     }
 }
diff --git a/include/exec/exec-all.h b/include/exec/exec-all.h
index 40d5931..73abf9e 100644
--- a/include/exec/exec-all.h
+++ b/include/exec/exec-all.h
@@ -84,7 +84,7 @@ bool cpu_restore_state(CPUState *cpu, uintptr_t searched_pc);
 
 void QEMU_NORETURN cpu_resume_from_signal(CPUArchState *env1, void *puc);
 void QEMU_NORETURN cpu_io_recompile(CPUState *cpu, uintptr_t retaddr);
-TranslationBlock *tb_gen_code(CPUArchState *env, 
+TranslationBlock *tb_gen_code(CPUState *cpu,
                               target_ulong pc, target_ulong cs_base, int flags,
                               int cflags);
 void cpu_exec_init(CPUArchState *env);
diff --git a/translate-all.c b/translate-all.c
index 1158681..6e5bfee 100644
--- a/translate-all.c
+++ b/translate-all.c
@@ -933,10 +933,11 @@ static void build_page_bitmap(PageDesc *p)
     }
 }
 
-TranslationBlock *tb_gen_code(CPUArchState *env,
+TranslationBlock *tb_gen_code(CPUState *cpu,
                               target_ulong pc, target_ulong cs_base,
                               int flags, int cflags)
 {
+    CPUArchState *env = cpu->env_ptr;
     TranslationBlock *tb;
     uint8_t *tc_ptr;
     tb_page_addr_t phys_pc, phys_page2;
@@ -1110,7 +1111,7 @@ void tb_invalidate_phys_page_range(tb_page_addr_t start, tb_page_addr_t end,
            modifying the memory. It will ensure that it cannot modify
            itself */
         cpu->current_tb = NULL;
-        tb_gen_code(env, current_pc, current_cs_base, current_flags, 1);
+        tb_gen_code(cpu, current_pc, current_cs_base, current_flags, 1);
         cpu_resume_from_signal(env, NULL);
     }
 #endif
@@ -1209,7 +1210,7 @@ static void tb_invalidate_phys_page(tb_page_addr_t addr,
            modifying the memory. It will ensure that it cannot modify
            itself */
         cpu->current_tb = NULL;
-        tb_gen_code(env, current_pc, current_cs_base, current_flags, 1);
+        tb_gen_code(cpu, current_pc, current_cs_base, current_flags, 1);
         if (locked) {
             mmap_unlock();
         }
@@ -1482,7 +1483,7 @@ void cpu_io_recompile(CPUState *cpu, uintptr_t retaddr)
     tb_phys_invalidate(tb, -1);
     /* FIXME: In theory this could raise an exception.  In practice
        we have already translated the block once so it's probably ok.  */
-    tb_gen_code(env, pc, cs_base, flags, cflags);
+    tb_gen_code(cpu, pc, cs_base, flags, cflags);
     /* TODO: If env->pc != tb->pc (i.e. the faulting instruction was not
        the first in the TB) then we end up generating a whole new TB and
        repeating the fault, which is horribly inefficient.
