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 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

(* ── 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
    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
    (* 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