(** Running a macro: the half of expansion that has to compile something.

    [Expand] is the image format, the quasiquote desugaring and the marshaller,
    and it depends on nothing above [Form]. This file is the part that cannot:
    expanding a macro means compiling it and dlopening it, so it needs [Check],
    [Build] and [Emit], and it therefore sits above the parser it feeds. The
    join is [Parse.expander], filled in at the bottom of this file. *)

(* ── Which names are macros ────────────────────────────────────────
   A [defmacro] is an [Ast.Defn] by the time [Parse] is finished with it, so
   the word only survives in the form and collecting them is a scan of the top
   level. It is the prelude's macros plus the file's, and not an imported
   package's: [Load] learns a package's imports by parsing it, so collecting
   from one would mean a second import resolver running over Forms. A defmacro
   in an imported package is refused by name instead. *)

let macro_name (f : Form.t) =
  match f.Form.v with
  | Form.List ({ Form.v = Form.Sym "defmacro"; _ }
               :: { Form.v = Form.Sym n; _ } :: _) -> Some n
  | _ -> None

let macros_in forms = List.filter_map macro_name forms

(* Does this form call one of these macros? A head position only, which is what
   a call is, and it is why the quasiquote desugaring has to have run first: a
   quasiquoted (cond ...) is a (Form.Sym {.s "cond"}) by now, and the name is a
   string in an argument rather than a head anything could mistake. *)
let rec names_macro (known : string list) (f : Form.t) =
  match f.Form.v with
  | Form.List ({ Form.v = Form.Sym n; _ } :: rest) ->
    List.mem n known || List.exists (names_macro known) rest
  | Form.List xs | Form.Vec xs | Form.Map xs ->
    List.exists (names_macro known) xs
  | _ -> false

(* ── The module ────────────────────────────────────────────────────
   The prelude plus the file's defmacros, and not the file's own functions.
   Compiling those would mean compiling a program that has not been expanded
   yet, which is the chicken and egg the pre-pass exists to avoid. The cost is
   that a macro body may call prelude functions and other macros and nothing
   else.

   Cached on disk under the object cache, keyed by a digest of exactly what
   goes into it. Every `flan build` is a fresh process, so without this the
   clang driver would be paid once per build of the same program instead of
   once per change to it. *)

type loaded = {
  handle : Dynload.handle;
  fns    : (string * Dynload.addr) list;
}

(* This compiler's own identity, and it belongs in the key for a reason the
   other caches do not have. A [.o] under the object cache is decided entirely
   by the C text and the C compiler that made it, so its key is total without
   naming flan at all. A macro module is not: it is *this* binary's codegen,
   dlopen'd back into *this* binary and called across a marshalled boundary.
   Change [Emit] or the runtime ABI and the .so on disk is wrong while the
   prelude text that keyed it has not moved — a stale macro expander, which
   fails as a crash inside [Expand.call] rather than as a compile error.

   It never showed because the cache sat under dune's per-run [TMPDIR] and so
   was empty on every run. Now that the cache outlives the run, the key has to
   carry what the directory used to hide.

   The stamp of the running binary is the identity, except in the one place
   where that binary is not a stable thing: a [flan dev] merged build lives at
   /tmp/flan-dev-<pid>/program, so its size-and-mtime is new on every start and
   keying on it would rebuild a macro module per session — measured at ~350ms
   of every dev start, which is most of what this cache exists to save. So
   [Dev.start_merged] passes its own stamp across the exec, and the merged
   binary uses the stamp of the compiler that built it, which is the one this
   key is actually about. *)
let self =
  lazy
    (match Sys.getenv_opt "FLAN_COMPILER_STAMP" with
     | Some s when s <> "" -> s
     | _ -> Build.stamp_of Sys.executable_name)

let key (extra : Form.t list) =
  Digest.to_hex
    (Digest.string
       (Lazy.force self ^ "\000" ^ Prelude.source ^ "\000"
        ^ String.concat "\000" (List.map Form.to_string extra)))

(* True while a macro module is being built. [Build.macro_module] goes through
   [Check.program], which parses the prelude, which calls back into
   [Parse.program] — and that would re-enter this and recurse forever. Nothing
   is lost by refusing to expand there: a macro compiled in round n calls only
   macros compiled in rounds before it, and those calls were already expanded
   before the build was entered. *)
let building = ref false

