/*
 * JIT compiler: RV64I basic-block -> x86-64 native code.
 *
 * Strategy:
 *   - Guest registers live in a u64[32] array pointed to by the first arg
 *     (rdi on entry).  Memory base is in rsi.  The third arg (rdx) is a
 *     pointer where we write the next guest PC on exit.
 *   - During block execution we pin:
 *       r12 = &regs[0]   (guest register file)
 *       r13 = memory base
 *       r14 = pc_out pointer
 *   - Guest registers are loaded/stored on demand from/to the array.
 *   - At block boundaries we store the next guest PC into *pc_out and
 *     return an exit-reason code in eax.
 *
 *  We emit raw x86-64 bytes into the code cache.
 */

#include <string.h>
#include <sys/mman.h>

#include "jit.h"
#include "riscv.h"
#include "types.h"

/* ---------- x86-64 code emitter helpers --------------------------------- */

/* A small code buffer that we fill up, then copy into the cache. */
struct emitter {
    u8  *buf;      /* Start of buffer. */
    u32  pos;      /* Current write position. */
    u32  capacity; /* Maximum bytes. */
    bool overflow; /* Set if we ran out of space. */
};

static inline void emit_u8(struct emitter *e, u8 b) {
    if (e->pos < e->capacity)
        e->buf[e->pos++] = b;
    else
        e->overflow = true;
}

static inline void emit_u32(struct emitter *e, u32 v) {
    emit_u8(e, (u8)(v));
    emit_u8(e, (u8)(v >> 8));
    emit_u8(e, (u8)(v >> 16));
    emit_u8(e, (u8)(v >> 24));
}

static inline void emit_u64(struct emitter *e, u64 v) {
    emit_u32(e, (u32)v);
    emit_u32(e, (u32)(v >> 32));
}

/* x86-64 register encoding (for ModR/M, SIB, REX). */
enum x86reg {
    X_RAX = 0,
    X_RCX = 1,
    X_RDX = 2,
    X_RBX = 3,
    X_RSP = 4,
    X_RBP = 5,
    X_RSI = 6,
    X_RDI = 7,
    X_R8  = 8,
    X_R9  = 9,
    X_R10 = 10,
    X_R11 = 11,
    X_R12 = 12,
    X_R13 = 13,
    X_R14 = 14,
    X_R15 = 15,
};

/* REX prefix byte. W=64-bit, R=reg extension, X=SIB index, B=rm/base. */
static inline u8 rex(bool w, bool r, bool x, bool b) {
    return (u8)(0x40 | (w ? 8 : 0) | (r ? 4 : 0) | (x ? 2 : 0) | (b ? 1 : 0));
}

/* ModR/M byte. */
static inline u8 modrm(u8 mod, u8 reg, u8 rm) {
    return (u8)((mod << 6) | ((reg & 7) << 3) | (rm & 7));
}

/* SIB byte. */
static inline u8 sib(u8 scale, u8 index, u8 base) {
    return (u8)((scale << 6) | ((index & 7) << 3) | (base & 7));
}

/* Pinned host registers. */
#define HREGS  X_R12 /* &regs[0]  */
#define HMEM   X_R13 /* memory    */
#define HPCOUT X_R14 /* pc_out    */

/* Scratch registers for codegen (caller-saved, not pinned). */
#define HTMP1 X_RAX
#define HTMP2 X_RCX

/* ---------- Common encoding helpers ------------------------------------ */

/* Emit REX.W prefix (64-bit operand size, no extended regs). */
static inline void emit_rexw(struct emitter *e) {
    emit_u8(e, 0x48);
}

/* Emit a 2-register ALU op: <op> rax, rcx (64-bit).
 * `opcode` is the x86 opcode byte (e.g. 0x01=add, 0x29=sub, etc). */
static void emit_alu_rax_rcx(struct emitter *e, u8 opcode) {
    emit_rexw(e);
    emit_u8(e, opcode);
    emit_u8(e, modrm(3, X_RCX, X_RAX));
}

/* Emit a shift: <shift> rax, cl (64-bit).
 * `ext` is the ModR/M extension (4=SHL, 5=SHR, 7=SAR). */
static void emit_shift_rax_cl(struct emitter *e, u8 ext) {
    emit_rexw(e);
    emit_u8(e, 0xD3);
    emit_u8(e, modrm(3, ext, X_RAX));
}

/* ---------- Load/store guest register from register file --------------- */

/* Emit mov with [r12 + disp] addressing (r12 needs SIB). */
static void emit_r12_disp(
    struct emitter *e, u8 opcode, enum x86reg reg, u32 off
) {
    emit_u8(e, rex(true, reg >= 8, false, true)); /* B=1 for r12 */
    emit_u8(e, opcode);
    if (off < 128) {
        emit_u8(e, modrm(1, reg, X_R12 & 7));
        emit_u8(e, sib(0, 4, X_R12 & 7));
        emit_u8(e, (u8)off);
    } else {
        emit_u8(e, modrm(2, reg, X_R12 & 7));
        emit_u8(e, sib(0, 4, X_R12 & 7));
        emit_u32(e, off);
    }
}

