enable notdirty_write for snapshots when possible

qemu/accel/tcg/cputlb.c

@@ -1192,9 +1192,7 @@ static void notdirty_write(CPUState *cpu, vaddr mem_vaddr, unsigned size,
 {
     ram_addr_t ram_addr = mem_vaddr + iotlbentry->addr;
 
-    if (mr == NULL) {
-        mr = cpu->uc->memory_mapping(cpu->uc, mem_vaddr);
-    }
+    assert(mr);
 
     if ((mr->perms & UC_PROT_EXEC) != 0) {
         struct page_collection *pages
@@ -1208,7 +1206,7 @@ static void notdirty_write(CPUState *cpu, vaddr mem_vaddr, unsigned size,
     // - have memory hooks installed
     // - or doing snapshot
     // , then never clean the tlb
-    if (!(cpu->uc->snapshot_level > 0 || mr->priority > 0) && 
+    if (!(mr->priority < cpu->uc->snapshot_level) &&
             !(HOOK_EXISTS(cpu->uc, UC_HOOK_MEM_READ) || HOOK_EXISTS(cpu->uc, UC_HOOK_MEM_WRITE))) {
         tlb_set_dirty(cpu, mem_vaddr);
     }
@@ -1232,6 +1230,8 @@ void *probe_access(CPUArchState *env, target_ulong addr, int size,
     target_ulong tlb_addr;
     size_t elt_ofs = 0;
     int wp_access = 0;
+    MemoryRegion *mr;
+    target_ulong paddr;
 
 #ifdef _MSC_VER
     g_assert(((target_ulong)0 - (addr | TARGET_PAGE_MASK)) >= size);
@@ -1288,7 +1288,9 @@ void *probe_access(CPUArchState *env, target_ulong addr, int size,
 
         /* Handle clean RAM pages.  */
         if (tlb_addr & TLB_NOTDIRTY) {
-            notdirty_write(env_cpu(env), addr, size, iotlbentry, retaddr, NULL);
+            paddr = entry->paddr | (addr & ~TARGET_PAGE_MASK);
+            mr = env->uc->memory_mapping(env->uc, paddr);
+            notdirty_write(env_cpu(env), addr, size, iotlbentry, retaddr, mr);
         }
     }
 
@@ -1362,6 +1364,8 @@ static void *atomic_mmu_lookup(CPUArchState *env, target_ulong addr,
     int a_bits = get_alignment_bits(mop);
     int s_bits = mop & MO_SIZE;
     void *hostaddr;
+    MemoryRegion *mr;
+    target_ulong paddr;
 
     /* Adjust the given return address.  */
     retaddr -= GETPC_ADJ;
@@ -1413,8 +1417,10 @@ static void *atomic_mmu_lookup(CPUArchState *env, target_ulong addr,
     hostaddr = (void *)((uintptr_t)addr + tlbe->addend);
 
     if (unlikely(tlb_addr & TLB_NOTDIRTY)) {
+        paddr = tlbe->paddr | (addr & ~TARGET_PAGE_MASK);
+        mr = env->uc->memory_mapping(env->uc, paddr);
         notdirty_write(env_cpu(env), addr, 1 << s_bits,
-                       &env_tlb(env)->d[mmu_idx].iotlb[index], retaddr, NULL);
+                       &env_tlb(env)->d[mmu_idx].iotlb[index], retaddr, mr);
     }
 
     return hostaddr;

uc.c

@@ -2147,6 +2147,7 @@ uc_err uc_context_save(uc_engine *uc, uc_context *context)
             restore_jit_state(uc);
             return ret;
         }
+        uc->tcg_flush_tlb(uc);
     }
 
     context->snapshot_level = uc->snapshot_level;
@@ -2418,6 +2419,7 @@ uc_err uc_context_restore(uc_engine *uc, uc_context *context)
             return ret;
         }
         uc_snapshot(uc);
+        uc->tcg_flush_tlb(uc);
     }
 
     if (uc->context_content & UC_CTL_CONTEXT_CPU) {