(* A spike: Tast -> x86-64 machine code, in memory, called. Not a backend.
   The point is to find out what breaks, so the subset is deliberately tiny
   and every case it cannot do raises with the node that defeated it -- an
   honest [Unsupported] is the measurement, and a silently wrong answer is
   the one outcome that would waste the exercise.

   Register allocation is the trivial one the brief allows: every slot is a
   stack slot at [rbp - 8*(i+1)], every value is computed into rax, and a
   binary operator pushes its left operand. Two registers are enough for
   everything below and nothing is kept live across a statement. That is what
   makes an instruction selector tractable in an afternoon; it is also why the
   code it produces is four times the size of clang -O0's.

   Conventions, all of them SysV's, because raylib is called from this code:
     - integer arguments in rdi rsi rdx rcx r8 r9, then right-to-left on the
       stack; integer result in rax.
     - rsp % 16 == 0 at the [call] instruction. raylib spills xmm registers
       with movaps and faults far from the cause when this is wrong.
     - rbx rbp r12-r15 are callee-saved. This emitter touches none of them
       except rbp, which it saves.
     - every Flan function takes the transfer channel as a trailing ptr
       (emit.ml, [signature]), so a Flan function of n parameters is an n+1
       argument C function. *)

exception Unsupported of string

let unsupported fmt = Printf.ksprintf (fun s -> raise (Unsupported s)) fmt

(* ── Bytes ───────────────────────────────────────────────────────────── *)

type buf = { mutable bytes : Buffer.t }

let create () = { bytes = Buffer.create 256 }
let len b = Buffer.length b.bytes
let contents b = Buffer.contents b.bytes
let u8 b n = Buffer.add_char b.bytes (Char.chr (n land 0xff))

let u32 b n =
  for i = 0 to 3 do u8 b ((n asr (i * 8)) land 0xff) done

let i32 b (n : int) =
  if n < -0x80000000 || n > 0x7fffffff then unsupported "displacement %d" n;
  u32 b n

let u64 b (n : int64) =
  for i = 0 to 7 do
    u8 b (Int64.to_int (Int64.logand (Int64.shift_right_logical n (i * 8)) 0xffL))
  done

(* ── Registers and modrm ─────────────────────────────────────────────── *)

(* The encoding order, not the ABI order: this numbering *is* the three bits
   the modrm byte wants, which is why rsp is 4 and rbp is 5 rather than
   anything more memorable. *)
let rax = 0 and rcx = 1 and rdx = 2 and _rbx = 3
let rsp = 4 and rbp = 5 and rsi = 6 and rdi = 7
let r8 = 8 and r9 = 9

(* REX.W is always set: everything here is 64-bit. R extends the reg field and
   B the r/m field, which is the whole of what r8-r15 need. *)
let rex b ~r ~m = u8 b (0x48 lor (if r >= 8 then 4 else 0) lor (if m >= 8 then 1 else 0))
let modrm b ~md ~r ~m = u8 b ((md lsl 6) lor ((r land 7) lsl 3) lor (m land 7))

(* reg, reg *)
let rr b op ~r ~m = rex b ~r ~m; u8 b op; modrm b ~md:3 ~r ~m

(* reg, [rbp + disp32]. Always disp32 rather than the shorter disp8 form: a
   frame can outgrow 128 bytes and a one-byte displacement that silently wraps
   is exactly the bug this spike would not find. *)
let rm_rbp b op ~r ~disp =
  rex b ~r ~m:rbp; u8 b op; modrm b ~md:2 ~r ~m:rbp; i32 b disp

let mov_rr b ~dst ~src = rr b 0x89 ~r:src ~m:dst          (* mov dst, src *)
let mov_load b ~dst ~disp = rm_rbp b 0x8b ~r:dst ~disp    (* mov dst, [rbp+d] *)
let mov_store b ~src ~disp = rm_rbp b 0x89 ~r:src ~disp   (* mov [rbp+d], src *)

let movabs b ~dst (n : int64) =
  rex b ~r:0 ~m:dst; u8 b (0xb8 lor (dst land 7)); u64 b n

let push b r = if r >= 8 then u8 b 0x41; u8 b (0x50 lor (r land 7))
let pop b r = if r >= 8 then u8 b 0x41; u8 b (0x58 lor (r land 7))

