/*
 * jit.c — x86_64 native code JIT compiler for lumbda
 *
 * Emits real machine code bytes into mmap'd executable memory.
 * Handles: integer arithmetic (+, -, *), comparisons (=, <, >, <=, >=),
 * if/cond branching, recursive function calls (System V AMD64 ABI),
 * tail call optimization (jmp instead of call).
 * Falls back to the C interpreter for anything it can't compile.
 *
 * Code outlasts authors.
 */
#include "lumbda.h"
#include "jit.h"
#include <sys/mman.h>

bool g_jit_enabled = false;

/* ═══════════════════════════════════════════════════════════════════════════
 * x86_64 encoding constants
 * ═══════════════════════════════════════════════════════════════════════════ */

#define REX_W   0x48
#define REX_WB  0x49

#define RAX 0
#define RCX 1
#define RDX 2
#define RBX 3
#define RSP 4
#define RBP 5
#define RSI 6
#define RDI 7
#define R8  0
#define R9  1
#define R10 2
#define R11 3
#define R12 4
#define R13 5
#define R14 6
#define R15 7

#define MODRM(mod, reg, rm) (((mod) << 6) | ((reg) << 3) | (rm))

/* System V AMD64 argument registers */
static const int ARG_REGS[]    = { RDI, RSI, RDX, RCX, R8, R9 };
static const bool ARG_REG_EXT[] = { false, false, false, false, true, true };

/* NaN-box prefix for integers: QNAN | (TAG_INT << 48) */
#define INT_BOX_PREFIX (QNAN | (TAG_INT << TAG_SHIFT))

typedef enum { CMP_EQ, CMP_NE, CMP_LT, CMP_LE, CMP_GT, CMP_GE } CmpKind;

/* ═══════════════════════════════════════════════════════════════════════════
 * JIT context
 * ═══════════════════════════════════════════════════════════════════════════ */

/* Maximum local variables in let/let* bindings */
#define MAX_JIT_LOCALS 16

typedef struct {
    Value sym;
    int   slot;    /* stack slot index (0-based from rbp) */
} JitLocal;

typedef struct {
    uint8_t *buf;
    size_t   len;
    size_t   cap;
    uint8_t *entry;         /* start of mmap'd block = function entry */
    Value   *params;
    int      nparams;
    const char *name;
    Value    self_sym;      /* interned name for recursive calls */
    int      body_start;    /* offset of body for TCO jumps */
    int      epilogue_jumps[512];
    int      n_epilogue_jumps;
    int      tco_jumps[512];
    int      n_tco_jumps;
    int      extra_stack;   /* bytes of extra stack allocated */
    /* Local variable tracking for let/let* */
    JitLocal locals[MAX_JIT_LOCALS];
    int      nlocals;
    int      local_base;    /* rbp offset base for locals */
    /* Named-let loop support */
    Value    loop_sym;      /* symbol of the named-let loop name */
    int      loop_start;    /* code offset for loop restart */
    int      loop_nparams;  /* number of loop variables */
    Value   *loop_params;   /* loop variable symbols */
    int      loop_slots[MAX_JIT_LOCALS]; /* stack slots for loop vars */
} JitCtx;

