Add Arm64 TSAN support and fix JIT cache coherency

This commit adds ThreadSanitizer (TSAN) support for ARM64/Apple Silicon
and fixes critical JIT instruction cache coherency issues.

ARM64 TSAN Support:
- Extended TSAN-compatible memory allocation to ARM64 architecture
- Main memory allocated at fixed address 0x150000000000 (21TB)
- JIT buffer allocated at 0x151000000000 with MAP_JIT for Apple Silicon
- Both allocations avoid TSAN shadow memory and enable race detection
- Note: Requires ASLR disabled on macOS (SIP restrictions may apply)

JIT Cache Coherency Fixes:
1. Fixed pthread_jit_write_protect_np() ordering in update_branch_imm
2. Added sys_icache_invalidate() in update_branch_imm
3. Added cache invalidation in resolve_jumps() for x86_64

Author: jserv

src/emulate.c

@@ -1205,10 +1205,14 @@ void rv_step(void *arg)
 #endif
         ) {
             jit_translate(rv, block);
-            ((exec_block_func_t) state->buf)(
-                rv, (uintptr_t) (state->buf + block->offset));
-            prev = NULL;
-            continue;
+            /* Only execute if translation succeeded (block is hot) */
+            if (block->hot) {
+                ((exec_block_func_t) state->buf)(
+                    rv, (uintptr_t) (state->buf + block->offset));
+                prev = NULL;
+                continue;
+            }
+            /* Fall through to interpreter if translation failed */
         }
         set_reset(&pc_set);
         has_loops = false;

src/io.c

@@ -27,18 +27,32 @@ memory_t *memory_new(uint32_t size)
         return NULL;
     assert(mem);
 #if HAVE_MMAP
-#if defined(TSAN_ENABLED) && defined(__x86_64__)
+#if defined(TSAN_ENABLED)
     /* ThreadSanitizer compatibility: Use MAP_FIXED to allocate at a specific
-     * address within TSAN's app range (0x7cf000000000 - 0x7ffffffff000).
+     * address to avoid conflicts with TSAN's shadow memory.
+     */
+#if defined(__x86_64__)
+    /* x86_64: Allocate within TSAN's range (0x7cf000000000 - 0x7ffffffff000).
      *
      * Fixed address: 0x7d0000000000
      * Size: up to 4GB (0x100000000)
      * End: 0x7d0100000000 (well within app range)
-     *
-     * This guarantees the allocation won't land in TSAN's shadow memory,
-     * preventing "unexpected memory mapping" errors.
      */
     void *fixed_addr = (void *) 0x7d0000000000UL;
