(* The spike's harness: run the real frontend, lower the functions it produced
   with [X86], put the bytes in executable memory, call them, and compare with
   what the language says they should answer.

   The comparison is the whole point. Reading the bytes proves nothing -- a
   disassembly that looks right and a program that returns the wrong number is
   the normal outcome of hand-encoding, which is why the oracle here is the
   arithmetic and not objdump. [oracle.sh] disassembles the same buffer, and
   that is a debugging aid, not the evidence. *)

external jit_alloc : int -> nativeint = "spike_jit_alloc"
external jit_write : nativeint -> string -> unit = "spike_jit_write"
external jit_protect : nativeint -> int -> unit = "spike_jit_protect"
external call1 : nativeint -> int64 -> int64 = "spike_call1"
external call2 : nativeint -> int64 -> int64 -> int64 = "spike_call2"
external sym : string -> nativeint = "spike_sym"

let failures = ref 0
let checks = ref 0

let check name got want =
  incr checks;
  if got = want then Printf.printf "  ok    %-28s = %Ld\n" name got
  else begin
    incr failures;
    Printf.printf "  FAIL  %-28s = %Ld, want %Ld\n" name got want
  end

(* One page per function, so that a function that runs off its own end lands in
   an unmapped page and segfaults at the fault rather than in the middle of the
   next function. This is the crudest possible version of the code-object
   question the whole exercise is really about. *)
let page = 4096

let install (code : string) : nativeint =
  if String.length code > page then failwith "function exceeds one page";
  let p = jit_alloc page in
  jit_write p code;
  jit_protect p page;
  p