(* ── The bootstrap, and what a prelude macro may not call ───────────
   [Check.program] prepends the prelude to every program, this one included, so
   the module that expands the prelude's macros is compiled *from* the prelude.
   A prelude function that calls a macro therefore cannot be compiled into it:
   the call is a name nothing defines yet. That is a cycle and not an ordering
   mistake — no amount of moving the prepend around removes it.

   It is broken at one level, which is the restriction already recorded and
   kept: a macro module is built from the prelude with every [defn] that
   depends on a macro *removed*. Directly or transitively, because a function
   calling a dropped one is as unbuildable as the dropped one itself.

   Only [defn]s are dropped. A [defstruct], [defunion], [defalias], [defenum]
   or [defvar] stays whatever it names: the functions that survive still
   mention those types, and a reduced prelude missing them would not check.
   There used to be a sharper reason — [Parse.prelude_types] memoised the
   prelude's type names for the parser's return-type guess, and a reduced
   answer cached during a bootstrap build would have been wrong for every
   compile after it. That set is gone with the guess: a defn states its return
   type, so nothing in the parser asks what the prelude declares.

   A [defmacro] that lands in the dropped set is the violation of the rule, and
   it is refused here by name rather than reaching clang as an unknown symbol. *)

let head_name (f : Form.t) =
  match f.Form.v with
  | Form.List ({ Form.v = Form.Sym h; _ } :: { Form.v = Form.Sym n; _ } :: _) ->
    Some (h, n)
  | _ -> None

let reduce (forms : Form.t list) : Form.t list =
  let macros = macros_in forms in
  (* Fixpoint: a form is out once it names something already out. Bounded by
     the number of forms, since the set only grows. *)
  let out = ref macros in
  let changed = ref true in
  while !changed do
    changed := false;
    List.iter
      (fun f ->
         match head_name f with
         | Some (("defn" | "defmacro"), n) when not (List.mem n !out) ->
           if names_macro !out f then begin out := n :: !out; changed := true end
         | _ -> ())
      forms
  done;
  (* The macros themselves are in [out] by construction; a macro that is there
     for any *other* reason called one, which is the thing that cannot work. *)
  List.iter
    (fun f ->
       match head_name f with
       | Some ("defmacro", n) when names_macro macros f ->
         Loc.fail f.Form.loc
           "the prelude macro %s calls a macro, and a prelude macro may not: \
            the module that expands it is compiled from the prelude, so the \
            call would have to be expanded by a module that does not exist \
            yet. Call a function instead"
           n
       | _ -> ())
    forms;
  List.filter
    (fun f ->
       match head_name f with
       | Some ("defn", n) -> not (List.mem n !out)
       | _ -> true)
    forms

let compile (names : string list) (extra : Form.t list) : loaded =
  let out =
    Filename.concat (Build.cachedir ()) ("flan-macros-" ^ key extra ^ ".so")
  in
  if not (Sys.file_exists out) then begin
    building := true;
    Prelude.bootstrap := reduce;
    Fun.protect
      ~finally:(fun () ->
        building := false;
        Prelude.bootstrap := (fun fs -> fs))
      (fun () ->
         (* [Check.program] prepends the prelude itself — reduced, for the one
            build that cannot have all of it — so only the file's own defmacros
            go in here. *)
         let p = Check.program (Parse.program extra) in
         (* Written beside the final name and renamed, so a second process
            reading the cache never sees a half-written object. *)
         let tmp = out ^ "." ^ string_of_int (Unix.getpid ()) in
         ignore (Build.macro_module ~macros:names p ~out:tmp);
         (try Sys.rename tmp out with Sys_error _ -> ()))
  end;
  let handle = Dynload.dl_open out in
  { handle;
    fns = List.map (fun n -> (n, Dynload.dl_sym handle ("flan.macro." ^ n))) names }