/* Load guest register `guest` into host `dst` from [r12 + guest*8]. */
static void emit_load_guest(struct emitter *e, enum x86reg dst, u32 guest) {
    if (guest == 0) {
        /* x0 is always zero -- xor dst, dst. */
        emit_u8(e, rex(true, dst >= 8, false, dst >= 8));
        emit_u8(e, 0x31);
        emit_u8(e, modrm(3, dst, dst));
        return;
    }
    emit_r12_disp(e, 0x8B, dst, guest * 8);
}

/* Store host `src` to guest register [r12 + guest*8]. */
static void emit_store_guest(struct emitter *e, u32 guest, enum x86reg src) {
    if (guest == 0)
        return;
    emit_r12_disp(e, 0x89, src, guest * 8);
}

/* mov reg, imm64 (REX.W + B8+rd) */
static void emit_mov_imm64(struct emitter *e, enum x86reg dst, u64 imm) {
    emit_u8(e, rex(true, false, false, dst >= 8));
    emit_u8(e, (u8)(0xB8 + (dst & 7)));
    emit_u64(e, imm);
}

/* mov reg, imm32 (sign-extended to 64 via mov r/m64, imm32) */
static void emit_mov_imm32_sx(struct emitter *e, enum x86reg dst, i32 imm) {
    emit_u8(e, rex(true, false, false, dst >= 8));
    emit_u8(e, 0xC7);
    emit_u8(e, modrm(3, 0, dst));
    emit_u32(e, (u32)imm);
}

/* Careful prologue using known encodings. */
static void emit_prologue(struct emitter *e) {
    emit_u8(e, 0x55); /* push rbp */
    emit_u8(e, 0x41);
    emit_u8(e, 0x54); /* push r12 */
    emit_u8(e, 0x41);
    emit_u8(e, 0x55); /* push r13 */
    emit_u8(e, 0x41);
    emit_u8(e, 0x56); /* push r14 */
    emit_u8(e, 0x49);
    emit_u8(e, 0x89);
    emit_u8(e, 0xFC); /* mov r12, rdi */
    emit_u8(e, 0x49);
    emit_u8(e, 0x89);
    emit_u8(e, 0xF5); /* mov r13, rsi */
    emit_u8(e, 0x49);
    emit_u8(e, 0x89);
    emit_u8(e, 0xD6); /* mov r14, rdx */
}

/* Emit function epilogue + ret with exit reason. */
static void emit_epilogue_with_exit(struct emitter *e, u32 exit_reason) {
    emit_u8(e, 0xB8);
    emit_u32(e, exit_reason); /* mov eax, reason */
    emit_u8(e, 0x41);
    emit_u8(e, 0x5E); /* pop r14 */
    emit_u8(e, 0x41);
    emit_u8(e, 0x5D); /* pop r13 */
    emit_u8(e, 0x41);
    emit_u8(e, 0x5C); /* pop r12 */
    emit_u8(e, 0x5D); /* pop rbp */
    emit_u8(e, 0xC3); /* ret */
}

/* Write next-PC to *r14, then epilogue+ret. */
static void emit_block_exit(
    struct emitter *e, enum jit_exit reason, u32 next_pc
) {
    /* mov dword [r14], next_pc */
    emit_u8(e, 0x41);
    emit_u8(e, 0xC7);
    emit_u8(e, modrm(0, 0, X_R14 & 7));
    emit_u32(e, next_pc);
    emit_epilogue_with_exit(e, (u32)reason);
}

/* ---------- Jump patching ---------------------------------------------- */

/* Emit jnz rel32 (0F 85), return position for patching. */
static u32 emit_jnz_placeholder(struct emitter *e) {
    u32 pos = e->pos;
    emit_u8(e, 0x0F);
    emit_u8(e, 0x85);
    emit_u32(e, 0);
    return pos;
}

/* Emit jne rel32 (0F 85), return position for patching. */
#define emit_jne_placeholder emit_jnz_placeholder

/* Emit jmp rel32 (E9), return position for patching. */
static u32 emit_jmp_placeholder(struct emitter *e) {
    u32 pos = e->pos;
    emit_u8(e, 0xE9);
    emit_u32(e, 0);
    return pos;
}

/* Patch a jcc rel32 (6-byte: 0F xx rel32) to jump to `target`. */
static void patch_jcc(struct emitter *e, u32 jcc_pos, u32 target) {
    u32 rel             = target - (jcc_pos + 6);
    e->buf[jcc_pos + 2] = (u8)rel;
    e->buf[jcc_pos + 3] = (u8)(rel >> 8);
    e->buf[jcc_pos + 4] = (u8)(rel >> 16);
    e->buf[jcc_pos + 5] = (u8)(rel >> 24);
}

/* Patch a jmp rel32 (5-byte: E9 rel32) to jump to `target`. */
static void patch_jmp(struct emitter *e, u32 jmp_pos, u32 target) {
    u32 rel             = target - (jmp_pos + 5);
    e->buf[jmp_pos + 1] = (u8)rel;
    e->buf[jmp_pos + 2] = (u8)(rel >> 8);
    e->buf[jmp_pos + 3] = (u8)(rel >> 16);
    e->buf[jmp_pos + 4] = (u8)(rel >> 24);
}

/* ---------- Per-instruction translation helpers ------------------------ */

/* Load rs1 into rax. */
static void emit_load_rs1(struct emitter *e, u32 rs1) {
    emit_load_guest(e, HTMP1, rs1);
}