static void jctx_init(JitCtx *j, size_t cap) {
    j->buf = (uint8_t *)mmap(NULL, cap,
                              PROT_READ | PROT_WRITE | PROT_EXEC,
                              MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
    if (j->buf == MAP_FAILED) { j->buf = NULL; j->cap = 0; return; }
    j->len = 0;
    j->cap = cap;
    j->entry = j->buf;
    j->n_epilogue_jumps = 0;
    j->n_tco_jumps = 0;
    j->body_start = 0;
    j->extra_stack = 0;
    j->nlocals = 0;
    j->local_base = 0;
    j->loop_sym = VAL_NIL;
    j->loop_start = 0;
    j->loop_nparams = 0;
    j->loop_params = NULL;
}

/* ═══════════════════════════════════════════════════════════════════════════
 * Low-level byte emitters
 * ═══════════════════════════════════════════════════════════════════════════ */

static inline void eb(JitCtx *j, uint8_t b) {
    if (j->len < j->cap) j->buf[j->len++] = b;
}

static inline void ei32(JitCtx *j, int32_t v) {
    eb(j, v & 0xFF); eb(j, (v>>8)&0xFF); eb(j, (v>>16)&0xFF); eb(j, (v>>24)&0xFF);
}

static inline void ei64(JitCtx *j, int64_t v) {
    for (int i = 0; i < 8; i++) eb(j, (v >> (i*8)) & 0xFF);
}

/* ═══════════════════════════════════════════════════════════════════════════
 * x86_64 instruction emitters
 * ═══════════════════════════════════════════════════════════════════════════ */

static void emit_mov_imm64(JitCtx *j, int reg, bool ext, uint64_t imm) {
    eb(j, ext ? REX_WB : REX_W);
    eb(j, 0xB8 + reg);
    ei64(j, (int64_t)imm);
}

static void emit_mov_rr(JitCtx *j, int dst, bool de, int src, bool se) {
    uint8_t rex = REX_W;
    if (se) rex |= 0x04; if (de) rex |= 0x01;
    eb(j, rex); eb(j, 0x89); eb(j, MODRM(3, src, dst));
}

static void emit_push(JitCtx *j, int reg, bool ext) {
    if (ext) eb(j, 0x41); eb(j, 0x50 + reg);
}

static void emit_pop(JitCtx *j, int reg, bool ext) {
    if (ext) eb(j, 0x41); eb(j, 0x58 + reg);
}

static void emit_add_rr(JitCtx *j, int dst, bool de, int src, bool se) {
    uint8_t rex = REX_W; if (se) rex|=4; if (de) rex|=1;
    eb(j, rex); eb(j, 0x01); eb(j, MODRM(3, src, dst));
}

static void emit_sub_rr(JitCtx *j, int dst, bool de, int src, bool se) {
    uint8_t rex = REX_W; if (se) rex|=4; if (de) rex|=1;
    eb(j, rex); eb(j, 0x29); eb(j, MODRM(3, src, dst));
}

static void emit_imul_rr(JitCtx *j, int dst, bool de, int src, bool se) {
    uint8_t rex = REX_W; if (de) rex|=4; if (se) rex|=1;
    eb(j, rex); eb(j, 0x0F); eb(j, 0xAF); eb(j, MODRM(3, dst, src));
}

static void emit_cmp_rr(JitCtx *j, int rm, bool rme, int reg, bool rege) {
    uint8_t rex = REX_W; if (rege) rex|=4; if (rme) rex|=1;
    eb(j, rex); eb(j, 0x39); eb(j, MODRM(3, reg, rm));
}

static void emit_neg(JitCtx *j, int reg, bool ext) {
    eb(j, ext ? REX_WB : REX_W);
    eb(j, 0xF7); eb(j, MODRM(3, 3, reg));
}

static void emit_sub_imm(JitCtx *j, int reg, bool ext, int32_t imm) {
    eb(j, ext ? REX_WB : REX_W);
    if (imm >= -128 && imm <= 127) {
        eb(j, 0x83); eb(j, MODRM(3, 5, reg)); eb(j, (uint8_t)(int8_t)imm);
    } else {
        eb(j, 0x81); eb(j, MODRM(3, 5, reg)); ei32(j, imm);
    }
}

static void emit_add_imm(JitCtx *j, int reg, bool ext, int32_t imm) {
    eb(j, ext ? REX_WB : REX_W);
    if (imm >= -128 && imm <= 127) {
        eb(j, 0x83); eb(j, MODRM(3, 0, reg)); eb(j, (uint8_t)(int8_t)imm);
    } else {
        eb(j, 0x81); eb(j, MODRM(3, 0, reg)); ei32(j, imm);
    }
}

/* ── Conditional jumps (rel32, return offset position for patching) ── */

static int emit_je(JitCtx *j)  { eb(j,0x0F); eb(j,0x84); int p=(int)j->len; ei32(j,0); return p; }
static int emit_jne(JitCtx *j) { eb(j,0x0F); eb(j,0x85); int p=(int)j->len; ei32(j,0); return p; }
static int emit_jl(JitCtx *j)  { eb(j,0x0F); eb(j,0x8C); int p=(int)j->len; ei32(j,0); return p; }
static int emit_jle(JitCtx *j) { eb(j,0x0F); eb(j,0x8E); int p=(int)j->len; ei32(j,0); return p; }
static int emit_jg(JitCtx *j)  { eb(j,0x0F); eb(j,0x8F); int p=(int)j->len; ei32(j,0); return p; }
static int emit_jge(JitCtx *j) { eb(j,0x0F); eb(j,0x8D); int p=(int)j->len; ei32(j,0); return p; }
static int emit_jmp(JitCtx *j) { eb(j,0xE9); int p=(int)j->len; ei32(j,0); return p; }

static void patch_jump(JitCtx *j, int pos) {
    int32_t rel = (int32_t)((int)j->len - (pos + 4));
    j->buf[pos]   = rel & 0xFF;
    j->buf[pos+1] = (rel >> 8) & 0xFF;
    j->buf[pos+2] = (rel >> 16) & 0xFF;
    j->buf[pos+3] = (rel >> 24) & 0xFF;
}

/* ── Epilogue jump (replaces bare ret) ── */
static void emit_return(JitCtx *j) {
    int pos = emit_jmp(j);
    if (j->n_epilogue_jumps < 512)
        j->epilogue_jumps[j->n_epilogue_jumps++] = pos;
}

/* ── TCO jump back to body start ── */
static void emit_tco_jump(JitCtx *j) {
    int pos = emit_jmp(j);
    if (j->n_tco_jumps < 512)
        j->tco_jumps[j->n_tco_jumps++] = pos;
}

/* ═══════════════════════════════════════════════════════════════════════════
 * NaN-boxing: unbox/box integers in registers
 * ═══════════════════════════════════════════════════════════════════════════ */

/* Unbox: extract lower 48 bits and sign-extend via shl 16; sar 16 */
static void emit_unbox_int(JitCtx *j, int reg, bool ext) {
    eb(j, ext ? REX_WB : REX_W);
    eb(j, 0xC1); eb(j, MODRM(3, 4, reg)); eb(j, 16); /* shl reg, 16 */
    eb(j, ext ? REX_WB : REX_W);
    eb(j, 0xC1); eb(j, MODRM(3, 7, reg)); eb(j, 16); /* sar reg, 16 */
}

/* Box: raw int64 → NaN-boxed integer. Uses r10, r11 as scratch. */
static void emit_box_int(JitCtx *j, int reg, bool ext) {
    emit_mov_imm64(j, R11, true, INT_BOX_PREFIX);
    emit_mov_imm64(j, R10, true, PAYLOAD_MASK);
    /* and reg, r10 */
    {
        uint8_t rex = REX_W;
        if (ext) rex |= 0x01;
        rex |= 0x04; /* REX.R for r10 */
        eb(j, rex); eb(j, 0x21); eb(j, MODRM(3, R10, reg));
    }
    /* or reg, r11 */
    {
        uint8_t rex = REX_W;
        if (ext) rex |= 0x01;
        rex |= 0x04;
        eb(j, rex); eb(j, 0x09); eb(j, MODRM(3, R11, reg));
    }
}

/* ═══════════════════════════════════════════════════════════════════════════
 * AST analysis — determine if a procedure is JIT-compilable
 * ═══════════════════════════════════════════════════════════════════════════ */

/* Extended can_jit_expr: checks with extra local variable names from let/let* */
static bool can_jit_expr_ext(Value expr, Value *params, int nparams,
                             Value self_sym, Value *locals, int nlocals);

/* Check if a symbol is in params or locals */
static bool is_known_var(Value sym, Value *params, int nparams,
                         Value *locals, int nlocals) {
    for (int i = 0; i < nparams; i++)
        if (params[i] == sym) return true;
    for (int i = 0; i < nlocals; i++)
        if (locals[i] == sym) return true;
    return false;
}

static bool can_jit_expr_ext(Value expr, Value *params, int nparams,
                             Value self_sym, Value *locals, int nlocals) {
    if (IS_INT(expr)) return true;

    /* Boolean literals */
    if (expr == VAL_TRUE || expr == VAL_FALSE) return true;

    /* NIL literal */
    if (IS_NIL(expr)) return true;

    if (IS_SYM(expr)) {
        return is_known_var(expr, params, nparams, locals, nlocals);
    }

    if (!IS_PAIR(expr)) return false;
    Value head = CAR(expr);

    if (head == SYM_IF) {
        Value rest = CDR(expr);
        if (!IS_PAIR(rest)) return false;
        if (!can_jit_expr_ext(CAR(rest), params, nparams, self_sym, locals, nlocals)) return false;
        rest = CDR(rest);
        if (!IS_PAIR(rest)) return false;
        if (!can_jit_expr_ext(CAR(rest), params, nparams, self_sym, locals, nlocals)) return false;
        rest = CDR(rest);
        if (IS_PAIR(rest))
            return can_jit_expr_ext(CAR(rest), params, nparams, self_sym, locals, nlocals);
        return true;
    }

    /* cond → chain of (test body...) clauses */
    if (head == SYM_COND) {
        Value clauses = CDR(expr);
        while (IS_PAIR(clauses)) {
            Value clause = CAR(clauses);
            if (!IS_PAIR(clause)) return false;
            Value test = CAR(clause);
            /* (else body) */
            if (IS_SYM(test) && strcmp(sym_name(test), "else") == 0) {
                Value body = CDR(clause);
                while (IS_PAIR(body)) {
                    if (!can_jit_expr_ext(CAR(body), params, nparams, self_sym, locals, nlocals)) return false;
                    body = CDR(body);
                }
                return true;
            }
            if (!can_jit_expr_ext(test, params, nparams, self_sym, locals, nlocals)) return false;
            Value body = CDR(clause);
            while (IS_PAIR(body)) {
                if (!can_jit_expr_ext(CAR(body), params, nparams, self_sym, locals, nlocals)) return false;
                body = CDR(body);
            }
            clauses = CDR(clauses);
        }
        return true;
    }

    /* and / or — short-circuit */
    if (head == SYM_AND || head == SYM_OR) {
        Value args = CDR(expr);
        while (IS_PAIR(args)) {
            if (!can_jit_expr_ext(CAR(args), params, nparams, self_sym, locals, nlocals))
                return false;
            args = CDR(args);
        }
        return true;
    }

    /* let / let* — check bindings and body with extended locals */
    if (head == SYM_LET || head == SYM_LET_STAR) {
        Value rest = CDR(expr);
        if (!IS_PAIR(rest)) return false;
        Value bindings_or_name = CAR(rest);

        /* Named let: (let name ((var init) ...) body ...) */
        if (head == SYM_LET && IS_SYM(bindings_or_name)) {
            Value loop_name = bindings_or_name;
            rest = CDR(rest);
            if (!IS_PAIR(rest)) return false;
            Value bindings = CAR(rest);
            Value body = CDR(rest);

            /* Collect binding variables */
            Value ext_locals[MAX_JIT_LOCALS];
            int n_ext = nlocals;
            if (n_ext > MAX_JIT_LOCALS) return false;
            for (int i = 0; i < nlocals; i++) ext_locals[i] = locals[i];

            Value bnd = bindings;
            int nbindings = 0;
            while (IS_PAIR(bnd)) {
                Value pair = CAR(bnd);
                if (!IS_PAIR(pair) || !IS_SYM(CAR(pair))) return false;
                if (!IS_PAIR(CDR(pair))) return false;
                /* Check init expr */
                if (!can_jit_expr_ext(CADR(pair), params, nparams, self_sym, locals, nlocals))
                    return false;
                if (n_ext < MAX_JIT_LOCALS) ext_locals[n_ext++] = CAR(pair);
                nbindings++;
                bnd = CDR(bnd);
            }
            if (nbindings == 0 || nbindings > MAX_JIT_LOCALS) return false;

            /* Body with extended locals + loop_name as self for recursive calls */
            while (IS_PAIR(body)) {
                if (!can_jit_expr_ext(CAR(body), params, nparams,
                                      loop_name, ext_locals, n_ext))
                    return false;
                body = CDR(body);
            }
            return true;
        }

        /* Regular let / let* */
        Value bindings = bindings_or_name;
        Value body = CDR(rest);

        Value ext_locals[MAX_JIT_LOCALS];
        int n_ext = nlocals;
        if (n_ext > MAX_JIT_LOCALS) return false;
        for (int i = 0; i < nlocals; i++) ext_locals[i] = locals[i];

        Value bnd = bindings;
        while (IS_PAIR(bnd)) {
            Value pair = CAR(bnd);
            if (!IS_PAIR(pair) || !IS_SYM(CAR(pair))) return false;
            if (!IS_PAIR(CDR(pair))) return false;
            /* For let*, init exprs can reference prior bindings */
            if (head == SYM_LET_STAR) {
                if (!can_jit_expr_ext(CADR(pair), params, nparams, self_sym, ext_locals, n_ext))
                    return false;
            } else {
                if (!can_jit_expr_ext(CADR(pair), params, nparams, self_sym, locals, nlocals))
                    return false;
            }
            if (n_ext < MAX_JIT_LOCALS) ext_locals[n_ext++] = CAR(pair);
            bnd = CDR(bnd);
        }

        while (IS_PAIR(body)) {
            if (!can_jit_expr_ext(CAR(body), params, nparams, self_sym, ext_locals, n_ext))
                return false;
            body = CDR(body);
        }
        return true;
    }

    if (IS_SYM(head)) {
        const char *name = sym_name(head);
        bool is_arith = (strcmp(name, "+") == 0 || strcmp(name, "-") == 0 ||
                         strcmp(name, "*") == 0 || strcmp(name, "=") == 0 ||
                         strcmp(name, "<") == 0 || strcmp(name, ">") == 0 ||
                         strcmp(name, "<=") == 0 || strcmp(name, ">=") == 0 ||
                         strcmp(name, "not") == 0 || strcmp(name, "zero?") == 0 ||
                         strcmp(name, "car") == 0 || strcmp(name, "cdr") == 0 ||
                         strcmp(name, "cons") == 0 || strcmp(name, "null?") == 0 ||
                         strcmp(name, "pair?") == 0);
        if (is_arith) {
            Value args = CDR(expr);
            while (IS_PAIR(args)) {
                if (!can_jit_expr_ext(CAR(args), params, nparams, self_sym, locals, nlocals))
                    return false;
                args = CDR(args);
            }
            return true;
        }

        /* Self-recursive call (or named-let loop call) */
        if (head == self_sym && !IS_NIL(self_sym)) {
            Value args = CDR(expr);
            int argc = 0;
            while (IS_PAIR(args)) {
                if (!can_jit_expr_ext(CAR(args), params, nparams, self_sym, locals, nlocals))
                    return false;
                argc++;
                args = CDR(args);
            }
            return argc == nparams;
        }
    }

    return false;
}

static bool can_jit_expr(Value expr, Value *params, int nparams, Value self_sym) {
    return can_jit_expr_ext(expr, params, nparams, self_sym, NULL, 0);
}

static bool can_jit_proc(Proc *proc) {
    if (proc->nparams > 6) return false;
    if (!IS_NIL(proc->rest)) return false;
    if (proc->has_defs) return false;
    if (proc->body.count < 1) return false;

    Value self_sym = proc->name ? intern(proc->name) : VAL_NIL;
    /* All body expressions must be JIT-compilable */
    for (int i = 0; i < proc->body.count; i++) {
        if (!can_jit_expr(proc->body.exprs[i], proc->params, proc->nparams, self_sym))
            return false;
    }
    return true;
}

/* ═══════════════════════════════════════════════════════════════════════════
 * Code generation
 *
 * Convention: result always in RAX (NaN-boxed).
 * Parameters live in callee-saved regs: r12-r15 (params 0-3),
 * stack slots for params 4-5.
 * r10, r11 used as scratch for boxing/unboxing.
 * ═══════════════════════════════════════════════════════════════════════════ */

/* ── Store to RBP-relative stack slot ── */
static void emit_store_rbp(JitCtx *j, int offset, int src_reg, bool src_ext) {
    uint8_t rex = REX_W;
    if (src_ext) rex |= 0x04;
    eb(j, rex); eb(j, 0x89);
    eb(j, MODRM(2, src_reg, RBP)); ei32(j, offset);
}

/* ── Load from RBP-relative stack slot ── */
static void emit_load_rbp(JitCtx *j, int dst_reg, bool dst_ext, int offset) {
    uint8_t rex = REX_W;
    if (dst_ext) rex |= 0x04;
    eb(j, rex); eb(j, 0x8B);
    eb(j, MODRM(2, dst_reg, RBP)); ei32(j, offset);
}

static bool emit_expr(JitCtx *j, Value expr, bool tail);

static int find_param(JitCtx *j, Value sym) {
    for (int i = 0; i < j->nparams; i++)
        if (j->params[i] == sym) return i;
    return -1;
}

/* Find a local variable, returns its stack slot offset or 0 if not found */
static int find_local(JitCtx *j, Value sym, bool *found) {
    /* Search in reverse so inner scopes shadow outer */
    for (int i = j->nlocals - 1; i >= 0; i--) {
        if (j->locals[i].sym == sym) {
            *found = true;
            return j->locals[i].slot;
        }
    }
    /* Also check named-let loop params */
    if (!IS_NIL(j->loop_sym)) {
        for (int i = 0; i < j->loop_nparams; i++) {
            if (j->loop_params[i] == sym) {
                *found = true;
                return j->loop_slots[i];
            }
        }
    }
    *found = false;
    return 0;
}

static void load_param(JitCtx *j, int idx) {
    if (idx < 4) {
        emit_mov_rr(j, RAX, false, R12 + idx, true);
    } else {
        int offset = -8 * (idx - 3);
        emit_load_rbp(j, RAX, false, offset);
    }
}

/* load_var unused for now — kept for future inlining work */

static bool emit_unboxed(JitCtx *j, Value expr) {
    if (!emit_expr(j, expr, false)) return false;
    emit_unbox_int(j, RAX, false);
    return true;
}

/* Evaluate a, b; compare; set flags. Returns comparison kind. */
static bool emit_comparison(JitCtx *j, Value a, Value b, CmpKind *kind, const char *op) {
    if (!emit_unboxed(j, a)) return false;
    emit_push(j, RAX, false);
    if (!emit_unboxed(j, b)) return false;
    emit_mov_rr(j, RCX, false, RAX, false);
    emit_pop(j, RAX, false);
    emit_cmp_rr(j, RAX, false, RCX, false);

    if      (strcmp(op, "=")  == 0) *kind = CMP_EQ;
    else if (strcmp(op, "<")  == 0) *kind = CMP_LT;
    else if (strcmp(op, ">")  == 0) *kind = CMP_GT;
    else if (strcmp(op, "<=") == 0) *kind = CMP_LE;
    else if (strcmp(op, ">=") == 0) *kind = CMP_GE;
    else return false;
    return true;
}

/* Emit CMOVcc: set RAX to VAL_TRUE or VAL_FALSE based on flags */
static void emit_cmov(JitCtx *j, CmpKind kind) {
    emit_mov_imm64(j, RAX, false, VAL_FALSE);
    emit_mov_imm64(j, RCX, false, VAL_TRUE);
    eb(j, REX_W); eb(j, 0x0F);
    switch (kind) {
        case CMP_EQ: eb(j, 0x44); break; /* CMOVE  */
        case CMP_NE: eb(j, 0x45); break; /* CMOVNE */
        case CMP_LT: eb(j, 0x4C); break; /* CMOVL  */
        case CMP_LE: eb(j, 0x4E); break; /* CMOVLE */
        case CMP_GT: eb(j, 0x4F); break; /* CMOVG  */
        case CMP_GE: eb(j, 0x4D); break; /* CMOVGE */
    }
    eb(j, MODRM(3, RAX, RCX));
}

/* Emit the negated conditional jump to else_target */
static int emit_negated_jump(JitCtx *j, CmpKind kind) {
    switch (kind) {
        case CMP_EQ: return emit_jne(j);
        case CMP_NE: return emit_je(j);
        case CMP_LT: return emit_jge(j);
        case CMP_LE: return emit_jg(j);
        case CMP_GT: return emit_jle(j);
        case CMP_GE: return emit_jl(j);
    }
    return emit_jne(j); /* unreachable */
}

static bool emit_expr(JitCtx *j, Value expr, bool tail) {
    /* ── Integer literal ── */
    if (IS_INT(expr)) {
        emit_mov_imm64(j, RAX, false, expr);
        if (tail) emit_return(j);
        return true;
    }

    /* ── Boolean / NIL literals ── */
    if (expr == VAL_TRUE || expr == VAL_FALSE || IS_NIL(expr)) {
        emit_mov_imm64(j, RAX, false, expr);
        if (tail) emit_return(j);
        return true;
    }

    /* ── Variable reference (local first for shadowing, then param) ── */
    if (IS_SYM(expr)) {
        bool found;
        int slot = find_local(j, expr, &found);
        if (found) {
            emit_load_rbp(j, RAX, false, slot);
            if (tail) emit_return(j);
            return true;
        }
        int idx = find_param(j, expr);
        if (idx >= 0) {
            load_param(j, idx);
            if (tail) emit_return(j);
            return true;
        }
        return false;
    }

    if (!IS_PAIR(expr)) return false;
    Value head = CAR(expr);
    Value rest = CDR(expr);

    /* ── (if test then else?) ── */
    if (head == SYM_IF) {
        Value test = CAR(rest);
        Value then_br = CADR(rest);
        Value else_rest = CDDR(rest);
        bool has_else = IS_PAIR(else_rest);
        Value else_br = has_else ? CAR(else_rest) : VAL_VOID;

        bool is_cmp = false;
        CmpKind cmp_kind = CMP_EQ;

        if (IS_PAIR(test) && IS_SYM(CAR(test))) {
            const char *tn = sym_name(CAR(test));
            if ((strcmp(tn,"=") == 0 || strcmp(tn,"<") == 0 || strcmp(tn,">") == 0 ||
                 strcmp(tn,"<=") == 0 || strcmp(tn,">=") == 0) &&
                IS_PAIR(CDR(test)) && IS_PAIR(CDDR(test)) && IS_NIL(CDR(CDDR(test)))) {
                if (!emit_comparison(j, CADR(test), CADDR(test), &cmp_kind, tn))
                    return false;
                is_cmp = true;
            } else if (strcmp(tn,"not") == 0 && IS_PAIR(CDR(test)) && IS_NIL(CDDR(test))) {
                if (!emit_expr(j, CADR(test), false)) return false;
                emit_mov_imm64(j, RCX, false, VAL_FALSE);
                emit_cmp_rr(j, RAX, false, RCX, false);
                cmp_kind = CMP_EQ; is_cmp = true;
            } else if (strcmp(tn,"zero?") == 0 && IS_PAIR(CDR(test)) && IS_NIL(CDDR(test))) {
                if (!emit_expr(j, CADR(test), false)) return false;
                emit_mov_imm64(j, RCX, false, VAL_INT(0));
                emit_cmp_rr(j, RAX, false, RCX, false);
                cmp_kind = CMP_EQ; is_cmp = true;
            }
        }

        if (!is_cmp) {
            if (!emit_expr(j, test, false)) return false;
            emit_mov_imm64(j, RCX, false, VAL_FALSE);
            emit_cmp_rr(j, RAX, false, RCX, false);
        }

        int else_jmp;
        if (is_cmp) {
            else_jmp = emit_negated_jump(j, cmp_kind);
        } else {
            /* Jump to else when result == #f (i.e., when ZF=1 from cmp) */
            else_jmp = emit_je(j);
        }

        if (!emit_expr(j, then_br, tail)) return false;
        if (!tail) {
            int end_jmp = emit_jmp(j);
            patch_jump(j, else_jmp);
            if (has_else) {
                if (!emit_expr(j, else_br, false)) return false;
            } else {
                emit_mov_imm64(j, RAX, false, VAL_VOID);
            }
            patch_jump(j, end_jmp);
        } else {
            patch_jump(j, else_jmp);
            if (has_else) {
                if (!emit_expr(j, else_br, true)) return false;
            } else {
                emit_mov_imm64(j, RAX, false, VAL_VOID);
                emit_return(j);
            }
        }
        return true;
    }

    /* ── cond → cascaded if/else ── */
    if (head == SYM_COND) {
        int end_jmps[32]; int n_ends = 0;
        Value clauses = rest;
        while (IS_PAIR(clauses)) {
            Value clause = CAR(clauses);
            Value test = CAR(clause);
            Value body = CDR(clause);
            clauses = CDR(clauses);
            bool is_else = IS_SYM(test) && strcmp(sym_name(test), "else") == 0;
            bool is_last = !IS_PAIR(clauses);

            if (is_else) {
                /* else clause — just emit the body */
                while (IS_PAIR(body) && IS_PAIR(CDR(body))) {
                    if (!emit_expr(j, CAR(body), false)) return false;
                    body = CDR(body);
                }
                if (IS_PAIR(body)) {
                    if (!emit_expr(j, CAR(body), tail)) return false;
                }
                break;
            }

            /* Emit test */
            bool is_cmp = false;
            CmpKind cmp_kind = CMP_EQ;
            if (IS_PAIR(test) && IS_SYM(CAR(test))) {
                const char *tn = sym_name(CAR(test));
                if ((strcmp(tn,"=") == 0 || strcmp(tn,"<") == 0 || strcmp(tn,">") == 0 ||
                     strcmp(tn,"<=") == 0 || strcmp(tn,">=") == 0) &&
                    IS_PAIR(CDR(test)) && IS_PAIR(CDDR(test)) && IS_NIL(CDR(CDDR(test)))) {
                    if (emit_comparison(j, CADR(test), CADDR(test), &cmp_kind, tn))
                        is_cmp = true;
                }
            }
            if (!is_cmp) {
                if (!emit_expr(j, test, false)) return false;
                emit_mov_imm64(j, RCX, false, VAL_FALSE);
                emit_cmp_rr(j, RAX, false, RCX, false);
            }
            int skip_jmp = is_cmp ? emit_negated_jump(j, cmp_kind) : emit_je(j);

            /* Emit body (last expr in tail position if this cond is tail) */
            while (IS_PAIR(body) && IS_PAIR(CDR(body))) {
                if (!emit_expr(j, CAR(body), false)) return false;
                body = CDR(body);
            }
            if (IS_PAIR(body)) {
                if (!emit_expr(j, CAR(body), tail && is_last)) return false;
            }
            if (!tail || !is_last) {
                if (n_ends < 32) end_jmps[n_ends++] = emit_jmp(j);
            }
            patch_jump(j, skip_jmp);
        }
        /* Patch all end jumps to here */
        for (int i = 0; i < n_ends; i++) patch_jump(j, end_jmps[i]);
        return true;
    }

    /* ── (and e1 e2 ...) — short-circuit: return first falsy, else last ── */
    if (head == SYM_AND) {
        Value args = rest;
        if (!IS_PAIR(args)) {
            /* (and) → #t */
            emit_mov_imm64(j, RAX, false, VAL_TRUE);
            if (tail) emit_return(j);
            return true;
        }
        int end_jmps[32]; int n_ends = 0;
        while (IS_PAIR(args)) {
            bool is_last = !IS_PAIR(CDR(args));
            if (!emit_expr(j, CAR(args), tail && is_last)) return false;
            if (!is_last) {
                /* If result is #f, short-circuit: jump to end */
                emit_mov_imm64(j, RCX, false, VAL_FALSE);
                emit_cmp_rr(j, RAX, false, RCX, false);
                if (n_ends < 32) end_jmps[n_ends++] = emit_je(j);
            }
            args = CDR(args);
        }
        for (int i = 0; i < n_ends; i++) patch_jump(j, end_jmps[i]);
        /* tail position already handled in the loop above */
        return true;
    }

    /* ── (or e1 e2 ...) — short-circuit: return first truthy, else last ── */
    if (head == SYM_OR) {
        Value args = rest;
        if (!IS_PAIR(args)) {
            /* (or) → #f */
            emit_mov_imm64(j, RAX, false, VAL_FALSE);
            if (tail) emit_return(j);
            return true;
        }
        int end_jmps[32]; int n_ends = 0;
        while (IS_PAIR(args)) {
            bool is_last = !IS_PAIR(CDR(args));
            if (!emit_expr(j, CAR(args), tail && is_last)) return false;
            if (!is_last) {
                /* If result is NOT #f, short-circuit: jump to end */
                emit_mov_imm64(j, RCX, false, VAL_FALSE);
                emit_cmp_rr(j, RAX, false, RCX, false);
                if (n_ends < 32) end_jmps[n_ends++] = emit_jne(j);
            }
            args = CDR(args);
        }
        for (int i = 0; i < n_ends; i++) patch_jump(j, end_jmps[i]);
        return true;
    }

    /* ── let / let* — allocate locals on stack ── */
    if (head == SYM_LET || head == SYM_LET_STAR) {
        Value r = rest;
        if (!IS_PAIR(r)) return false;
        Value bindings_or_name = CAR(r);

        /* ── Named let: (let name ((var init) ...) body ...) ── */
        if (head == SYM_LET && IS_SYM(bindings_or_name)) {
            Value loop_name = bindings_or_name;
            r = CDR(r);
            if (!IS_PAIR(r)) return false;
            Value bindings = CAR(r);
            Value body = CDR(r);

            /* Count bindings */
            int nbindings = 0;
            Value bnd = bindings;
            while (IS_PAIR(bnd)) { nbindings++; bnd = CDR(bnd); }
            if (nbindings == 0 || nbindings > MAX_JIT_LOCALS) return false;

            /* Save current state */
            int saved_nlocals = j->nlocals;
            Value saved_loop_sym = j->loop_sym;
            int saved_loop_start = j->loop_start;
            int saved_loop_nparams = j->loop_nparams;
            Value *saved_loop_params = j->loop_params;

            /* Pre-compute stack slots (but don't register yet — inits use outer scope) */
            Value loop_var_syms[MAX_JIT_LOCALS];
            int loop_var_slots[MAX_JIT_LOCALS];
            bnd = bindings;
            for (int i = 0; i < nbindings; i++) {
                Value pair = CAR(bnd);
                loop_var_syms[i] = CAR(pair);
                loop_var_slots[i] = j->local_base - 8 * (saved_nlocals + i + 1);
                bnd = CDR(bnd);
            }

            /* Evaluate init expressions BEFORE registering locals (outer scope) */
            bnd = bindings;
            for (int i = 0; i < nbindings; i++) {
                Value pair = CAR(bnd);
                if (!emit_expr(j, CADR(pair), false)) {
                    return false;
                }
                emit_store_rbp(j, loop_var_slots[i], RAX, false);
                bnd = CDR(bnd);
            }

            /* NOW register locals so body sees loop variables */
            for (int i = 0; i < nbindings; i++) {
                j->locals[j->nlocals].sym = loop_var_syms[i];
                j->locals[j->nlocals].slot = loop_var_slots[i];
                j->nlocals++;
            }

            /* Set up loop context */
            j->loop_sym = loop_name;
            j->loop_nparams = nbindings;
            j->loop_params = loop_var_syms;
            memcpy(j->loop_slots, loop_var_slots, sizeof(int) * nbindings);
            j->loop_start = (int)j->len;  /* mark loop entry point */

            /* Emit body */
            while (IS_PAIR(body) && IS_PAIR(CDR(body))) {
                if (!emit_expr(j, CAR(body), false)) {
                    j->nlocals = saved_nlocals;
                    j->loop_sym = saved_loop_sym;
                    return false;
                }
                body = CDR(body);
            }
            if (IS_PAIR(body)) {
                if (!emit_expr(j, CAR(body), tail)) {
                    j->nlocals = saved_nlocals;
                    j->loop_sym = saved_loop_sym;
                    return false;
                }
            }

            /* Restore state */
            j->nlocals = saved_nlocals;
            j->loop_sym = saved_loop_sym;
            j->loop_start = saved_loop_start;
            j->loop_nparams = saved_loop_nparams;
            j->loop_params = saved_loop_params;
            return true;
        }

        /* ── Regular let / let* ── */
        Value bindings = bindings_or_name;
        Value body = CDR(r);

        /* Count bindings */
        int nbindings = 0;
        Value bnd = bindings;
        while (IS_PAIR(bnd)) { nbindings++; bnd = CDR(bnd); }

        int saved_nlocals = j->nlocals;

        /* Evaluate and store bindings */
        bnd = bindings;
        for (int i = 0; i < nbindings; i++) {
            Value pair = CAR(bnd);
            if (!emit_expr(j, CADR(pair), false)) {
                j->nlocals = saved_nlocals;
                return false;
            }
            int slot = j->local_base - 8 * (j->nlocals + 1);
            emit_store_rbp(j, slot, RAX, false);
            j->locals[j->nlocals].sym = CAR(pair);
            j->locals[j->nlocals].slot = slot;
            j->nlocals++;
            bnd = CDR(bnd);
        }

        /* Emit body */
        while (IS_PAIR(body) && IS_PAIR(CDR(body))) {
            if (!emit_expr(j, CAR(body), false)) {
                j->nlocals = saved_nlocals;
                return false;
            }
            body = CDR(body);
        }
        if (IS_PAIR(body)) {
            if (!emit_expr(j, CAR(body), tail)) {
                j->nlocals = saved_nlocals;
                return false;
            }
        }

        j->nlocals = saved_nlocals;
        return true;
    }

    if (!IS_SYM(head)) return false;
    const char *opname = sym_name(head);

    /* ── Arithmetic: + - * ── */
    if (strcmp(opname, "+") == 0 || strcmp(opname, "-") == 0 || strcmp(opname, "*") == 0) {
        Value args = rest;
        int argc = 0;
        Value av[16];
        while (IS_PAIR(args) && argc < 16) { av[argc++] = CAR(args); args = CDR(args); }
        char op = opname[0];

        if (argc == 0) {
            emit_mov_imm64(j, RAX, false, op == '*' ? VAL_INT(1) : VAL_INT(0));
        } else if (argc == 1 && op == '-') {
            if (!emit_unboxed(j, av[0])) return false;
            emit_neg(j, RAX, false);
            emit_box_int(j, RAX, false);
        } else if (argc == 1) {
            if (!emit_expr(j, av[0], false)) return false;
        } else {
            if (!emit_unboxed(j, av[0])) return false;
            for (int i = 1; i < argc; i++) {
                emit_push(j, RAX, false);
                if (!emit_unboxed(j, av[i])) return false;
                emit_mov_rr(j, RCX, false, RAX, false);
                emit_pop(j, RAX, false);
                if      (op == '+') emit_add_rr(j, RAX, false, RCX, false);
                else if (op == '-') emit_sub_rr(j, RAX, false, RCX, false);
                else                emit_imul_rr(j, RAX, false, RCX, false);
            }
            emit_box_int(j, RAX, false);
        }
        if (tail) emit_return(j);
        return true;
    }

    /* ── Comparisons: = < > <= >= ── */
    if (strcmp(opname,"=") == 0 || strcmp(opname,"<") == 0 || strcmp(opname,">") == 0 ||
        strcmp(opname,"<=") == 0 || strcmp(opname,">=") == 0) {
        if (!IS_PAIR(rest) || !IS_PAIR(CDR(rest)) || !IS_NIL(CDDR(rest))) return false;
        CmpKind kind;
        if (!emit_comparison(j, CAR(rest), CADR(rest), &kind, opname)) return false;
        emit_cmov(j, kind);
        if (tail) emit_return(j);
        return true;
    }

    /* ── (not x) ── */
    if (strcmp(opname, "not") == 0) {
        if (!IS_PAIR(rest) || !IS_NIL(CDR(rest))) return false;
        if (!emit_expr(j, CAR(rest), false)) return false;
        emit_mov_imm64(j, RCX, false, VAL_FALSE);
        emit_cmp_rr(j, RAX, false, RCX, false);
        emit_cmov(j, CMP_EQ);
        if (tail) emit_return(j);
        return true;
    }

    /* ── (zero? x) ── */
    if (strcmp(opname, "zero?") == 0) {
        if (!IS_PAIR(rest) || !IS_NIL(CDR(rest))) return false;
        if (!emit_expr(j, CAR(rest), false)) return false;
        emit_mov_imm64(j, RCX, false, VAL_INT(0));
        emit_cmp_rr(j, RAX, false, RCX, false);
        emit_cmov(j, CMP_EQ);
        if (tail) emit_return(j);
        return true;
    }

    /* ── (null? x) — test against VAL_NIL ── */
    if (strcmp(opname, "null?") == 0) {
        if (!IS_PAIR(rest) || !IS_NIL(CDR(rest))) return false;
        if (!emit_expr(j, CAR(rest), false)) return false;
        emit_mov_imm64(j, RCX, false, VAL_NIL);
        emit_cmp_rr(j, RAX, false, RCX, false);
        emit_cmov(j, CMP_EQ);
        if (tail) emit_return(j);
        return true;
    }

    /* ── (pair? x) — check IS_PTR && obj_type == OBJ_PAIR ── */
    if (strcmp(opname, "pair?") == 0) {
        if (!IS_PAIR(rest) || !IS_NIL(CDR(rest))) return false;
        if (!emit_expr(j, CAR(rest), false)) return false;

        /* Inline IS_PAIR check:
         * Check QNAN bits, check TAG_PTR (0), dereference and check obj_type == OBJ_PAIR (0) */

        /* mov rcx, QNAN */
        emit_mov_imm64(j, RCX, false, QNAN);
        /* mov rdx, rax (save orig) */
        emit_mov_rr(j, RDX, false, RAX, false);
        /* and rax, rcx → should equal QNAN if NaN-boxed */
        {
            uint8_t rex = REX_W;
            eb(j, rex); eb(j, 0x21); eb(j, MODRM(3, RCX, RAX)); /* and rax, rcx */
        }
        emit_cmp_rr(j, RAX, false, RCX, false);
        int not_nan = emit_jne(j); /* if not NaN → is a double → false */
        /* Check tag bits: GET_TAG(rdx) == TAG_PTR (0) */
        emit_mov_rr(j, RAX, false, RDX, false);
        /* shr rax, 48 */
        eb(j, REX_W); eb(j, 0xC1); eb(j, MODRM(3, 5, RAX)); eb(j, 48);
        /* and rax, 7 */
        eb(j, REX_W); eb(j, 0x83); eb(j, MODRM(3, 4, RAX)); eb(j, 7);
        /* test rax, rax (should be 0 for TAG_PTR) */
        eb(j, REX_W); eb(j, 0x85); eb(j, MODRM(3, RAX, RAX));
        int not_ptr = emit_jne(j);
        /* It's a pointer — extract payload and check ObjType */
        emit_mov_imm64(j, RCX, false, PAYLOAD_MASK);
        emit_mov_rr(j, RAX, false, RDX, false);
        {
            eb(j, REX_W); eb(j, 0x21); eb(j, MODRM(3, RCX, RAX)); /* and rax, rcx */
        }
        /* Dereference: mov eax, [rax] (load ObjType, first 4 bytes) */
        eb(j, 0x8B); eb(j, MODRM(0, RAX, RAX));
        /* cmp eax, OBJ_PAIR (0) */
        eb(j, REX_W); eb(j, 0x83); eb(j, MODRM(3, 7, RAX)); eb(j, OBJ_PAIR);
        int not_pair = emit_jne(j);
        /* It's a pair! */
        emit_mov_imm64(j, RAX, false, VAL_TRUE);
        int done = emit_jmp(j);
        /* Not a pair */
        patch_jump(j, not_nan);
        patch_jump(j, not_ptr);
        patch_jump(j, not_pair);
        emit_mov_imm64(j, RAX, false, VAL_FALSE);
        patch_jump(j, done);

        if (tail) emit_return(j);
        return true;
    }

    /* ── (car x) — extract Pair.car from NaN-boxed pointer ── */
    if (strcmp(opname, "car") == 0) {
        if (!IS_PAIR(rest) || !IS_NIL(CDR(rest))) return false;
        if (!emit_expr(j, CAR(rest), false)) return false;
        /* Extract pointer: rax & PAYLOAD_MASK */
        emit_mov_imm64(j, RCX, false, PAYLOAD_MASK);
        { eb(j, REX_W); eb(j, 0x21); eb(j, MODRM(3, RCX, RAX)); }
        /* Load car field: offset = sizeof(ObjHeader) = 4, but aligned to 8 */
        /* Pair layout: ObjHeader(4 bytes) + 4 pad + Value car + Value cdr + int line */
        /* offsetof(Pair, car) */
        int car_off = (int)__builtin_offsetof(Pair, car);
        eb(j, REX_W); eb(j, 0x8B); eb(j, MODRM(2, RAX, RAX)); ei32(j, car_off);
        if (tail) emit_return(j);
        return true;
    }

    /* ── (cdr x) — extract Pair.cdr from NaN-boxed pointer ── */
    if (strcmp(opname, "cdr") == 0) {
        if (!IS_PAIR(rest) || !IS_NIL(CDR(rest))) return false;
        if (!emit_expr(j, CAR(rest), false)) return false;
        emit_mov_imm64(j, RCX, false, PAYLOAD_MASK);
        { eb(j, REX_W); eb(j, 0x21); eb(j, MODRM(3, RCX, RAX)); }
        int cdr_off = (int)__builtin_offsetof(Pair, cdr);
        eb(j, REX_W); eb(j, 0x8B); eb(j, MODRM(2, RAX, RAX)); ei32(j, cdr_off);
        if (tail) emit_return(j);
        return true;
    }

    /* ── (cons a b) — call C cons() function ── */
    if (strcmp(opname, "cons") == 0) {
        if (!IS_PAIR(rest) || !IS_PAIR(CDR(rest)) || !IS_NIL(CDDR(rest))) return false;
        /* Evaluate first arg */
        if (!emit_expr(j, CAR(rest), false)) return false;
        emit_push(j, RAX, false);
        /* Evaluate second arg */
        if (!emit_expr(j, CADR(rest), false)) return false;
        emit_mov_rr(j, RSI, false, RAX, false); /* arg1 = cdr */
        emit_pop(j, RDI, false);                 /* arg0 = car */
        /* Call cons() — must preserve callee-saved regs */
        emit_mov_imm64(j, RAX, false, (uint64_t)(uintptr_t)cons);
        eb(j, 0xFF); eb(j, MODRM(3, 2, RAX)); /* call rax */
        /* Result in rax (NaN-boxed pointer) */
        if (tail) emit_return(j);
        return true;
    }

    /* ── Named-let loop call — compile to jmp back to loop start ── */
    if (!IS_NIL(j->loop_sym) && head == j->loop_sym) {
        Value args = rest;
        int argc = 0;
        Value av[MAX_JIT_LOCALS];
        while (IS_PAIR(args) && argc < MAX_JIT_LOCALS) {
            av[argc++] = CAR(args);
            args = CDR(args);
        }
        if (argc != j->loop_nparams) return false;

        /* Evaluate all args to temp stack, then store into loop slots */
        for (int i = 0; i < argc; i++) {
            if (!emit_expr(j, av[i], false)) return false;
            emit_push(j, RAX, false);
        }
        for (int i = argc - 1; i >= 0; i--) {
            emit_pop(j, RAX, false);
            emit_store_rbp(j, j->loop_slots[i], RAX, false);
        }
        /* Jump back to loop start — true zero-overhead native loop */
        eb(j, 0xE9);
        int32_t rel = (int32_t)(j->loop_start - ((int)j->len + 4));
        ei32(j, rel);
        return true;
    }

    /* ── Self-recursive call ── */
    if (head == j->self_sym && !IS_NIL(j->self_sym)) {
        Value args = rest;
        int argc = 0;
        Value av[6];
        while (IS_PAIR(args) && argc < 6) { av[argc++] = CAR(args); args = CDR(args); }
        if (argc != j->nparams) return false;

        if (tail) {
            /* TCO: evaluate args to stack, pop into param regs, jmp body_start */
            for (int i = 0; i < argc; i++) {
                if (!emit_expr(j, av[i], false)) return false;
                emit_push(j, RAX, false);
            }
            for (int i = argc - 1; i >= 0; i--) {
                if (i < 4) {
                    emit_pop(j, R12 + i, true);
                } else {
                    emit_pop(j, RAX, false);
                    int offset = -8 * (i - 3);
                    eb(j, REX_W); eb(j, 0x89);
                    eb(j, MODRM(2, RAX, RBP)); ei32(j, offset);
                }
            }
            emit_tco_jump(j);
            return true;
        } else {
            /* Non-tail: callee-saved r12-r15 survive the call.
             * Evaluate args, push, pop into System V arg regs, call self. */
            for (int i = 0; i < argc; i++) {
                if (!emit_expr(j, av[i], false)) return false;
                emit_push(j, RAX, false);
            }
            for (int i = argc - 1; i >= 0; i--) {
                emit_pop(j, ARG_REGS[i], ARG_REG_EXT[i]);
            }
            /* call self via absolute address in rax */
            emit_mov_imm64(j, RAX, false, (uint64_t)(uintptr_t)j->entry);
            eb(j, 0xFF); eb(j, MODRM(3, 2, RAX)); /* call rax */
            /* Result in rax. Our r12-r15 params are preserved by callee. */
            return true;
        }
    }

    return false;
}

/* ═══════════════════════════════════════════════════════════════════════════
 * Top-level JIT compilation
 * ═══════════════════════════════════════════════════════════════════════════ */

JitBlock *jit_compile(Proc *proc) {
    if (!can_jit_proc(proc)) return NULL;

    JitCtx j;
    jctx_init(&j, 4096 * 4); /* 16KB */
    if (!j.buf) return NULL;

    j.params  = proc->params;
    j.nparams = proc->nparams;
    j.name    = proc->name;
    j.self_sym = proc->name ? intern(proc->name) : VAL_NIL;

    /* ═══ Prologue ═══ */
    emit_push(&j, RBP, false);              /* push rbp */
    emit_mov_rr(&j, RBP, false, RSP, false); /* mov rbp, rsp */
    emit_push(&j, RBX, false);              /* push rbx */
    emit_push(&j, R12, true);               /* push r12 */
    emit_push(&j, R13, true);               /* push r13 */
    emit_push(&j, R14, true);               /* push r14 */
    emit_push(&j, R15, true);               /* push r15 */

    /* Stack alignment: 6 pushes = 48 bytes. Entry was 16-aligned - 8 (ret addr).
     * After 6 pushes: offset = 56 = 16*3 + 8. Need sub 8 to align to 16.
     * Reserve extra space for local variables (let, named-let). */
    int param_slots = (proc->nparams > 4) ? 8 * (proc->nparams - 4) : 0;
    int local_slots = MAX_JIT_LOCALS * 8; /* reserve space for locals */
    j.extra_stack = 8 + param_slots + local_slots;
    j.extra_stack = (j.extra_stack + 15) & ~15;
    emit_sub_imm(&j, RSP, false, j.extra_stack);

    /* Set local_base: locals start after param spill area */
    j.local_base = -(param_slots + 8);

    /* Copy System V arg regs → callee-saved param regs */
    for (int i = 0; i < proc->nparams && i < 4; i++) {
        emit_mov_rr(&j, R12 + i, true, ARG_REGS[i], ARG_REG_EXT[i]);
    }
    for (int i = 4; i < proc->nparams; i++) {
        int offset = -8 * (i - 3);
        uint8_t rex = REX_W;
        if (ARG_REG_EXT[i]) rex |= 0x04;
        eb(&j, rex); eb(&j, 0x89);
        eb(&j, MODRM(2, ARG_REGS[i], RBP)); ei32(&j, offset);
    }

    /* Mark body start for TCO */
    j.body_start = (int)j.len;

    /* ═══ Emit body ═══ */
    for (int bi = 0; bi < proc->body.count; bi++) {
        bool is_last = (bi == proc->body.count - 1);
        if (!emit_expr(&j, proc->body.exprs[bi], is_last)) {
            munmap(j.buf, j.cap);
            return NULL;
        }
    }

    /* ═══ Epilogue ═══ */
    int epilogue = (int)j.len;
    emit_add_imm(&j, RSP, false, j.extra_stack);
    emit_pop(&j, R15, true);
    emit_pop(&j, R14, true);
    emit_pop(&j, R13, true);
    emit_pop(&j, R12, true);
    emit_pop(&j, RBX, false);
    emit_pop(&j, RBP, false);
    eb(&j, 0xC3); /* ret */

    /* ═══ Patch jumps ═══ */
    for (int i = 0; i < j.n_epilogue_jumps; i++) {
        int pos = j.epilogue_jumps[i];
        int32_t rel = (int32_t)(epilogue - (pos + 4));
        j.buf[pos]   = rel & 0xFF;
        j.buf[pos+1] = (rel >> 8) & 0xFF;
        j.buf[pos+2] = (rel >> 16) & 0xFF;
        j.buf[pos+3] = (rel >> 24) & 0xFF;
    }
    for (int i = 0; i < j.n_tco_jumps; i++) {
        int pos = j.tco_jumps[i];
        int32_t rel = (int32_t)(j.body_start - (pos + 4));
        j.buf[pos]   = rel & 0xFF;
        j.buf[pos+1] = (rel >> 8) & 0xFF;
        j.buf[pos+2] = (rel >> 16) & 0xFF;
        j.buf[pos+3] = (rel >> 24) & 0xFF;
    }

    /* ═══ Create JitBlock ═══ */
    JitBlock *block = (JitBlock *)ul_malloc(sizeof(JitBlock));
    block->code = j.buf;
    block->size = j.cap;
    block->func = (JitFunc)(void *)j.buf;
    block->name = proc->name ? ul_strdup(proc->name) : NULL;

    return block;
}

void jit_free(JitBlock *block) {
    if (!block) return;
    if (block->code) munmap(block->code, block->size);
    if (block->name) ul_free((char *)block->name);
    ul_free(block);
}