let add_rr b ~dst ~src = rr b 0x01 ~r:src ~m:dst
let sub_rr b ~dst ~src = rr b 0x29 ~r:src ~m:dst
let and_rr b ~dst ~src = rr b 0x21 ~r:src ~m:dst
let or_rr  b ~dst ~src = rr b 0x09 ~r:src ~m:dst
let xor_rr b ~dst ~src = rr b 0x31 ~r:src ~m:dst
let imul_rr b ~dst ~src =                                  (* 0f af /r *)
  rex b ~r:dst ~m:src; u8 b 0x0f; u8 b 0xaf; modrm b ~md:3 ~r:dst ~m:src
let cmp_rr b ~a ~bb = rr b 0x39 ~r:bb ~m:a                 (* cmp a, b *)

let add_imm32 b ~dst n = rex b ~r:0 ~m:dst; u8 b 0x81; modrm b ~md:3 ~r:0 ~m:dst; i32 b n
let sub_imm32 b ~dst n = rex b ~r:0 ~m:dst; u8 b 0x81; modrm b ~md:3 ~r:5 ~m:dst; i32 b n

let call_r b r = if r >= 8 then u8 b 0x41; u8 b 0xff; modrm b ~md:3 ~r:2 ~m:r
let leave b = u8 b 0xc9
let ret b = u8 b 0xc3
let ud2 b = u8 b 0x0f; u8 b 0x0b