/* Load rs2 into rcx. */
static void emit_load_rs2(struct emitter *e, u32 rs2) {
    emit_load_guest(e, HTMP2, rs2);
}

/* Store rax to rd. */
static void emit_store_rd(struct emitter *e, u32 rd) {
    emit_store_guest(e, rd, HTMP1);
}

/* add rax, imm32 (sign-extended) */
static void emit_add_rax_imm32(struct emitter *e, i32 imm) {
    if (imm == 0)
        return;
    emit_rexw(e);
    if (imm >= -128 && imm <= 127) {
        emit_u8(e, 0x83);
        emit_u8(e, modrm(3, 0, X_RAX));
        emit_u8(e, (u8)(i8)imm);
    } else {
        emit_u8(e, 0x05); /* add rax, imm32 (short form) */
        emit_u32(e, (u32)imm);
    }
}

/* imul rax, rcx */
static void emit_imul_rax_rcx(struct emitter *e) {
    emit_rexw(e);
    emit_u8(e, 0x0F);
    emit_u8(e, 0xAF);
    emit_u8(e, modrm(3, X_RAX, X_RCX));
}

/* cqo (sign-extend rax into rdx:rax) */
static void emit_cqo(struct emitter *e) {
    emit_rexw(e);
    emit_u8(e, 0x99);
}

/* idiv rcx (signed divide rdx:rax by rcx, 64-bit) */
static void emit_idiv_rcx(struct emitter *e) {
    emit_rexw(e);
    emit_u8(e, 0xF7);
    emit_u8(e, modrm(3, 7, X_RCX));
}

/* div rcx (unsigned divide rdx:rax by rcx, 64-bit) */
static void emit_div_rcx(struct emitter *e) {
    emit_rexw(e);
    emit_u8(e, 0xF7);
    emit_u8(e, modrm(3, 6, X_RCX));
}

/* xor rdx, rdx (zero rdx for unsigned division) */
static void emit_xor_rdx_rdx(struct emitter *e) {
    emit_rexw(e);
    emit_u8(e, 0x31);
    emit_u8(e, modrm(3, X_RDX, X_RDX));
}

/* movsxd rax, eax (sign-extend 32-bit result to 64-bit) */
static void emit_movsxd_rax_eax(struct emitter *e) {
    emit_rexw(e);
    emit_u8(e, 0x63);
    emit_u8(e, modrm(3, X_RAX, X_RAX));
}

/* test rcx, rcx (64-bit) */
static void emit_test_rcx(struct emitter *e) {
    emit_rexw(e);
    emit_u8(e, 0x85);
    emit_u8(e, modrm(3, X_RCX, X_RCX));
}

/* test ecx, ecx (32-bit) */
static void emit_test_ecx(struct emitter *e) {
    emit_u8(e, 0x85);
    emit_u8(e, modrm(3, X_RCX, X_RCX));
}

/* cmp rcx, -1 (64-bit) */
static void emit_cmp_rcx_neg1_64(struct emitter *e) {
    emit_rexw(e);
    emit_u8(e, 0x83);
    emit_u8(e, modrm(3, 7, X_RCX));
    emit_u8(e, 0xFF);
}

/* cmp ecx, -1 (32-bit) */
static void emit_cmp_ecx_neg1_32(struct emitter *e) {
    emit_u8(e, 0x83);
    emit_u8(e, modrm(3, 7, X_RCX));
    emit_u8(e, 0xFF);
}

/* mov rax, r11 (via: REX.WR mov rax, r11) */
static void emit_mov_rax_r11(struct emitter *e) {
    emit_u8(e, 0x4C);
    emit_u8(e, 0x89);
    emit_u8(e, modrm(3, X_R11 & 7, X_RAX));
}

/* mov rax, rdx (64-bit) */
static void emit_mov_rax_rdx(struct emitter *e) {
    emit_rexw(e);
    emit_u8(e, 0x89);
    emit_u8(e, modrm(3, X_RDX, X_RAX));
}

/* cmp rax, rcx (64-bit) */
static void emit_cmp_rax_rcx(struct emitter *e) {
    emit_rexw(e);
    emit_u8(e, 0x39);
    emit_u8(e, modrm(3, X_RCX, X_RAX));
}

/* setCC al + movzx rax, al (64-bit).
 * `cc` is the setcc secondary opcode (e.g. 0x9C=setl, 0x92=setb). */
static void emit_setcc_rax(struct emitter *e, u8 cc) {
    emit_u8(e, 0x0F);
    emit_u8(e, cc);
    emit_u8(e, modrm(3, 0, X_RAX));
    emit_rexw(e);
    emit_u8(e, 0x0F);
    emit_u8(e, 0xB6);
    emit_u8(e, modrm(3, X_RAX, X_RAX));
}

/* and ecx, imm8 (for masking shift amounts). */
static void emit_and_ecx_imm8(struct emitter *e, u8 mask) {
    emit_u8(e, 0x83);
    emit_u8(e, modrm(3, 4, X_RCX));
    emit_u8(e, mask);
}

/* ---------- Division helpers ------------------------------------------- */

