(** 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;
}

let key (extra : Form.t list) =
  Digest.to_hex
    (Digest.string
       (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

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;
    Fun.protect
      ~finally:(fun () -> building := false)
      (fun () ->
         (* [Check.program] prepends the prelude itself, 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
    settle l n loc (Expand.call ~loc (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
      settle l first loc (Expand.call ~loc (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
    let mine = List.filter_map (fun f -> Option.map (fun n -> (n, f)) (macro_name f)) 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