diff --git a/qemu/cputlb.c b/qemu/cputlb.c
index cf146b89..c194dad5 100644
--- a/qemu/cputlb.c
+++ b/qemu/cputlb.c
@@ -159,8 +159,9 @@ void cpu_tlb_reset_dirty_all(struct uc_struct *uc,
 
 /* update the TLB corresponding to virtual page vaddr
    so that it is no longer dirty */
-void tlb_set_dirty(CPUArchState *env, target_ulong vaddr)
+void tlb_set_dirty(CPUState *cpu, target_ulong vaddr)
 {
+    CPUArchState *env = cpu->env_ptr;
     int i;
     int mmu_idx;
 
diff --git a/qemu/exec.c b/qemu/exec.c
index 91a74c94..8f160000 100644
--- a/qemu/exec.c
+++ b/qemu/exec.c
@@ -1382,8 +1382,7 @@ static void notdirty_mem_write(struct uc_struct* uc, void *opaque, hwaddr ram_ad
     /* we remove the notdirty callback only if the code has been
        flushed */
     if (!cpu_physical_memory_is_clean(uc, ram_addr)) {
-        CPUArchState *env = uc->current_cpu->env_ptr;
-        tlb_set_dirty(env, uc->current_cpu->mem_io_vaddr);
+        tlb_set_dirty(uc->current_cpu, uc->current_cpu->mem_io_vaddr);
     }
 }
 
diff --git a/qemu/include/exec/cputlb.h b/qemu/include/exec/cputlb.h
index 0d8e6c16..0d0720cd 100644
--- a/qemu/include/exec/cputlb.h
+++ b/qemu/include/exec/cputlb.h
@@ -26,7 +26,7 @@ void tlb_unprotect_code(CPUState *cpu, ram_addr_t ram_addr);
 void tlb_reset_dirty_range(CPUTLBEntry *tlb_entry,
     uintptr_t start, uintptr_t length);
 void cpu_tlb_reset_dirty_all(struct uc_struct *uc, ram_addr_t start1, ram_addr_t length);
-void tlb_set_dirty(CPUArchState *env, target_ulong vaddr);
+void tlb_set_dirty(CPUState *env, target_ulong vaddr);
 //extern int tlb_flush_count;
 
 /* exec.c */