/*
 * Shared skeleton for 64-bit div/rem with RISC-V corner cases:
 *   - divisor == 0: emit_zero_case (custom per variant)
 *   - signed overflow (INT_MIN / -1): emit_overflow_case
 *   - otherwise: normal division
 *
 * `is_signed`: whether to check for INT_MIN/-1 overflow.
 * `is_rem`:    whether to move rdx->rax after division (remainder result).
 * `save_dividend`: whether to save rax to r11 before testing (needed for
 *                  rem(x,0)=x semantics).
 */
static void emit_div64(struct emitter *e, bool is_signed, bool is_rem) {
    /* For rem: save dividend to r11 (needed if divisor==0). */
    if (is_rem) {
        /* mov r11, rax */
        emit_u8(e, rex(true, false, false, true));
        emit_u8(e, 0x89);
        emit_u8(e, modrm(3, X_RAX, X_R11 & 7));
    }

    emit_test_rcx(e);
    u32 jnz = emit_jnz_placeholder(e);

    /* Divisor == 0. */
    if (is_rem) {
        emit_mov_rax_r11(e); /* result = dividend */
    } else {
        emit_mov_imm32_sx(e, HTMP1, -1); /* result = -1 (all ones) */
    }
    u32 jmp_end1 = emit_jmp_placeholder(e);

    /* .nonzero: */
    u32 nonzero  = e->pos;
    u32 jne_safe = 0, jmp_end2 = 0;

    if (is_signed) {
        emit_cmp_rcx_neg1_64(e);
        jne_safe = emit_jne_placeholder(e);
        /* rcx == -1: overflow case. */
        if (is_rem) {
            /* result = 0 */
            emit_rexw(e);
            emit_u8(e, 0x31);
            emit_u8(e, modrm(3, X_RAX, X_RAX));
        }
        /* else: result = rax (already INT_MIN, which is correct) */
        jmp_end2 = emit_jmp_placeholder(e);
    }

    /* .safe: perform the actual division. */
    u32 safe = e->pos;
    if (is_signed) {
        emit_cqo(e);
        emit_idiv_rcx(e);
    } else {
        emit_xor_rdx_rdx(e);
        emit_div_rcx(e);
    }
    if (is_rem) {
        emit_mov_rax_rdx(e); /* remainder is in rdx */
    }

    /* .end: */
    u32 end = e->pos;
    patch_jcc(e, jnz, nonzero);
    patch_jmp(e, jmp_end1, end);
    if (is_signed) {
        patch_jcc(e, jne_safe, safe);
        patch_jmp(e, jmp_end2, end);
    }
}

/* Same pattern for 32-bit division (W-suffix instructions). */
static void emit_div32(struct emitter *e, bool is_signed, bool is_rem) {
    /* For rem: save dividend in edx. */
    if (is_rem) {
        /* mov edx, eax */
        emit_u8(e, 0x89);
        emit_u8(e, modrm(3, X_RAX, X_RDX));
    }

    emit_test_ecx(e);
    u32 jnz = emit_jnz_placeholder(e);

    /* Divisor == 0. */
    if (is_rem) {
        /* mov eax, edx (result = dividend) */
        emit_u8(e, 0x89);
        emit_u8(e, modrm(3, X_RDX, X_RAX));
    } else {
        /* mov eax, -1 */
        emit_u8(e, 0xB8);
        emit_u32(e, 0xFFFFFFFF);
    }
    u32 jmp_end1 = emit_jmp_placeholder(e);

    /* .nonzero: */
    u32 nonzero  = e->pos;
    u32 jne_safe = 0, jmp_end2 = 0;

    if (is_signed) {
        emit_cmp_ecx_neg1_32(e);
        jne_safe = emit_jne_placeholder(e);
        if (is_rem) {
            /* result = 0 */
            emit_u8(e, 0x31);
            emit_u8(e, modrm(3, X_RAX, X_RAX));
        }
        /* else: result = eax (INT_MIN stays INT_MIN) */
        jmp_end2 = emit_jmp_placeholder(e);
    }

    /* .safe: */
    u32 safe = e->pos;
    if (is_rem) {
        /* Restore dividend to eax from edx. */
        emit_u8(e, 0x89);
        emit_u8(e, modrm(3, X_RDX, X_RAX));
    }
    if (is_signed) {
        emit_u8(e, 0x99); /* cdq */
        emit_u8(e, 0xF7);
        emit_u8(e, modrm(3, 7, X_RCX)); /* idiv ecx */
    } else {
        emit_u8(e, 0x31);
        emit_u8(e, modrm(3, X_RDX, X_RDX)); /* xor edx,edx */
        emit_u8(e, 0xF7);
        emit_u8(e, modrm(3, 6, X_RCX)); /* div ecx */
    }
    if (is_rem) {
        emit_u8(e, 0x89);
        emit_u8(e, modrm(3, X_RDX, X_RAX)); /* mov eax,edx */
    }

    /* .end: */
    u32 end = e->pos;
    patch_jcc(e, jnz, nonzero);
    patch_jmp(e, jmp_end1, end);
    if (is_signed) {
        patch_jcc(e, jne_safe, safe);
        patch_jmp(e, jmp_end2, end);
    }
}

/* ---------- Memory access helpers -------------------------------------- */

/* add rax, r13 (compute host address from guest address). */
static void emit_add_rax_r13(struct emitter *e) {
    emit_u8(e, 0x4C);
    emit_u8(e, 0x01);
    emit_u8(e, 0xE8);
}