(* ── The walk ──────────────────────────────────────────────────────
   Bottom up: a macro's arguments are expanded before it is called, so nothing
   a macro is handed contains a call to another macro. Then what it answers is
   expanded again, because a macro that expands into a call to itself — which
   is what a recursive [cond] is — has to keep going.

   That re-expansion is what needs a bound. [(defmacro loop [args] `(loop))]
   settles at nothing, and the honest answer to a macro that will not settle is
   to say which one it was, at the call site, rather than to run out of
   memory. *)

let fuel = 200

let rec expand_form (l : loaded) (f : Form.t) : Form.t =
  let loc = f.Form.loc in
  match f.Form.v with
  | Form.List ({ Form.v = Form.Sym n; _ } :: args) when List.mem_assoc n l.fns ->
    let args = List.map (expand_form l) args in
    (* The call site, tagged with the macro it is a call to. [Expand.unmarshal]
       stamps this onto every node the macro answers with, so from here down
       every form it produced knows where it came from and an error on one of
       them can say so. *)
    let from = Loc.from_macro n loc in
    settle l n loc (Expand.call ~loc:from (List.assoc n l.fns) args) fuel
  | Form.List xs -> Form.make (Form.List (List.map (expand_form l) xs)) loc
  | Form.Vec xs -> Form.make (Form.Vec (List.map (expand_form l) xs)) loc
  | Form.Map xs -> Form.make (Form.Map (List.map (expand_form l) xs)) loc
  | _ -> f

and settle l first loc (f : Form.t) left =
  match f.Form.v with
  | Form.List ({ Form.v = Form.Sym m; _ } :: args) when List.mem_assoc m l.fns ->
    if left <= 0 then
      Loc.fail loc
        "expanding %s did not settle after %d rounds — a macro that expands \
         into a call to a macro has to get smaller each time, and this one is \
         not"
        first fuel
    else begin
      let args = List.map (expand_form l) args in
      let from = Loc.from_macro m loc in
      settle l first loc (Expand.call ~loc:from (List.assoc m l.fns) args)
        (left - 1)
    end
  (* Settled at the head. The rest of it may still hold macro calls — a cond
     expands to an if whose else-branch is another cond — so the ordinary walk
     finishes the job. *)
  | _ -> expand_form l f

(* ── The rounds ────────────────────────────────────────────────────
   A macro's body may call a macro, so one sweep is not enough: a macro with an
   unexpanded call in its body cannot be compiled at all, because that call is
   a name nothing defines.

   So the module is built in rounds. Round 0 takes every macro whose body names
   no macro that is still waiting. Round 1 expands what is left against round
   0's module and takes whatever became clean. A round that takes nothing while
   macros remain is a cycle, and it is named rather than looped on.

   The prelude's own macros are in every round by construction — they are in
   every module this builds — so a prelude macro may not call a macro. It would
   fail to compile with an unknown name rather than with a reason, which is
   worth fixing the day the prelude wants one. *)

let rounds ~(prelude : string list) (pending : (string * Form.t) list)
  : (string * Form.t) list =
  let rec go ~taken ~pending =
    if pending = [] then taken
    else
      let waiting = List.map fst pending in
      let now, blocked =
        List.partition (fun (_, f) -> not (names_macro waiting f)) pending
      in
      if now = [] then
        Loc.fail (snd (List.hd pending)).Form.loc
          "these macros call each other and none can be compiled first: %s. A \
           defmacro has to be compiled before the call it expands, so a ring \
           has no order to be compiled in — one of them has to call a function \
           instead"
          (String.concat ", " waiting)
      else
        let taken = taken @ now in
        (* Nothing is waiting on this round, so there is nothing to expand it
           against and no module to build here. The common case is this one:
           every macro in the file is clean and round 0 is the only round. *)
        if blocked = [] then taken
        else begin
          let l = compile (prelude @ List.map fst taken) (List.map snd taken) in
          let blocked = List.map (fun (n, f) -> (n, expand_form l f)) blocked in
          Dynload.dl_close l.handle;
          Dynload.release ();
          go ~taken ~pending:blocked
        end
  in
  go ~taken:[] ~pending

(* ── The whole pass ────────────────────────────────────────────────── *)

(* Read once. The prelude is a constant string, and asking whether a file uses
   a macro would otherwise re-read the whole of it on every parse. *)
let prelude_macros = lazy (macros_in (Prelude.forms ()))

let program (forms : Form.t list) : Form.t list =
  if !building then forms
  else
    let prelude = Lazy.force prelude_macros in
    (* The prelude's own macros are dropped from [mine], and the reason is that
       these forms may *be* the prelude: [Check.program] prepends it, so a
       prelude macro handed back as [extra] would be declared twice and refused
       as a redefinition. They are already in [prelude], which is where the
       module gets them from. *)
    let mine =
      List.filter_map
        (fun f ->
           match macro_name f with
           | Some n when not (List.mem n prelude) -> Some (n, f)
           | _ -> None)
        forms
    in
    let all = prelude @ List.map fst mine in
    (* The common case by a wide margin, and the reason a build that uses no
       macro pays nothing: a file that calls none costs one scan and no
       compiler. Without it every build in the suite would link a macro module
       for the prelude's macros and pay a clang driver to answer nothing. *)
    if all = [] || not (List.exists (names_macro all) forms) then forms
    else begin
      let extra = rounds ~prelude mine in
      let l = compile all (List.map snd extra) in
      let out = List.map (expand_form l) forms in
      Dynload.dl_close l.handle;
      Dynload.release ();
      out
    end

let () = Parse.expander := program