flan/test/test_reload.ml

(* The reload primitive, measured (NEXT.md, dev loop step 1).

   One function is recompiled into its own object and loaded into a process
   that is already running. Everything after this — indirection cells, the
   agent in the game, the daemon — assumes this works and is fast; nothing in
   the codebase had ever done it, and plan.org's 16ms was measured with clang
   in isolation somewhere else.

   The parts, all of them new here:

     Emit.program ~dev   a cell per function; every call goes through one
     Emit.redefinition   a form list defined, everything else [external],
                         plus [flan_reload_install] to publish it into its cell
     flan_dev.c          the by-name registry a run-time-new name needs
     Build.shared        that IR text through llc + ld -shared, timed
     reload_host.c       dlopen, install, call — twice, in one process

   The host is C rather than OCaml because that is where it has to end up: the
   agent of step 3 lives in the game process, next to flan_rt.c, and there is
   no OCaml runtime there. *)

open Flan

let failures = ref 0
let fail fmt = Printf.ksprintf (fun s -> incr failures; print_endline ("FAIL  " ^ s)) fmt

let scratch = Filename.get_temp_dir_name ()
let tmp name = Filename.concat scratch ("flan-reload-" ^ name)

let checked path =
  Check.program
    (Load.program ~file:path (Parse.program (Reader.read_file path))).Load.decls

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