/* Load from [r13+rax] into rax, with sign/zero extension. */
static void emit_load_mem_i8(struct emitter *e) {
    emit_add_rax_r13(e);
    emit_rexw(e);
    emit_u8(e, 0x0F);
    emit_u8(e, 0xBE);
    emit_u8(e, modrm(0, X_RAX, X_RAX));
}

static void emit_load_mem_u8(struct emitter *e) {
    emit_add_rax_r13(e);
    emit_u8(e, 0x0F);
    emit_u8(e, 0xB6);
    emit_u8(e, modrm(0, X_RAX, X_RAX));
}

static void emit_load_mem_i16(struct emitter *e) {
    emit_add_rax_r13(e);
    emit_rexw(e);
    emit_u8(e, 0x0F);
    emit_u8(e, 0xBF);
    emit_u8(e, modrm(0, X_RAX, X_RAX));
}

static void emit_load_mem_u16(struct emitter *e) {
    emit_add_rax_r13(e);
    emit_u8(e, 0x0F);
    emit_u8(e, 0xB7);
    emit_u8(e, modrm(0, X_RAX, X_RAX));
}

static void emit_load_mem_i32(struct emitter *e) {
    emit_add_rax_r13(e);
    emit_rexw(e);
    emit_u8(e, 0x63);
    emit_u8(e, modrm(0, X_RAX, X_RAX));
}

static void emit_load_mem_u32(struct emitter *e) {
    emit_add_rax_r13(e);
    emit_u8(e, 0x8B);
    emit_u8(e, modrm(0, X_RAX, X_RAX));
}

static void emit_load_mem_u64(struct emitter *e) {
    emit_add_rax_r13(e);
    emit_rexw(e);
    emit_u8(e, 0x8B);
    emit_u8(e, modrm(0, X_RAX, X_RAX));
}

/* Store from rcx to [r13+rax]. */
static void emit_store_mem_u8(struct emitter *e) {
    emit_add_rax_r13(e);
    emit_u8(e, 0x88);
    emit_u8(e, modrm(0, X_RCX, X_RAX));
}

static void emit_store_mem_u16(struct emitter *e) {
    emit_add_rax_r13(e);
    emit_u8(e, 0x66);
    emit_u8(e, 0x89);
    emit_u8(e, modrm(0, X_RCX, X_RAX));
}

static void emit_store_mem_u32(struct emitter *e) {
    emit_add_rax_r13(e);
    emit_u8(e, 0x89);
    emit_u8(e, modrm(0, X_RCX, X_RAX));
}

static void emit_store_mem_u64(struct emitter *e) {
    emit_add_rax_r13(e);
    emit_rexw(e);
    emit_u8(e, 0x89);
    emit_u8(e, modrm(0, X_RCX, X_RAX));
}

/* ---------- Translate one RV64I instruction ----------------------------- */

/*
 * Returns: true if the instruction ends the basic block (branch/jump/ecall),
 *          false if execution should continue to the next instruction.
 */
