(* What redefining a generic function costs the dev loop, measured.

   The question the spike exists to answer: C-c C-c on a concrete function is
   about 35 ms today, and a generic function that is redefined has to rebuild
   *every* instantiation. So the sweep is one generic called at N concrete
   types, N = 1..8, against the handwritten N-copies program it replaces, and
   the three things a C-c C-c actually pays for are timed separately:

     check   Check.program_with_env over the whole accumulated program —
             which is what Session.eval does on every evaluation, so this is
             paid whether the redefined function is generic or not.
     emit    Emit.redefinition for the fns being installed.
     build   llc + ld -shared, from Build.shared — the dominant term.

   Nothing here modifies the session or the dev loop; it drives the real ones. *)

let tys = [| "i8"; "i16"; "i32"; "i64"; "u8"; "u16"; "u32"; "f32" |]

let time f =
  let t0 = Unix.gettimeofday () in
  let x = f () in
  (x, (Unix.gettimeofday () -. t0) *. 1000.)

(* Best of k, because llc and the linker are processes and the machine is
   noisy; a median would hide a systematic cost and a mean would report the
   scheduler. *)
let best k f =
  let rec go i acc = if i = 0 then acc else
      let _, ms = time f in go (i - 1) (min acc ms) in
  go k infinity

let generic_src n =
  let b = Buffer.create 1024 in
  Buffer.add_string b
    "(defn gswap [xs [$t] i i32 j i32] ()\n\
    \  (let [tmp (at xs i)]\n\
    \    (set (at xs i) (at xs j))\n\
    \    (set (at xs j) tmp)))\n\n\
     (defn gsort [s [$t] before? (Fn [$t $t] bool)] ()\n\
    \  (let [i 1]\n\
    \    (while (< i (len s))\n\
    \      (let [j i]\n\
    \        (while (and (> j 0) (before? (at s j) (at s (- j 1))))\n\
    \          (gswap s (- j 1) j)\n\
    \          (set j (- j 1))))\n\
    \      (set i (+ i 1)))))\n\n";
  for i = 0 to n - 1 do
    Buffer.add_string b (Printf.sprintf "(defvar xs-%s [8 %s])\n" tys.(i) tys.(i))
  done;
  Buffer.add_string b "\n(defn main [] ()\n";
  for i = 0 to n - 1 do
    Buffer.add_string b
      (Printf.sprintf "  (gsort (slice xs-%s 0 8) (fn [a b] (< a b)))\n" tys.(i))
  done;
  Buffer.add_string b "  )\n";
  Buffer.contents b