+#elif defined(__aarch64__)
+    /* ARM64 (macOS/Apple Silicon): Use higher address range.
+     *
+     * Fixed address: 0x150000000000 (21TB)
+     * Size: up to 4GB (0x100000000)
+     * End: 0x150100000000
+     *
+     * This avoids TSAN's shadow memory and typical process allocations.
+     * Requires ASLR disabled via: setarch $(uname -m) -R
+     */
+    void *fixed_addr = (void *) 0x150000000000UL;
+#else
+#error "TSAN is only supported on x86_64 and aarch64"
+#endif
     data_memory_base = mmap(fixed_addr, size, PROT_READ | PROT_WRITE,
                             MAP_PRIVATE | MAP_ANONYMOUS | MAP_FIXED, -1, 0);
     if (data_memory_base == MAP_FAILED) {

src/jit.c

@@ -593,6 +593,10 @@ static void update_branch_imm(struct jit_state *state,
     assert((imm & 3) == 0);
     uint32_t insn;
     imm >>= 2;
+#if defined(__APPLE__) && defined(__aarch64__)
+    /* Must be in write mode to read/write MAP_JIT memory on Apple ARM64 */
+    pthread_jit_write_protect_np(false);
+#endif
     memcpy(&insn, state->buf + offset, sizeof(uint32_t));
     if ((insn & 0xfe000000U) == 0x54000000U /* Conditional branch immediate. */
         || (insn & 0x7e000000U) ==
@@ -607,10 +611,8 @@ static void update_branch_imm(struct jit_state *state,
         assert(false);
         insn = BAD_OPCODE;
     }
-#if defined(__APPLE__) && defined(__aarch64__)
-    pthread_jit_write_protect_np(false);
-#endif
     memcpy(state->buf + offset, &insn, sizeof(uint32_t));
+    sys_icache_invalidate(state->buf + offset, sizeof(uint32_t));
 #if defined(__APPLE__) && defined(__aarch64__)
     pthread_jit_write_protect_np(true);
 #endif
@@ -2167,6 +2169,7 @@ static void code_cache_flush(struct jit_state *state, riscv_t *rv)
     should_flush = false;
     state->offset = state->org_size;
     state->n_blocks = 0;
+    state->n_jumps = 0; /* Reset jump count when flushing */
     set_reset(&state->set);
     clear_cache_hot(rv->block_cache, (clear_func_t) clear_hot);
 #if RV32_HAS(T2C)
@@ -2229,6 +2232,7 @@ static void resolve_jumps(struct jit_state *state)
 
         uint8_t *offset_ptr = &state->buf[jump.offset_loc];
         memcpy(offset_ptr, &rel, sizeof(uint32_t));
+        sys_icache_invalidate(offset_ptr, sizeof(uint32_t));
 #elif defined(__aarch64__)
         int32_t rel = target_loc - jump.offset_loc;
         update_branch_imm(state, jump.offset_loc, rel);
@@ -2320,6 +2324,9 @@ void jit_translate(riscv_t *rv, block_t *block)
     block->offset = state->offset;
     translate_chained_block(state, rv, block);
     if (unlikely(should_flush)) {
+        /* Mark block as not translated since translation was incomplete */
+        block->hot = false;
+        /* Don't reset offset - it will be set correctly on restart */
         code_cache_flush(state, rv);
         goto restart;
     }
@@ -2336,10 +2343,12 @@ struct jit_state *jit_state_init(size_t size)
 
     state->offset = 0;
     state->size = size;
-#if defined(TSAN_ENABLED) && defined(__x86_64__)
+#if defined(TSAN_ENABLED)
     /* ThreadSanitizer compatibility: Allocate JIT code buffer at a fixed
      * address above the main memory region to avoid conflicts.
-     *
+     */
+#if defined(__x86_64__)
+    /* x86_64 memory layout:
      * Main memory: 0x7d0000000000 - 0x7d0100000000 (4GB for FULL4G)
      * JIT buffer:  0x7d1000000000 + size
      *
@@ -2348,7 +2357,32 @@ struct jit_state *jit_state_init(size_t size)
      */
     void *jit_addr = (void *) 0x7d1000000000UL;
     state->buf = mmap(jit_addr, size, PROT_READ | PROT_WRITE | PROT_EXEC,
-                      MAP_PRIVATE | MAP_ANONYMOUS | MAP_FIXED, -1, 0);
+                      MAP_PRIVATE | MAP_ANONYMOUS | MAP_FIXED
+#if defined(__APPLE__)
+                          | MAP_JIT
+#endif
+                      ,
+                      -1, 0);
+#elif defined(__aarch64__)
+    /* ARM64 memory layout (macOS/Apple Silicon):
+     * Main memory: 0x150000000000 - 0x150100000000 (4GB for FULL4G)
+     * JIT buffer:  0x151000000000 + size
+     *
+     * Apple Silicon requires MAP_JIT for executable memory. The fixed
+     * address is chosen to avoid TSAN's shadow memory and typical process
+     * allocations. Requires ASLR disabled via: setarch $(uname -m) -R
+     */
+    void *jit_addr = (void *) 0x151000000000UL;
+    state->buf = mmap(jit_addr, size, PROT_READ | PROT_WRITE | PROT_EXEC,
+                      MAP_PRIVATE | MAP_ANONYMOUS | MAP_FIXED
+#if defined(__APPLE__)
+                          | MAP_JIT
+#endif
+                      ,
+                      -1, 0);
+#else
+#error "TSAN is only supported on x86_64 and aarch64"
+#endif
     if (state->buf == MAP_FAILED) {
         free(state);
         return NULL;