KZT, fix: Preserve callback control-flow boundaries

accel/tcg/cpu-exec.c

@@ -672,6 +672,13 @@ inline void tb_add_jump(TranslationBlock *tb, int n,
 
     qemu_thread_jit_write();
     assert(n < ARRAY_SIZE(tb->jmp_list_next));
+#ifdef CONFIG_LATX_KZT
+    /* The KZT callback owns this entry slot; never relink over its patch. */
+    if (!tb->canlink[n] && tb->jmp_target_arg[n] == 0 &&
+        tb->jmp_reset_offset[n] == TB_JMP_RESET_OFFSET_INVALID) {
+        return;
+    }
+#endif
     qemu_spin_lock(&tb_next->jmp_lock);
 
     /* make sure the destination TB is valid */
@@ -864,7 +871,14 @@ TranslationBlock * kzt_tb_find_exp(
             TranslationBlock *last_tb,
             int tb_exit, uint32_t cflags)
 {
-    return tb_find(cpu, last_tb, tb_exit, cflags);
+    CPUArchState *env = cpu->env_ptr;
+    target_ulong pc, cs_base;
+    uint32_t flags;
+
+    /* The callback installer owns mmap_lock and explicitly generates a fresh
+     * JIT TB when necessary; do not load AOT or publish unpatched code here. */
+    cpu_get_tb_cpu_state(env, &pc, &cs_base, &flags);
+    return tb_lookup(cpu, pc, cs_base, flags, cflags);
 }
 
 #endif

accel/tcg/translate-all.c

@@ -1921,6 +1921,12 @@ TranslationBlock *tb_gen_code(CPUState *cpu,
 #endif
     void *host_pc;
 
+#ifdef CONFIG_LATX_KZT
+    bool kzt_private = cflags & CF_KZT_PRIVATE;
+
+    cflags &= ~CF_KZT_PRIVATE;
+#endif
+
     assert_memory_lock();
     qemu_thread_jit_write();
 
@@ -2321,6 +2327,13 @@ TranslationBlock *tb_gen_code(CPUState *cpu,
     }
 #endif
 
+#ifdef CONFIG_LATX_KZT
+    /* KZT patches this fresh TB and publishes it from its own installer. */
+    if (kzt_private) {
+        return tb;
+    }
+#endif
+
     tcg_tb_insert(tb);
 
     /*

include/exec/exec-all.h

@@ -600,6 +600,10 @@ struct TranslationBlock {
 #define CF_USE_ICOUNT  0x00020000
 #define CF_INVALID     0x00040000 /* TB is stale. Set with @jmp_lock held */
 #define CF_PARALLEL    0x00080000 /* Generate code for a parallel context */
+#ifdef CONFIG_LATX_KZT
+/* KZT installation request only; consumed before assigning TB identity. */
+#define CF_KZT_PRIVATE 0x00100000 /* Return an unpublished TB to the installer */
+#endif
 #define CF_CLUSTER_MASK 0xff000000 /* Top 8 bits are cluster ID */
 #define CF_CLUSTER_SHIFT 24
 

linux-user/i386/cpu_loop.c

@@ -22,6 +22,9 @@
 #include "qemu.h"
 #include "cpu_loop-common.h"
 #include "latx-options.h"
+#if defined(CONFIG_LATX_KZT)
+#include "callback.h"
+#endif
 #include <linux/audit.h>
 
 /***********************************************************/
@@ -222,7 +225,8 @@ void cpu_loop(CPUX86State *env)
         process_queued_cpu_work(cs);
 #if defined(CONFIG_LATX_KZT)
         kzt_reclaim_unloaded_headers();
-        if (latx_kzt_runtime_enabled() && trapnr == 0xCC)
+        if (latx_kzt_runtime_enabled() && trapnr == 0xCC &&
+            env->eip == (uintptr_t)&RunFunctionWithState)
             break;
 #endif
         switch(trapnr) {