(* The same program as it is written today: one copy of each function per
   element type, by hand. This is prelude.ml's shape. *)
let mono_src n =
  let b = Buffer.create 1024 in
  for i = 0 to n - 1 do
    let t = tys.(i) in
    Buffer.add_string b
      (Printf.sprintf
         "(defn mswap-%s [xs [%s] i i32 j i32] ()\n\
         \  (let [tmp (at xs i)]\n\
         \    (set (at xs i) (at xs j))\n\
         \    (set (at xs j) tmp)))\n\n\
          (defn msort-%s [s [%s] before? (Fn [%s %s] bool)] ()\n\
         \  (let [i 1]\n\
         \    (while (< i (len s))\n\
         \      (let [j i]\n\
         \        (while (and (> j 0) (before? (at s j) (at s (- j 1))))\n\
         \          (mswap-%s s (- j 1) j)\n\
         \          (set j (- j 1))))\n\
         \      (set i (+ i 1)))))\n\n"
         t t t t t t t);
    Buffer.add_string b (Printf.sprintf "(defvar xs-%s [8 %s])\n\n" t t)
  done;
  Buffer.add_string b "(defn main [] ()\n";
  for i = 0 to n - 1 do
    Buffer.add_string b
      (Printf.sprintf "  (msort-%s (slice xs-%s 0 8) (fn [a b] (< a b)))\n"
         tys.(i) tys.(i))
  done;
  Buffer.add_string b "  )\n";
  Buffer.contents b

let write path s =
  let oc = open_out path in output_string oc s; close_out oc

let dir =
  let d = Filename.concat (Filename.get_temp_dir_name ()) "flan-generics-spike" in
  (try Unix.mkdir d 0o700 with Unix.Unix_error (Unix.EEXIST, _, _) -> ());
  d

let decls_of path =
  (Flan.Load.program ~file:path
     (Flan.Parse.program (Flan.Reader.read_file path))).Flan.Load.decls

(* Every function the program ended up with whose name starts with one of the
   generic names — the instantiations, which is exactly what a redefinition of
   the generic would have to rebuild. *)
let instantiations (p : Flan.Tast.program) =
  List.filter_map
    (fun (f : Flan.Tast.fn) ->
       let n = f.Flan.Tast.name in
       if String.length n > 6 && String.sub n 0 6 = "gswap-" then Some n
       else if String.length n > 6 && String.sub n 0 6 = "gsort-" then Some n
       else None)
    p.Flan.Tast.fns

let build_ms ir =
  let out = Filename.concat dir "redef.so" in
  best 3 (fun () ->
    ignore
      (Flan.Build.shared
         ~opts:{ Flan.Build.default with dev = true } ~ir ~out ()))

let () =
  Printf.printf
    "n   check-gen  check-mono   emit-1  emit-N   build-1  build-N   fns\n";
  (try
  for n = 1 to 8 do
    let gpath = Filename.concat dir (Printf.sprintf "gen%d.flan" n) in
    let mpath = Filename.concat dir (Printf.sprintf "mono%d.flan" n) in
    write gpath (generic_src n);
    write mpath (mono_src n);
    let gd = decls_of gpath and md = decls_of mpath in
    let check_gen = best 3 (fun () -> ignore (Flan.Check.program_with_env gd)) in
    let check_mono = best 3 (fun () -> ignore (Flan.Check.program_with_env md)) in
    let p, _ = Flan.Check.program_with_env gd in
    let insts = instantiations p in
    let one = [ List.hd insts ] in
    let ir_one =
      Flan.Emit.redefinition ~dev:true ~known:(fun _ -> true) p ~fns:one
    in
    let ir_all =
      Flan.Emit.redefinition ~dev:true ~known:(fun _ -> true) p ~fns:insts
    in
    let emit1 =
      best 3 (fun () ->
        ignore (Flan.Emit.redefinition ~dev:true ~known:(fun _ -> true) p ~fns:one))
    and emitn =
      best 3 (fun () ->
        ignore (Flan.Emit.redefinition ~dev:true ~known:(fun _ -> true) p ~fns:insts))
    in
    let b1 = build_ms ir_one and bn = build_ms ir_all in
    Printf.printf "%d   %8.1f  %10.1f  %7.1f %7.1f  %8.1f %8.1f   %d\n%!"
      n check_gen check_mono emit1 emitn b1 bn (List.length insts)
  done
   with Flan.Loc.Error d -> prerr_endline (Flan.Loc.report d); exit 1);
  (* And what the session actually does when the generic itself is redefined.
     This is the real C-c C-c path — Session.eval on the form the editor sent
     — and what it reports is the finding, not the timing. *)
  let gpath = Filename.concat dir "gen4.flan" in
  let t, _ = Flan.Session.create ~file:gpath () in
  let form =
    "(defn gswap [xs [$t] i i32 j i32] ()\n\
    \  (let [tmp (at xs i)]\n\
    \    (set (at xs i) (at xs j))\n\
    \    (set (at xs j) tmp)))\n"
  in
  let c, ms = time (fun () -> Flan.Session.eval ~origin:gpath t form) in
  Printf.printf
    "\nSession.eval on the generic gswap itself: %.1f ms, installs=%b, \
     fns=[%s], names=[%s]\n"
    ms c.Flan.Session.installs
    (String.concat " " c.Flan.Session.fns)
    (String.concat " " c.Flan.Session.names)