let () =
  match Sys.command "command -v clang > /dev/null 2>&1 && command -v llc > /dev/null 2>&1" with
  | 0 ->
    let p1 = checked "programs/reload.flan" in
    let p2 = checked "programs/reload-v2.flan" in
    let p3 = checked "programs/reload-v3.flan" in
    let p4 = checked "programs/reload-v4.flan" in

    (* What the running process was built with. Everything else — v3's [extra]
       and [added] — has no symbol to bind to and goes through the registry.
       A session would keep this set and grow it; the test states it. *)
    let host_names =
      List.map (fun (f : Tast.fn) -> f.Tast.name) p1.Tast.fns
      @ List.map (fun (g : Tast.global) -> g.Tast.gname) p1.Tast.globals
    in
    let known n = List.exists (String.equal n) host_names in

    (* [dev] is the two halves of a reloadable build together: cells, so a
       call site can be made to follow a redefinition, and [-rdynamic], so the
       cells and globals are visible to a dlopen'd object at all. [-ldl] is the
       host's own, for its dlopen. *)
    let dev = { Build.default with Build.dev = true } in
    let host = tmp "host" in
    ignore
      (Build.executable ~opts:dev ~csrcs:[ "reload_host.c" ]
         ~lflags:[ "-ldl" ] p1 ~out:host);

    (* Two paths, not one rewritten in place: dlopen caches by path and would
       hand back the first handle, so the swap would silently not happen. *)
    let module_of p fns name =
      let out = tmp name in
      let ir, emit_ms = ms (fun () -> Emit.redefinition ~dev:true ~known p ~fns) in
      let t = Build.shared ~opts:dev ~ir ~out () in
      (out, ir, emit_ms, t)
    in
    let so1, _ir1, emit_ms, t1 = module_of p1 [ "bump" ] "v1.so" in
    let so2, ir2, emit2_ms, t2 = module_of p2 [ "bump" ] "v2.so" in
    (* One module, two forms: the var and the function that uses it have to
       arrive together or the intermediate state refers to storage that does
       not exist. This is the C-c C-k unit. *)
    let so3, ir3, _, _ = module_of p3 [ "bump"; "added" ] "v3.so" in
    let so4, ir4, _, _ = module_of p4 [ "added" ] "v4.so" in

    (* A redefinition module must not define what the host already owns:
       defining [counter] would give the loaded object a private copy and the
       state would reset on every reload, and defining [helper] would freeze a
       stale copy of it into the module. *)
    let has hay needle =
      let n = String.length needle and h = String.length hay in
      let rec go i = i + n <= h && (String.sub hay i n = needle || go (i + 1)) in
      go 0
    in
    (* No Str, for the same reason the reader is hand-written. [`First] and
       [`Last] are which occurrence; the pair shape keeps the match below
       readable. *)
    let find hay needle which =
      let n = String.length needle and h = String.length hay in
      let rec go i acc =
        if i + n > h then acc
        else if String.sub hay i n = needle then
          match which with `First -> Some i | `Last -> go (i + 1) (Some i)
        else go (i + 1) acc
      in
      go 0 None
    in
    if not (has ir2 "@\"flan.counter\" = external global i64") then
      fail "redefinition defines the global instead of declaring it";
    (* In a dev module a sibling is reached only through its cell, so there is
       nothing to declare and a [define] would be a private copy. *)
    if has ir2 "declare i64 @\"flan.helper\"" then
      fail "dev redefinition declares a sibling it should reach by cell";
    if has ir2 "define i64 @\"flan.helper\"" then
      fail "redefinition emitted a second body for a function it does not own";
    if has ir2 "define i32 @main" then fail "redefinition emitted an entry point";
    (* [hidden], or the module's own [@"flan.bump"] is interposed by the host's
       and the installer publishes the very function it is replacing. *)
    if not (has ir2 "define hidden i64 @\"flan.bump\"") then
      fail "redefinition's own body is interposable";
    if not (has ir2 "@\"flan.cell.helper\" = external global ptr") then
      fail "redefinition defines a cell instead of using the host's";
    (* A name the host has is a symbol; a name it lacks is a registry lookup
       cached in a module-local slot. Getting this backwards either fails to
       link or silently gives each module its own copy. *)
    if not (has ir3 "@\"flan.cellp.added\" = internal global ptr null") then
      fail "a run-time-new function did not get a slot";
    if not (has ir3 "@\"flan.gp.extra\" = internal global ptr null") then
      fail "a run-time-new global did not get a slot";
    if has ir3 "@\"flan.extra\" = " then
      fail "a run-time-new global was given storage in the module";
    (* Every lookup is resolved before any body is published: publishing first
       exposes a function whose module-local slots are still null to anything
       that calls it. Not race-testable, so it is asserted on the text. *)
    (* v4 redefines a name that exists only in the registry, so it publishes
       through the cell it looked up rather than into a symbol — there is no
       [@"flan.cell.added"] anywhere to store into. *)
    if has ir4 "@\"flan.cell.added\"" then
      fail "a run-time-new function was published into a symbol";
    if not (has ir4 "call ptr @flan_dev_cell") then
      fail "v4 did not look its target up by name";
    (match find ir3 "store ptr @\"flan." `First, find ir3 "@flan_dev_(" `Last with
     | Some publish, Some resolve when resolve > publish ->
       fail "flan_reload_install publishes a body before resolving a lookup"
     | None, _ -> fail "flan_reload_install publishes nothing"
     | _ -> ());
    (* A dev host's calls are indirect; a release host's are not. That is the
       only difference between the two, and the whole of C-c C-c rests on it. *)
    let host_ir = Emit.program ~dev:true p1 in
    if not (has host_ir "@\"flan.cell.bump\" = global ptr @\"flan.bump\"") then
      fail "dev build emitted no cell";
    if has (Emit.program p1) "flan.cell." then
      fail "release build emitted a cell";

    let out = tmp "out" in
    let cmd =
      (* stderr kept apart from stdout: the host times its own dlopen there,
         and stdout is what the expected transcript is compared against. *)
      Printf.sprintf "%s %s %s %s %s > %s 2> %s" (Filename.quote host)
        (Filename.quote so1) (Filename.quote so2) (Filename.quote so3)
        (Filename.quote so4) (Filename.quote out)
        (Filename.quote (tmp "err"))
    in
    let (code, dlopen_ms) = ms (fun () -> Sys.command cmd) in
    let text = In_channel.with_open_bin out In_channel.input_all in
    let timings = In_channel.with_open_bin (tmp "err") In_channel.input_all in

    (* Every call is [outer], compiled once into the host and never rebuilt, so
       a changed answer can only mean its call site followed the redefinition.
       The arithmetic, in order:

         host   counter 0 -> 1,   helper 1 = 2
         v1     counter 1 -> 2,   helper 2 = 4        (a rebuild of the same)
         v2     +10 and +1000, recursing through its own cell until the
                counter passes 100: 2 -> 12 -> ... -> 102, ten "v2" lines,
                helper 102 = 204, so 1204. An interposed self-call would reach
                the host's v1 body, print "v1", and land nowhere near it.
         v3     extra 0 -> 7, counter 102 -> 109, helper 109 = 218. [extra]
                and [added] are new names, so both came from the registry.
         v4     redefines [added] only. v3's [bump] is still the installed one
                and is not rebuilt here, so it reaches v4 only through a cell
                the two modules found by the same name: extra 7 -> 107,
                counter 109 -> 216, helper 216 = 432. Had v3 cached the
                function's address rather than its cell's, this would be 246.

       The "v1"/"v2"/"v3" lines come from inside each [bump] and are what
       exercise a redefinition module's own string constants. 1204 rather than
       1236 is [helper]: v2's text for it multiplies by three, and the module
       declares it rather than defining it, so the host's copy is the one that
       ran. *)
    let v2s = String.concat "" (List.init 10 (fun _ -> "v2\n")) in
    let want =
      "v1\nhost 2\nv1\nafter1 4\n" ^ v2s ^ "after2 1204\n\
       v3\nafter3 218\nv3\nafter4 432\ncounter 216\n"
    in
    if code <> 0 || text <> want then
      fail "reload\n        got: %S (exit %d)\n     wanted: %S" text code want;

    Printf.printf
      "reload: emit %.1fms  llc %.1fms  ld %.1fms  (v2: emit %.1fms  llc %.1fms  ld %.1fms)  host run %.1fms\n"
      emit_ms t1.Build.llc_ms t1.Build.link_ms emit2_ms t2.Build.llc_ms
      t2.Build.link_ms dlopen_ms;
    print_string timings;

    List.iter (fun p -> try Sys.remove p with Sys_error _ -> ())
      [ host; so1; so2; so3; so4; out; tmp "err" ];
    if !failures = 0 then print_endline "reload: all tests passed"
    else begin
      Printf.printf "\n%d failure(s)\n" !failures;
      exit 1
    end
  | _ -> print_endline "reload: skipped (no clang or llc on PATH)"