static bool translate_insn(struct emitter *e, instr_t ins, u32 pc) {
    u32 opcode  = ins.r.opcode;
    u32 pc_next = pc + INSTR_SIZE;

    switch (opcode) {

    case OP_LUI:
        if (ins.u.rd != 0) {
            emit_mov_imm32_sx(e, HTMP1, (i32)(ins.u.imm_31_12 << 12));
            emit_store_rd(e, ins.u.rd);
        }
        return false;

    case OP_AUIPC:
        if (ins.u.rd != 0) {
            i64 result = (i64)pc + (i64)(i32)(ins.u.imm_31_12 << 12);
            emit_mov_imm64(e, HTMP1, (u64)result);
            emit_store_rd(e, ins.u.rd);
        }
        return false;

    case OP_JAL: {
        u32 target = pc + (u32)get_j_imm(ins);
        if (ins.j.rd != 0) {
            emit_mov_imm64(e, HTMP1, (u64)pc_next);
            emit_store_guest(e, ins.j.rd, HTMP1);
        }
        emit_block_exit(e, JIT_EXIT_BRANCH, target);
        return true;
    }

    case OP_JALR: {
        i32 imm = get_i_imm(ins);

        /* Special-case RET: jalr x0, ra, 0. */
        if (ins.i.rd == 0 && ins.i.rs1 == RA && imm == 0) {
            emit_load_guest(e, HTMP1, RA);

            /* test rax, rax */
            emit_rexw(e);
            emit_u8(e, 0x85);
            emit_u8(e, modrm(3, X_RAX, X_RAX));

            u32 jnz_pos = emit_jnz_placeholder(e);

            /* RA == 0: program exit. */
            emit_block_exit(e, JIT_EXIT_RET, 0);

            /* RA != 0: compute target = RA & ~1. */
            patch_jcc(e, jnz_pos, e->pos);
            emit_rexw(e);
            emit_u8(e, 0x83);
            emit_u8(e, modrm(3, 4, X_RAX));
            emit_u8(e, 0xFE); /* and rax, ~1 */

            /* mov dword [r14], eax -- write PC */
            emit_u8(e, 0x41);
            emit_u8(e, 0x89);
            emit_u8(e, modrm(0, X_RAX, X_R14 & 7));

            emit_epilogue_with_exit(e, (u32)JIT_EXIT_BRANCH);
            return true;
        }

        /* General JALR: target = (rs1 + imm) & ~1. */
        emit_load_rs1(e, ins.i.rs1);
        emit_add_rax_imm32(e, imm);
        emit_rexw(e);
        emit_u8(e, 0x83);
        emit_u8(e, modrm(3, 4, X_RAX));
        emit_u8(e, 0xFE); /* and rax, ~1 */

        /* mov rcx, rax (save target) */
        emit_rexw(e);
        emit_u8(e, 0x89);
        emit_u8(e, modrm(3, X_RAX, X_RCX));

        if (ins.i.rd != 0) {
            emit_mov_imm64(e, HTMP1, (u64)pc_next);
            emit_store_guest(e, ins.i.rd, HTMP1);
        }

        /* mov dword [r14], ecx */
        emit_u8(e, 0x41);
        emit_u8(e, 0x89);
        emit_u8(e, modrm(0, X_RCX, X_R14 & 7));

        emit_epilogue_with_exit(e, (u32)JIT_EXIT_BRANCH);
        return true;
    }

    case OP_BRANCH: {
        u32 target = pc + (u32)get_b_imm(ins);

        emit_load_rs1(e, ins.b.rs1);
        emit_load_rs2(e, ins.b.rs2);
        emit_cmp_rax_rcx(e);

        /* jCC to .taken */
        u8 cc;
        switch (ins.b.funct3) {
        case 0x0:
            cc = 0x84;
            break; /* beq -> je */
        case 0x1:
            cc = 0x85;
            break; /* bne -> jne */
        case 0x4:
            cc = 0x8C;
            break; /* blt -> jl */
        case 0x5:
            cc = 0x8D;
            break; /* bge -> jge */
        case 0x6:
            cc = 0x82;
            break; /* bltu -> jb */
        case 0x7:
            cc = 0x83;
            break; /* bgeu -> jae */
        default:
            cc = 0x84;
            break;
        }

        u32 jcc_pos = e->pos;
        emit_u8(e, 0x0F);
        emit_u8(e, cc);
        emit_u32(e, 0);

        /* Not taken: fall through. */
        emit_block_exit(e, JIT_EXIT_BRANCH, pc_next);

        /* .taken: */
        patch_jcc(e, jcc_pos, e->pos);
        emit_block_exit(e, JIT_EXIT_BRANCH, target);
        return true;
    }

    case OP_LOAD: {
        i32 imm = get_i_imm(ins);
        if (ins.i.rd == 0)
            return false;

        emit_load_rs1(e, ins.i.rs1);
        emit_add_rax_imm32(e, imm);

        switch (ins.i.funct3) {
        case 0x0:
            emit_load_mem_i8(e);
            break; /* lb */
        case 0x1:
            emit_load_mem_i16(e);
            break; /* lh */
        case 0x2:
            emit_load_mem_i32(e);
            break; /* lw */
        case 0x3:
            emit_load_mem_u64(e);
            break; /* ld */
        case 0x4:
            emit_load_mem_u8(e);
            break; /* lbu */
        case 0x5:
            emit_load_mem_u16(e);
            break; /* lhu */
        case 0x6:
            emit_load_mem_u32(e);
            break; /* lwu */
        default:
            emit_block_exit(e, JIT_EXIT_FAULT, pc);
            return true;
        }
        emit_store_rd(e, ins.i.rd);
        return false;
    }

    case OP_STORE: {
        i32 imm = get_s_imm(ins);
        emit_load_guest(e, HTMP1, ins.s.rs1);
        emit_add_rax_imm32(e, imm);
        emit_load_guest(e, HTMP2, ins.s.rs2);

        switch (ins.s.funct3) {
        case 0x0:
            emit_store_mem_u8(e);
            break; /* sb */
        case 0x1:
            emit_store_mem_u16(e);
            break; /* sh */
        case 0x2:
            emit_store_mem_u32(e);
            break; /* sw */
        case 0x3:
            emit_store_mem_u64(e);
            break; /* sd */
        default:
            emit_block_exit(e, JIT_EXIT_FAULT, pc);
            return true;
        }
        return false;
    }

    case OP_IMM: {
        i32 imm = get_i_imm(ins);
        if (ins.i.rd == 0)
            return false;

        emit_load_rs1(e, ins.i.rs1);

        switch (ins.i.funct3) {
        case 0x0: /* addi */
            emit_add_rax_imm32(e, imm);
            break;
        case 0x1: /* slli */
            emit_rexw(e);
            emit_u8(e, 0xC1);
            emit_u8(e, modrm(3, 4, X_RAX));
            emit_u8(e, (u8)(imm & 0x3F));
            break;
        case 0x2: /* slti */
            emit_mov_imm32_sx(e, HTMP2, imm);
            emit_cmp_rax_rcx(e);
            emit_setcc_rax(e, 0x9C); /* setl */
            break;
        case 0x3: /* sltiu */
            emit_mov_imm32_sx(e, HTMP2, imm);
            emit_cmp_rax_rcx(e);
            emit_setcc_rax(e, 0x92); /* setb */
            break;
        case 0x4: /* xori */
            emit_mov_imm32_sx(e, HTMP2, imm);
            emit_alu_rax_rcx(e, 0x31); /* xor */
            break;
        case 0x5: /* srli/srai */
            emit_rexw(e);
            emit_u8(e, 0xC1);
            emit_u8(e, modrm(3, (imm & 0x400) ? 7 : 5, X_RAX));
            emit_u8(e, (u8)(imm & 0x3F));
            break;
        case 0x6: /* ori */
            emit_mov_imm32_sx(e, HTMP2, imm);
            emit_alu_rax_rcx(e, 0x09); /* or */
            break;
        case 0x7: /* andi */
            emit_mov_imm32_sx(e, HTMP2, imm);
            emit_alu_rax_rcx(e, 0x21); /* and */
            break;
        }
        emit_store_rd(e, ins.i.rd);
        return false;
    }

    case OP_IMM_32: {
        i32 imm = get_i_imm(ins);
        if (ins.i.rd == 0)
            return false;

        emit_load_rs1(e, ins.i.rs1);

        switch (ins.i.funct3) {
        case 0x0: /* addiw */
            emit_add_rax_imm32(e, imm);
            break;
        case 0x1: /* slliw */
            emit_u8(e, 0xC1);
            emit_u8(e, modrm(3, 4, X_RAX));
            emit_u8(e, (u8)(imm & 0x1F));
            break;
        case 0x5: /* srliw/sraiw */
            emit_u8(e, 0xC1);
            emit_u8(e, modrm(3, (imm & 0x400) ? 7 : 5, X_RAX));
            emit_u8(e, (u8)(imm & 0x1F));
            break;
        }
        emit_movsxd_rax_eax(e);
        emit_store_rd(e, ins.i.rd);
        return false;
    }

    case OP_OP: {
        if (ins.r.rd == 0)
            return false;
        emit_load_rs1(e, ins.r.rs1);
        emit_load_rs2(e, ins.r.rs2);

        switch (ins.r.funct7) {
        case FUNCT7_NORMAL:
            switch (ins.r.funct3) {
            case 0x0:
                emit_alu_rax_rcx(e, 0x01);
                break; /* add */
            case 0x1:  /* sll */
                emit_and_ecx_imm8(e, 0x3F);
                emit_shift_rax_cl(e, 4);
                break;
            case 0x2: /* slt */
                emit_cmp_rax_rcx(e);
                emit_setcc_rax(e, 0x9C); /* setl */
                break;
            case 0x3: /* sltu */
                emit_cmp_rax_rcx(e);
                emit_setcc_rax(e, 0x92); /* setb */
                break;
            case 0x4:
                emit_alu_rax_rcx(e, 0x31);
                break; /* xor */
            case 0x5:  /* srl */
                emit_and_ecx_imm8(e, 0x3F);
                emit_shift_rax_cl(e, 5);
                break;
            case 0x6:
                emit_alu_rax_rcx(e, 0x09);
                break; /* or */
            case 0x7:
                emit_alu_rax_rcx(e, 0x21);
                break; /* and */
            }
            break;

        case FUNCT7_SUB:
            switch (ins.r.funct3) {
            case 0x0:
                emit_alu_rax_rcx(e, 0x29);
                break; /* sub */
            case 0x5:  /* sra */
                emit_and_ecx_imm8(e, 0x3F);
                emit_shift_rax_cl(e, 7);
                break;
            }
            break;

        case FUNCT7_MUL:
            switch (ins.r.funct3) {
            case 0x0:
                emit_imul_rax_rcx(e);
                break; /* mul */
            case 0x4:
                emit_div64(e, true, false);
                break; /* div */
            case 0x5:
                emit_div64(e, false, false);
                break; /* divu */
            case 0x6:
                emit_div64(e, true, true);
                break; /* rem */
            case 0x7:
                emit_div64(e, false, true);
                break; /* remu */
            }
            break;
        }
        emit_store_rd(e, ins.r.rd);
        return false;
    }

    case OP_OP_32: {
        if (ins.r.rd == 0)
            return false;
        emit_load_rs1(e, ins.r.rs1);
        emit_load_rs2(e, ins.r.rs2);

        switch (ins.r.funct7) {
        case FUNCT7_NORMAL:
            switch (ins.r.funct3) {
            case 0x0: /* addw */
                emit_u8(e, 0x01);
                emit_u8(e, modrm(3, X_RCX, X_RAX));
                break;
            case 0x1: /* sllw */
                emit_and_ecx_imm8(e, 0x1F);
                emit_u8(e, 0xD3);
                emit_u8(e, modrm(3, 4, X_RAX));
                break;
            case 0x5: /* srlw */
                emit_and_ecx_imm8(e, 0x1F);
                emit_u8(e, 0xD3);
                emit_u8(e, modrm(3, 5, X_RAX));
                break;
            }
            break;

        case FUNCT7_SUB:
            switch (ins.r.funct3) {
            case 0x0: /* subw */
                emit_u8(e, 0x29);
                emit_u8(e, modrm(3, X_RCX, X_RAX));
                break;
            case 0x5: /* sraw */
                emit_and_ecx_imm8(e, 0x1F);
                emit_u8(e, 0xD3);
                emit_u8(e, modrm(3, 7, X_RAX));
                break;
            }
            break;

        case FUNCT7_MUL:
            switch (ins.r.funct3) {
            case 0x0: /* mulw */
                emit_u8(e, 0x0F);
                emit_u8(e, 0xAF);
                emit_u8(e, modrm(3, X_RAX, X_RCX));
                break;
            case 0x4:
                emit_div32(e, true, false);
                break; /* divw */
            case 0x5:
                emit_div32(e, false, false);
                break; /* divuw */
            case 0x6:
                emit_div32(e, true, true);
                break; /* remw */
            case 0x7:
                emit_div32(e, false, true);
                break; /* remuw */
            }
            break;
        }
        emit_movsxd_rax_eax(e);
        emit_store_rd(e, ins.r.rd);
        return false;
    }

    case OP_SYSTEM: {
        u32 funct12 = ins.i.imm_11_0;
        if (funct12 == 0) {
            emit_block_exit(e, JIT_EXIT_ECALL, pc);
        } else if (funct12 == 1) {
            emit_block_exit(e, JIT_EXIT_EBREAK, pc);
        } else {
            emit_block_exit(e, JIT_EXIT_FAULT, pc);
        }
        return true;
    }

    case OP_FENCE:
        return false;

    default:
        emit_block_exit(e, JIT_EXIT_FAULT, pc);
        return true;
    }
}