(* setcc al, then movzx rax, al -- a compare's result is a bool, which is one
   byte in Flan's layout (i1 in LLVM, and the ABI zero-extends it). *)
let setcc b cc = u8 b 0x0f; u8 b (0x90 lor cc); modrm b ~md:3 ~r:0 ~m:rax
let movzx_al b = u8 b 0x48; u8 b 0x0f; u8 b 0xb6; modrm b ~md:3 ~r:rax ~m:rax

(* jcc rel32 and jmp rel32, patched once the target is known. *)
let jcc b cc = u8 b 0x0f; u8 b (0x80 lor cc); let at = len b in u32 b 0; at
let jmp b = u8 b 0xe9; let at = len b in u32 b 0; at

let patch b ~at ~target =
  let rel = target - (at + 4) in
  let s = Buffer.contents b.bytes in
  let s = Bytes.of_string s in
  for i = 0 to 3 do
    Bytes.set s (at + i) (Char.chr ((rel asr (i * 8)) land 0xff))
  done;
  let nb = Buffer.create (Bytes.length s) in
  Buffer.add_bytes nb s;
  b.bytes <- nb

(* ── Lowering ────────────────────────────────────────────────────────── *)

type fnctx = {
  b : buf;
  nslots : int;
  (* How many 8-byte words this expression's evaluation has pushed since the
     prologue. rsp is 16-aligned at the end of the prologue, so [depth] even
     means rsp is aligned and [depth] odd means it is 8 out.

     This counter is the answer to the one bug the ABI probe found. Alignment
     is not a property of the prologue: the evaluator spills the left operand
     across the right one's evaluation, so a call written in the right operand
     runs with one word outstanding. Deriving it from a count kept here is the
     only way that stays correct as the evaluator grows cases, and it is what
     clang's [sub rsp, 8] before a call is doing. *)
  mutable depth : int;
  (* A symbol the code calls, resolved to an absolute address by the driver
     before emission. movabs + call r is what a JIT does anyway: a rel32 call
     cannot reach an arbitrary mmap, and the 2-byte indirect call is cheaper
     than the relocation machinery a real backend would grow here. *)
  resolve : string -> int64;
}

let slot_disp i = -8 * (i + 1)

(* Every stack movement goes through these two, so that nothing can move rsp
   without the counter noticing. *)
let pushv f r = push f.b r; f.depth <- f.depth + 1
let popv f r = pop f.b r; f.depth <- f.depth - 1

(* Every type this spike handles is one 8-byte integer register. Everything
   else is the real backend's problem and is enumerated in the verdict rather
   than guessed at here. *)
let word_ty (t : Flan.Types.t) =
  match t with
  | Flan.Types.Int _ | Flan.Types.Bool | Flan.Types.Ptr _ -> true
  | _ -> false

let check_word what (t : Flan.Types.t) =
  if not (word_ty t) then
    unsupported "%s of type %s: not a single integer register" what
      (Flan.Types.to_string t)

let cc_of signed (p : Flan.Tast.prim) =
  match p, signed with
  | Flan.Tast.Eq, _ -> 0x4 | Flan.Tast.Ne, _ -> 0x5
  | Flan.Tast.Lt, true -> 0xc | Flan.Tast.Lt, false -> 0x2
  | Flan.Tast.Le, true -> 0xe | Flan.Tast.Le, false -> 0x6
  | Flan.Tast.Gt, true -> 0xf | Flan.Tast.Gt, false -> 0x7
  | Flan.Tast.Ge, true -> 0xd | Flan.Tast.Ge, false -> 0x3
  | _ -> assert false

let arg_regs = [| rdi; rsi; rdx; rcx; r8; r9 |]

(* Value into rax. Everything is a subexpression of something that will
   immediately consume rax, so nothing is kept live and no allocator is
   needed. *)
let rec value f (e : Flan.Tast.expr) : unit =
  let b = f.b in
  match e.Flan.Tast.e with
  | Flan.Tast.Int (n, _) -> movabs b ~dst:rax n
  | Flan.Tast.Bool v -> movabs b ~dst:rax (if v then 1L else 0L)
  | Flan.Tast.Local i ->
    check_word "local" e.Flan.Tast.ty;
    if i >= f.nslots then unsupported "slot %d out of range" i;
    mov_load b ~dst:rax ~disp:(slot_disp i)
  | Flan.Tast.Do body -> block f body
  | Flan.Tast.Let (binds, body) ->
    List.iter
      (fun (i, e) ->
         value f e;
         check_word "binding" e.Flan.Tast.ty;
         mov_store b ~src:rax ~disp:(slot_disp i))
      binds;
    block f body
  | Flan.Tast.Set (Flan.Tast.Plocal i, rhs) ->
    value f rhs;
    check_word "assignment" rhs.Flan.Tast.ty;
    mov_store b ~src:rax ~disp:(slot_disp i)
  | Flan.Tast.If (c, t, e') -> emit_if f c t e'
  | Flan.Tast.Return (Some x) ->
    value f x;
    leave b; ret b
  | Flan.Tast.Return None -> leave b; ret b
  | Flan.Tast.Prim (p, args) -> prim f e p args
  | Flan.Tast.Call (name, args) -> call f (f.resolve name) args ~xfer:true
  | Flan.Tast.Unit -> ()
  | k -> unsupported "expression: %s" (node_name k)

and block f body =
  match body with
  | [] -> ()
  | [ last ] -> value f last
  | x :: rest -> value f x; block f rest

and prim f e (p : Flan.Tast.prim) args =
  let b = f.b in
  match p, args with
  | (Flan.Tast.Add | Flan.Tast.Sub | Flan.Tast.Mul
    | Flan.Tast.BitAnd | Flan.Tast.BitOr | Flan.Tast.BitXor), [ x; y ] ->
    check_word "arithmetic" x.Flan.Tast.ty;
    binop f x y;
    (* left in rax, right in rcx *)
    (match p with
     | Flan.Tast.Add -> add_rr b ~dst:rax ~src:rcx
     | Flan.Tast.Sub -> sub_rr b ~dst:rax ~src:rcx
     | Flan.Tast.Mul -> imul_rr b ~dst:rax ~src:rcx
     | Flan.Tast.BitAnd -> and_rr b ~dst:rax ~src:rcx
     | Flan.Tast.BitOr -> or_rr b ~dst:rax ~src:rcx
     | _ -> xor_rr b ~dst:rax ~src:rcx)
  | (Flan.Tast.Eq | Flan.Tast.Ne | Flan.Tast.Lt | Flan.Tast.Le
    | Flan.Tast.Gt | Flan.Tast.Ge), [ x; y ] ->
    let signed =
      match x.Flan.Tast.ty with
      | Flan.Types.Int k -> Flan.Types.signed k
      | Flan.Types.Bool -> false
      | t -> unsupported "comparison on %s" (Flan.Types.to_string t)
    in
    binop f x y;
    cmp_rr b ~a:rax ~bb:rcx;
    setcc b (cc_of signed p);
    movzx_al b
  | Flan.Tast.Rt sym, args -> call f (f.resolve sym) args ~xfer:false
  | _ -> unsupported "primitive in %s" (Flan.Types.to_string e.Flan.Tast.ty)

(* Left into rax, right into rcx, with the left spilled across the right's
   evaluation. Left-to-right, which emit.ml's [map_lr] is explicit about being
   required rather than a preference -- a call in either operand has effects.
   The push/pop pair keeps rsp 16-aligned in pairs, which matters only because
   [call] below re-derives alignment from a counter rather than tracking rsp. *)
and binop f x y =
  let b = f.b in
  value f x;
  pushv f rax;
  value f y;
  mov_rr b ~dst:rcx ~src:rax;
  popv f rax

and emit_if f c t e =
  let b = f.b in
  value f c;
  (* cmp rax, 0: 48 83 f8 00 -- written out because the helper above takes
     registers only and a zero-compare is the one immediate form worth having. *)
  u8 b 0x48; u8 b 0x83; modrm b ~md:3 ~r:7 ~m:rax; u8 b 0x00;
  let to_else = jcc b 0x4 in     (* je *)
  value f t;
  let to_end = jmp b in
  patch b ~at:to_else ~target:(len b);
  value f e;
  patch b ~at:to_end ~target:(len b)

(* A call, and this is the part that has to be exactly right.

   [xfer] appends the transfer channel, which every Flan function's signature
   carries and a C entry point does not. The spike passes NULL: nothing here
   signals, and a real backend would pass the caller's own channel pointer.

   Alignment: rsp is 16-aligned at function entry minus the 8 the [call]
   pushed, so after [push rbp] it is aligned again, and the frame is rounded to
   a multiple of 16. Every push here is paired with a pop before the next call
   can happen, so rsp is aligned at every call site by construction. Stack
   arguments are pushed in pairs to keep it that way -- an odd count gets a
   dummy push, which is what clang's [sub rsp, 8] is doing when you see it. *)
and call f (addr : int64) args ~xfer =
  let b = f.b in
  let n = List.length args + (if xfer then 1 else 0) in
  (* Bring rsp to 16 first, so everything below can count in pairs. *)
  let pad = f.depth land 1 = 1 in
  if pad then (sub_imm32 b ~dst:rsp 8; f.depth <- f.depth + 1);
  let stacked = List.filteri (fun i _ -> i >= 6) args in
  let nstack = List.length stacked + (if xfer && n > 6 then 1 else 0) in
  (* The stack half, evaluated right to left so that the seventh argument ends
     up at [rsp] and the eighth above it. The transfer channel is the last
     argument of all, so it is pushed first. *)
  if nstack land 1 = 1 then (sub_imm32 b ~dst:rsp 8; f.depth <- f.depth + 1);
  if xfer && n > 6 then (movabs b ~dst:rax 0L; pushv f rax);
  List.iter (fun a -> value f a; pushv f rax) (List.rev stacked);
  (* The register half needs a spill of its own: rdi..r9 are argument registers
     and rax is where every value lands, so an earlier argument would be
     clobbered by a later one's evaluation. Push each, then pop them into their
     registers in reverse. *)
  let inreg = List.filteri (fun i _ -> i < 6) args in
  List.iter (fun a -> value f a; pushv f rax) inreg;
  let nreg = List.length inreg in
  List.iteri (fun i _ -> popv f arg_regs.(nreg - 1 - i)) inreg;
  if xfer && n <= 6 then movabs b ~dst:arg_regs.(nreg) 0L;
  (* al = the number of vector registers used. Required only for a variadic
     callee and set unconditionally because it is two bytes: a wrong al on a
     printf-shaped entry point -- raylib's TraceLog is one -- is a crash that
     looks like anything else. After the argument registers, since al is rax's
     low byte. *)
  u8 b 0xb0; u8 b 0x00;          (* mov al, 0 *)
  (* r11 always, never r9: r11 is the scratch register SysV reserves and is the
     one register guaranteed not to be carrying an argument. Choosing the
     target conditionally is how a six-argument call gets quietly wrong. *)
  u8 b 0x49; u8 b 0xbb; u64 b addr;   (* movabs r11, addr *)
  assert (f.depth land 1 = 0);
  call_r b 11;
  let back = 8 * (nstack + (nstack land 1)) in
  if back > 0 then (add_imm32 b ~dst:rsp back; f.depth <- f.depth - (back / 8));
  if pad then (add_imm32 b ~dst:rsp 8; f.depth <- f.depth - 1)

and node_name (k : Flan.Tast.expr_kind) =
  match k with
  | Flan.Tast.Int _ -> "Int" | Flan.Tast.Float _ -> "Float"
  | Flan.Tast.Bool _ -> "Bool" | Flan.Tast.Str _ -> "Str"
  | Flan.Tast.Unit -> "Unit" | Flan.Tast.Zero _ -> "Zero"
  | Flan.Tast.Uninit _ -> "Uninit" | Flan.Tast.Local _ -> "Local"
  | Flan.Tast.Global _ -> "Global" | Flan.Tast.Prim _ -> "Prim"
  | Flan.Tast.Call _ -> "Call" | Flan.Tast.FnAddr _ -> "FnAddr"
  | Flan.Tast.CallPtr _ -> "CallPtr" | Flan.Tast.Do _ -> "Do"
  | Flan.Tast.Let _ -> "Let" | Flan.Tast.If _ -> "If"
  | Flan.Tast.While _ -> "While" | Flan.Tast.Return _ -> "Return"
  | Flan.Tast.Break _ -> "Break" | Flan.Tast.Continue _ -> "Continue"
  | Flan.Tast.Set _ -> "Set" | Flan.Tast.Field _ -> "Field"
  | Flan.Tast.Addr _ -> "Addr" | Flan.Tast.Deref _ -> "Deref"
  | Flan.Tast.Make _ -> "Make" | Flan.Tast.MakeCase _ -> "MakeCase"
  | Flan.Tast.CaseField _ -> "CaseField" | Flan.Tast.Arr _ -> "Arr"
  | Flan.Tast.Some_ _ -> "Some" | Flan.Tast.None_ -> "None"
  | Flan.Tast.Match _ -> "Match" | Flan.Tast.UnwrapSome _ -> "UnwrapSome"
  | Flan.Tast.Signal _ -> "Signal" | Flan.Tast.Handled _ -> "Handled"
  | Flan.Tast.RestartCase _ -> "RestartCase"
  | Flan.Tast.WithAlloc _ -> "WithAlloc"
  | Flan.Tast.InvokeRestart _ -> "InvokeRestart"

(* ── A whole function ────────────────────────────────────────────────── *)

let fn ~resolve (fd : Flan.Tast.fn) : string =
  let b = create () in
  let nslots = Array.length fd.Flan.Tast.slots in
  let f = { b; nslots; resolve; depth = 0 } in
  push b rbp;
  mov_rr b ~dst:rbp ~src:rsp;
  (* Round the frame to 16 so that rsp is aligned at every call site. One
     extra word for the transfer channel's slot, which is not a Flan slot and
     has no index -- the spike never reads it, but a real backend must, and
     leaving no room for it is the kind of thing that is cheap now and
     expensive later. *)
  let frame = (nslots + 1) * 8 in
  let frame = (frame + 15) land lnot 15 in
  if frame > 0 then sub_imm32 b ~dst:rsp frame;
  (* Parameters arrive in registers and are stored into their slots at once,
     which is also emit.ml's rule: slots 0..n-1 are the parameters, in order. *)
  let np = List.length fd.Flan.Tast.params in
  if np > 6 then unsupported "more than six parameters";
  List.iteri
    (fun i ty ->
       check_word "parameter" ty;
       mov_store b ~src:arg_regs.(i) ~disp:(slot_disp i))
    fd.Flan.Tast.params;
  (* The transfer channel is the last argument and goes just past the slots. *)
  if np < 6 then mov_store b ~src:arg_regs.(np) ~disp:(slot_disp nslots);
  block f fd.Flan.Tast.body;
  leave b; ret b;
  (* Anything that falls off the end of a Never-returning body lands here and
     traps rather than running into the next function. LLVM's [unreachable] is
     undefined behaviour; ud2 is a defined SIGILL, and the difference is one of
     the audit's findings. *)
  ud2 b;
  contents b