Support flushing translation blocks and flush when we don't need count hook

include/uc_priv.h

@@ -139,6 +139,9 @@ typedef void (*uc_invalidate_tb_t)(struct uc_struct *uc, uint64_t start,
 // Request generating TB at given address
 typedef uc_err (*uc_gen_tb_t)(struct uc_struct *uc, uint64_t pc, uc_tb *out_tb);
 
+// tb flush
+typedef uc_tcg_flush_tlb uc_tb_flush_t;
+
 struct hook {
     int type;       // UC_HOOK_*
     int insn;       // instruction for HOOK_INSN
@@ -272,6 +275,7 @@ struct uc_struct {
     uc_tcg_flush_tlb tcg_flush_tlb;
     uc_invalidate_tb_t uc_invalidate_tb;
     uc_gen_tb_t uc_gen_tb;
+    uc_tb_flush_t tb_flush;
     uc_add_inline_hook_t add_inline_hook;
     uc_del_inline_hook_t del_inline_hook;
 

include/unicorn/unicorn.h

@@ -530,7 +530,10 @@ typedef enum uc_control_type {
     UC_CTL_TB_REQUEST_CACHE,
     // Invalidate a tb cache at a specific address
     // Write: @args = (uint64_t, uint64_t)
-    UC_CTL_TB_REMOVE_CACHE
+    UC_CTL_TB_REMOVE_CACHE,
+    // Invalidate all translation blocks.
+    // No arguments.
+    UC_CTL_TB_FLUSH
 
 } uc_control_type;
 
@@ -605,7 +608,7 @@ See sample_ctl.c for a detailed example.
     uc_ctl(uc, UC_CTL_WRITE(UC_CTL_TB_REMOVE_CACHE, 2), (address), (end))
 #define uc_ctl_request_cache(uc, address, tb)                                  \
     uc_ctl(uc, UC_CTL_READ_WRITE(UC_CTL_TB_REQUEST_CACHE, 2), (address), (tb))
-
+#define uc_ctl_flush_tlb(uc) uc_ctl(uc, UC_CTL_WRITE(UC_CTL_TB_FLUSH, 0))
 // Opaque storage for CPU context, used with uc_context_*()
 struct uc_context;
 typedef struct uc_context uc_context;

qemu/accel/tcg/translate-all.c

@@ -985,6 +985,10 @@ static void tb_htable_init(struct uc_struct *uc)
 }
 
 
+static void uc_tb_flush(struct uc_struct *uc) {
+    tb_flush(uc->cpu);
+}
+
 static void uc_invalidate_tb(struct uc_struct *uc, uint64_t start_addr, size_t len) 
 {
     tb_page_addr_t start, end;
@@ -1095,6 +1099,7 @@ void tcg_exec_init(struct uc_struct *uc, unsigned long tb_size)
     /* Invalidate / Cache TBs */
     uc->uc_invalidate_tb = uc_invalidate_tb;
     uc->uc_gen_tb = uc_gen_tb;
+    uc->tb_flush = uc_tb_flush;
 
     /* Inline hooks optimization */
     uc->add_inline_hook = uc_add_inline_hook;

uc.c

@@ -808,6 +808,9 @@ uc_err uc_emu_start(uc_engine *uc, uint64_t begin, uint64_t until,
     if (count <= 0 && uc->count_hook != 0) {
         uc_hook_del(uc, uc->count_hook);
         uc->count_hook = 0;
+
+        // In this case, we have to drop all translated blocks.
+        uc->tb_flush(uc);
     }
     // set up count hook to count instructions.
     if (count > 0 && uc->count_hook == 0) {
@@ -2300,6 +2303,16 @@ uc_err uc_ctl(uc_engine *uc, uc_control_type control, ...)
         break;
     }
 
+    case UC_CTL_TB_FLUSH:
+
+        UC_INIT(uc);
+
+        if (rw == UC_CTL_IO_WRITE) {
+            uc->tb_flush(uc);
+        } else {
+            err = UC_ERR_ARG;
+        }
+
     default:
         err = UC_ERR_ARG;
         break;