/* ---------- Block compiler --------------------------------------------- */

static struct jit_block *jit_alloc_block(struct jit_state *jit) {
    if (jit->block_count >= JIT_MAX_BLOCKS)
        return NULL;
    return &jit->blocks[jit->block_count++];
}

static void jit_insert_block(struct jit_state *jit, struct jit_block *block) {
    u32 h              = (block->guest_pc >> 2) & (JIT_BLOCK_HASH_SIZE - 1);
    block->hash_next   = jit->block_hash[h];
    jit->block_hash[h] = block;
}

struct jit_block *jit_compile_block(
    struct jit_state *jit, u32 guest_pc, u8 *mem, u32 prog_base, u32 prog_bytes
) {

#define JIT_EMIT_BUF_SIZE (JIT_MAX_BLOCK_INSNS * JIT_MAX_INSN_BYTES + 256)
    static u8 emit_buf[JIT_EMIT_BUF_SIZE];

    struct emitter em = {
        .buf      = emit_buf,
        .pos      = 0,
        .capacity = JIT_EMIT_BUF_SIZE,
        .overflow = false,
    };

    emit_prologue(&em);

    u32 pc          = guest_pc;
    u32 insn_count  = 0;
    u32 prog_end_pc = prog_base + prog_bytes;

    while (insn_count < JIT_MAX_BLOCK_INSNS) {
        if (pc < prog_base || pc >= prog_end_pc) {
            emit_block_exit(&em, JIT_EXIT_FAULT, pc);
            break;
        }
        instr_t ins;
        memcpy(&ins, &mem[pc], sizeof(instr_t));
        insn_count++;

        bool ends_block  = translate_insn(&em, ins, pc);
        pc              += INSTR_SIZE;

        if (ends_block)
            break;

        if (em.pos + JIT_MAX_INSN_BYTES + 64 > em.capacity) {
            emit_block_exit(&em, JIT_EXIT_BRANCH, pc);
            break;
        }
    }

    if (insn_count >= JIT_MAX_BLOCK_INSNS) {
        emit_block_exit(&em, JIT_EXIT_BRANCH, pc);
    }

    if (em.overflow)
        return NULL;

    u32 code_size = em.pos;
    if (jit->code_cache_used + code_size > JIT_CODE_CACHE_SIZE)
        return NULL;

    u8 *dest = jit->code_cache + jit->code_cache_used;
    memcpy(dest, emit_buf, code_size);
    jit->code_cache_used += code_size;

    struct jit_block *block = jit_alloc_block(jit);
    if (!block)
        return NULL;

    block->guest_pc     = guest_pc;
    block->guest_end_pc = pc;
    block->insn_count   = insn_count;
    block->code         = dest;
    block->code_size    = code_size;
    block->hash_next    = NULL;

    jit_insert_block(jit, block);
    jit->blocks_compiled++;

    return block;
}

/* ---------- Init / Destroy / Flush ------------------------------------ */

bool jit_init(struct jit_state *jit) {
    memset(jit, 0, sizeof(*jit));

#if defined(__x86_64__) || defined(_M_X64)
    jit->code_cache = mmap(
        NULL,
        JIT_CODE_CACHE_SIZE,
        PROT_READ | PROT_WRITE | PROT_EXEC,
        MAP_PRIVATE | MAP_ANONYMOUS,
        -1,
        0
    );
    if (jit->code_cache == MAP_FAILED) {
        jit->code_cache = NULL;
        jit->available  = false;
        return false;
    }
    jit->available = true;
    return true;
#else
    jit->available = false;
    return false;
#endif
}

void jit_destroy(struct jit_state *jit) {
    if (jit->code_cache) {
        munmap(jit->code_cache, JIT_CODE_CACHE_SIZE);
        jit->code_cache = NULL;
    }
    jit->block_count = 0;
}

void jit_flush(struct jit_state *jit) {
    memset(jit->block_hash, 0, sizeof(jit->block_hash));
    jit->block_count     = 0;
    jit->code_cache_used = 0;
}