let run src =
  let decls =
    Flan.Load.program ~file:src (Flan.Parse.program_all (Flan.Reader.read_file src))
  in
  let prog = Flan.Check.program_all decls.Flan.Load.decls in
  Printf.printf "frontend: %d fns, %d globals, %d structs, %d externs\n"
    (List.length prog.Flan.Tast.fns) (List.length prog.Flan.Tast.globals)
    (List.length prog.Flan.Tast.structs) (List.length prog.Flan.Tast.externs);

  (* Two passes, because [spike-calls] calls functions whose addresses are not
     known until they are installed. Pass one installs every function at a
     fixed page; pass two emits the real code into it. A real backend does this
     with relocations; the spike does it by emitting twice, which is the same
     answer with none of the machinery. *)
  let addrs : (string, nativeint) Hashtbl.t = Hashtbl.create 16 in
  let unsupported = ref [] in
  let lowerable =
    List.filter
      (fun (fd : Flan.Tast.fn) ->
         try
           ignore (X86.fn ~resolve:(fun _ -> 0L) fd);
           true
         with X86.Unsupported m ->
           unsupported := (fd.Flan.Tast.name, m) :: !unsupported;
           false)
      prog.Flan.Tast.fns
  in
  List.iter
    (fun (fd : Flan.Tast.fn) ->
       Hashtbl.replace addrs fd.Flan.Tast.name (jit_alloc page))
    lowerable;
  let resolve name =
    match Hashtbl.find_opt addrs name with
    | Some p -> Int64.of_nativeint p
    | None ->
      (* Not a Flan function: a runtime entry point, looked up the way a dev
         build already reaches the host's symbols -- through the dynamic symbol
         table, which --dev links with -rdynamic. *)
      Int64.of_nativeint (sym name)
  in
  let bytes = Hashtbl.create 16 in
  List.iter
    (fun (fd : Flan.Tast.fn) ->
       let code = X86.fn ~resolve fd in
       Hashtbl.replace bytes fd.Flan.Tast.name code;
       let p = Hashtbl.find addrs fd.Flan.Tast.name in
       jit_write p code;
       jit_protect p page)
    lowerable;

  Printf.printf "lowered: %d of %d functions\n"
    (List.length lowerable) (List.length prog.Flan.Tast.fns);
  List.iter (fun (n, m) -> Printf.printf "  skipped %-20s %s\n" n m)
    (List.rev !unsupported);
  Hashtbl.iter (fun n c -> Printf.printf "  %-20s %4d bytes at %nx\n"
                   n (String.length c) (Hashtbl.find addrs n)) bytes;

  (* The bytes that actually ran, dumped where run.sh can objdump them.
     A debugging aid and not the evidence: a disassembly that reads correctly
     next to a function that answers 656 when it should answer 650 is the
     normal outcome of hand-encoding, which is why the checks below compare
     numbers. *)
  (match Sys.getenv_opt "SPIKE_DUMP" with
   | None -> ()
   | Some dir ->
     Hashtbl.iter
       (fun n c ->
          let oc = open_out_bin (Filename.concat dir (n ^ ".bin")) in
          output_string oc c; close_out oc)
       bytes);

  (* ── The SysV boundary ──────────────────────────────────────────────
     Three synthetic functions, built as Tast by hand rather than written in
     Flan, because the surface language has no way to spell a call to an
     arbitrary C symbol with eight arguments. [Tast.Rt] is the node a runtime
     call already uses and the one a [declare-c] shim lands on, so this is the
     real path with a made-up callee. *)
  let loc = Flan.Loc.unknown in
  let i64 = Flan.Types.Int Flan.Types.I64 in
  let ex e = { Flan.Tast.e; ty = i64; loc } in
  let lit n = ex (Flan.Tast.Int (Int64.of_int n, Flan.Types.I64)) in
  let probe name params body =
    { Flan.Tast.name; params; slots = Array.make (List.length params) i64;
      snames = Array.make (List.length params) None; ret = i64;
      body = [ body ]; fdefers = []; fparent = None; floc = loc }
  in
  let arg0 = ex (Flan.Tast.Local 0) in
  let probes = [
    (* Eight integers: six in registers and two on the stack, which is the case
       a register-only convention gets silently wrong. *)
    probe "abi-8" [ i64 ]
      (ex (Flan.Tast.Prim (Flan.Tast.Rt "spike_probe8",
         [ arg0; lit 2; lit 3; lit 4; lit 5; lit 6; lit 7; lit 8 ])));
    (* rsp % 16 == 0 at the call. The callee does an aligned 16-byte spill and
       answers -1 if it was entered misaligned. *)
    probe "abi-align" [ i64 ]
      (ex (Flan.Tast.Prim (Flan.Tast.Rt "spike_probe_align", [ arg0 ])));
    (* The same call, but underneath a binary operator -- so it is evaluated
       with the left operand spilled on the stack. This is the one that matters:
       alignment at a call site is not a property of the prologue, it is a
       property of how much the expression evaluator has pushed. *)
    probe "abi-align-nested" [ i64 ]
      (ex (Flan.Tast.Prim (Flan.Tast.Add,
         [ lit 0;
           ex (Flan.Tast.Prim (Flan.Tast.Rt "spike_probe_align", [ arg0 ])) ])));
  ] in
  List.iter
    (fun (fd : Flan.Tast.fn) ->
       let code = X86.fn ~resolve fd in
       let p = jit_alloc page in
       jit_write p code; jit_protect p page;
       Hashtbl.replace addrs fd.Flan.Tast.name p)
    probes;

  print_endline "results:";
  let at n = Hashtbl.find addrs n in
  check "spike-add 3 4" (call2 (at "spike-add") 3L 4L) 7L;
  check "spike-add -5 2" (call2 (at "spike-add") (-5L) 2L) (-3L);
  check "spike-arith 10 4" (call2 (at "spike-arith") 10L 4L) 19L;
  check "spike-let 6" (call1 (at "spike-let") 6L) 1332L;
  check "spike-if 1 2" (call2 (at "spike-if") 1L 2L) 1L;
  check "spike-if 9 2" (call2 (at "spike-if") 9L 2L) 7L;
  check "spike-calls 5" (call1 (at "spike-calls") 5L) 656L;
  check "abi-8 1" (call1 (at "abi-8") 1L) 87654321L;
  check "abi-align 10" (call1 (at "abi-align") 10L) 13L;
  check "abi-align-nested 10" (call1 (at "abi-align-nested") 10L) 13L;

  Printf.printf "\n%d checks, %d failures\n" !checks !failures;
  exit (if !failures = 0 then 0 else 1)

(* The frontend's diagnostics printed rather than swallowed: a spike that says
   [Fatal error: exception Errors(_)] costs an hour. *)
let () =
  try run Sys.argv.(1) with
  | Flan.Loc.Error d -> prerr_endline (Flan.Loc.report d); exit 2
  | Flan.Loc.Errors ds -> prerr_endline (Flan.Loc.report_all ds); exit 2