flan/lib/macro.ml

(** 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, the file's own, and an imported
   package's.

   The third of those does not arrive by scanning anything here. [Load] is the
   one thing that knows what a package is called from outside, and it now hands
   its answer over in [Parse.imported_macros] — qualified under the alias and
   desugared — so this file's only new job is to put that set in the same pot
   as the other two. See the header of [Load.import]. *)

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], [defdata], [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 ()))

(* The module the forms below will be expanded against, or [None] when there is
   nothing to expand them with.

   Split out of [program] rather than copied into the editor's path, because
   everything in it is a decision with a paragraph attached — which names are
   ambient, which shadow which, and the shortcut that keeps a macro-free build
   free of a clang driver. Two copies of that would drift, and the second copy
   is the one a reader would not know to distrust. [program] below is the whole
   of what used to be here; [expand_step] and [expand_all] are the editor's. *)
let loaded_for (forms : Form.t list) : loaded option =
  if !building then None
  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
    (* An import's macros, and they go in beside [mine] rather than beside
       [prelude]: a package macro may call another macro of its own package, so
       it is exactly as much a candidate for the rounds below as one written
       here.

       This used to say that an import's names carry a slash, so nothing
       ambient could collide with a name written here. That stopped being true
       when a session started holding the buffer's own [defmacro]s — see
       [Session.eval] — and it stopped being true on the most ordinary action
       there is: [C-c C-c] over a [defmacro] the session already knows sends a
       form declaring a name the ambient set also has. Two forms declaring one
       name reach [Check.program] as a duplicate declaration, refused with a
       sentence nobody would connect to this.

       So the merge dedupes, and the direction is the one [Load.macro_union]
       already uses a level up: a name declared in the forms being parsed
       shadows the ambient copy. That is also what makes an *edited* macro
       expand with its new body rather than with the session's stale one. *)
    let imported =
      List.filter_map
        (fun f -> Option.map (fun n -> (n, f)) (macro_name f))
        !Parse.imported_macros
    in
    let imported =
      List.filter (fun (n, _) -> not (List.mem_assoc n mine)) imported
    in
    let mine = imported @ mine 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 None
    else begin
      let extra = rounds ~prelude mine in
      Some (compile all (List.map snd extra))
    end

(* Closing the handle and freeing the marshaller's scratch, on the way out of
   an expansion however it ends.

   [Fun.protect] and not two statements after the call, which is what this was
   while a build was the only caller: a build that raised was a compile that
   failed and a process about to exit, so a leaked handle and a few unfreed
   argument buffers cost nothing anybody could measure. The daemon is not that
   process. A macro that does not settle raises out of here, the editor is told
   and stays connected, and the next [C-c C-m] does it again — so the one thing
   [Dynload.owned] must not become is a list that only ever grows over a
   session's lifetime. *)
let with_module (l : loaded) (f : unit -> 'a) : 'a =
  Fun.protect
    ~finally:(fun () ->
      Dynload.dl_close l.handle;
      Dynload.release ())
    f

let program (forms : Form.t list) : Form.t list =
  match loaded_for forms with
  | None -> forms
  | Some l -> with_module l (fun () -> List.map (expand_form l) forms)

let () = Parse.expander := program

(* ── What the editor asks ──────────────────────────────────────────
   [C-c C-m]. Two questions and not one, because a macro that quasiquotes a
   call to another macro makes the difference real: [mac/quad] answers
   [(mac/twice (mac/twice n))], and the fixpoint of that says nothing about
   which macro produced what.

   Both answer the name at the head when it is a macro, so the editor can say
   *which* macro it just ran rather than only that something changed. That name
   is the one thing the printed text cannot carry: [Loc.from_macro] is
   outermost-wins, so every node of a full expansion is stamped with the macro
   the author wrote and the intermediate names are gone by the time it settles.

   One step is outermost-only, and that is a deliberate difference from
   [expand_form], which expands a call's arguments *before* calling it. So
   [(mac/twice (mac/twice 1))] one-stepped here is [(+ (mac/twice 1)
   (mac/twice 1))], where the compiler's own first move is [(mac/twice (+ 1
   1))]. Different intermediates, the same fixpoint. One step is a view of what
   this macro did; all the way is the answer the compiler acts on. *)

let head_macro (l : loaded) (f : Form.t) : string option =
  match f.Form.v with
  | Form.List ({ Form.v = Form.Sym n; _ } :: _) when List.mem_assoc n l.fns ->
    Some n
  | _ -> None

(** One round of the outermost call, or the form unchanged when its head does
    not name a macro. Nothing here needs the fuel [settle] carries: one call is
    one call, and what it answers is not looked at again. *)
let expand_step (f : Form.t) : Form.t * string option =
  match loaded_for [ f ] with
  | None -> (f, None)
  | Some l ->
    with_module l (fun () ->
        match f.Form.v with
        | Form.List ({ Form.v = Form.Sym n; _ } :: args)
          when List.mem_assoc n l.fns ->
          ( Expand.call ~loc:(Loc.from_macro n f.Form.loc) (List.assoc n l.fns)
              args,
            Some n )
        | _ -> (f, None))

(** To the fixpoint, through exactly the walk a build goes through — so the
    text this answers is the text the checker is about to be handed, and the
    refusals are the build's own. A macro that does not settle raises
    [Loc.Error] out of [settle] at the bound, and a ring was refused a level up
    in [rounds] before anything was compiled at all. *)
let expand_all (f : Form.t) : Form.t * string option =
  match loaded_for [ f ] with
  | None -> (f, None)
  | Some l ->
    with_module l (fun () ->
        (* The name is read off the *input*, so it has to be taken before the
           walk replaces it. Bound rather than written as a tuple: OCaml does
           not promise the order a tuple's components are evaluated in, and
           [Dev.serve] already carries a comment about the one place that bit.
           Here it would be silent — the wrong macro named, never an error. *)
        let name = head_macro l f in
        (expand_form l f, name))