diff --git a/tcg/tcg.c b/tcg/tcg.c
index 4c8c2d2991..f1f1640c9b 100644
--- a/tcg/tcg.c
+++ b/tcg/tcg.c
@@ -2328,7 +2328,7 @@ void tcg_dump_op_count(FILE *f, fprintf_function cpu_fprintf)
 #endif
 
 
-int tcg_gen_code(TCGContext *s, tcg_insn_unit *gen_code_buf)
+int tcg_gen_code(TCGContext *s, TranslationBlock *tb)
 {
     int i, oi, oi_next, num_insns;
 
@@ -2387,8 +2387,8 @@ int tcg_gen_code(TCGContext *s, tcg_insn_unit *gen_code_buf)
 
     tcg_reg_alloc_start(s);
 
-    s->code_buf = gen_code_buf;
-    s->code_ptr = gen_code_buf;
+    s->code_buf = tb->tc_ptr;
+    s->code_ptr = tb->tc_ptr;
 
     tcg_out_tb_init(s);
 
diff --git a/tcg/tcg.h b/tcg/tcg.h
index b83f76351c..40c8fbe2ae 100644
--- a/tcg/tcg.h
+++ b/tcg/tcg.h
@@ -636,7 +636,7 @@ void tcg_context_init(TCGContext *s);
 void tcg_prologue_init(TCGContext *s);
 void tcg_func_start(TCGContext *s);
 
-int tcg_gen_code(TCGContext *s, tcg_insn_unit *gen_code_buf);
+int tcg_gen_code(TCGContext *s, TranslationBlock *tb);
 
 void tcg_set_frame(TCGContext *s, TCGReg reg, intptr_t start, intptr_t size);
 
diff --git a/translate-all.c b/translate-all.c
index e9f409b762..b3207fe274 100644
--- a/translate-all.c
+++ b/translate-all.c
@@ -1120,7 +1120,7 @@ TranslationBlock *tb_gen_code(CPUState *cpu,
        the tcg optimization currently hidden inside tcg_gen_code.  All
        that should be required is to flush the TBs, allocate a new TB,
        re-initialize it per above, and re-do the actual code generation.  */
-    gen_code_size = tcg_gen_code(&tcg_ctx, gen_code_buf);
+    gen_code_size = tcg_gen_code(&tcg_ctx, tb);
     if (unlikely(gen_code_size < 0)) {
         goto buffer_overflow;
     }