flan/lib/expand.ml

(** Macro expansion: the pass between the reader and [Parse].

    There is no interpreter and there is not going to be one (docs/BUILT.md, "Why
    there is no interpreter"), so running a macro at compile time means
    compiling it and loading it into this process. Every piece of that is
    already built and measured — [Emit.macro_thunk], [Build.macro_module],
    [Dynload] — and this file is the two halves nobody had written: the image
    format the two sides share, and the walk that finds macro calls and
    replaces them.

    Expansion runs over [Form], before [Parse]. Not over [Ast]: [Parse] refuses
    [defmacro] outright and there is no [Ast.Defmacro], so an Ast-level pass
    would have nothing to work with. That refusal is the ordering. It is also
    Clojure's ordering, and it is why a macro expanding to a special form is
    ordinary here rather than a special case. *)

(* ── The image format ──────────────────────────────────────────────
   A Form is { i32 tag, [2 x i64] payload }: 24 bytes, align 8, payload at
   offset 8. Those three numbers are the whole agreement between this file and
   the compiled macro, and they are not taken on trust — test_acceptance.ml's
   "Form's image format" asks LLVM for each of them through the same ptrtoint
   oracle the DWARF offsets go through. Change the prelude's defunion and that
   test says which number moved.

   The tag is the case's position in the prelude's (defunion Form ...), which
   is why that list is a layout contract and says so. *)

let form_size = 24
let payload = 8

(* A string and a slice are both %slice = { ptr, i64 }: two words at the start
   of the payload. Every case of Form holds one member, so there is no third
   offset anywhere below. *)
let ptr_off = payload
let len_off = payload + 8

type tag =
  | TSym | TKw | TInt | TFloat | TStr | TByte | TList | TVec | TMap

let tag_int = function
  | TSym -> 0l | TKw -> 1l | TInt -> 2l | TFloat -> 3l | TStr -> 4l
  | TByte -> 5l | TList -> 6l | TVec -> 7l | TMap -> 8l

let tag_of_int = function
  | 0l -> TSym | 1l -> TKw | 2l -> TInt | 3l -> TFloat | 4l -> TStr
  | 5l -> TByte | 6l -> TList | 7l -> TVec | 8l -> TMap
  | n ->
    failwith
      (Printf.sprintf
         "a macro returned a Form with tag %ld, and Form has nine cases. The \
          prelude's (defunion Form ...) and lib/expand.ml's tag list are one \
          contract and have come apart"
         n)

(* ── Writing a Form into memory a macro can read ───────────────────
   OCaml cannot address raw memory, so this goes through the poke family in
   dynload_stubs.c, one field at a time. Everything allocated here is owned by
   [Dynload] and released together after the call. *)

let rec marshal (f : Form.t) : Dynload.addr =
  let p = Dynload.take form_size in
  write p f;
  p

(* Into an existing 24 bytes, which is what an argument array needs: the macro
   takes a [Form] slice, and a slice is contiguous elements and not an array of
   pointers. *)
and write p (f : Form.t) =
  let tag t = Dynload.poke_i32 p 0 (tag_int t) in
  let str t s =
    tag t;
    let n = String.length s in
    (* A zero-length string still gets a pointer, because a slice with a null
       base is not the same value as one with a live base and a zero length --
       the difference shows the day something concatenates onto it. *)
    let b = Dynload.take (max n 1) in
    if n > 0 then Dynload.poke_bytes b 0 s;
    Dynload.poke_ptr p ptr_off b;
    Dynload.poke_i64 p len_off (Int64.of_int n)
  in
  let seq t xs =
    tag t;
    let n = List.length xs in
    let b = Dynload.take (max (n * form_size) 1) in
    List.iteri (fun i x -> write (Nativeint.add b (Nativeint.of_int (i * form_size))) x) xs;
    Dynload.poke_ptr p ptr_off b;
    Dynload.poke_i64 p len_off (Int64.of_int n)
  in
  match f.Form.v with
  | Form.Sym s -> str TSym s
  | Form.Kw s -> str TKw s
  | Form.Str s -> str TStr s
  | Form.Int i -> tag TInt; Dynload.poke_i64 p payload i
  | Form.Float x -> tag TFloat; Dynload.poke_f64 p payload x
  | Form.Byte b -> tag TByte; Dynload.poke_i32 p payload (Int32.of_int b)
  | Form.List xs -> seq TList xs
  | Form.Vec xs -> seq TVec xs
  | Form.Map xs -> seq TMap xs

(* ── Reading one back ──────────────────────────────────────────────
   [loc] is the call site's, stamped onto every node. A macro cannot invent a
   source location and the image has no room for one: Form on the Flan side
   mirrors [Form.value], not [Form.t]. So an error inside an expansion points
   at the call that produced it, which is the part of "the error carries the
   expansion" that can be had now without the structured-error rewrite. *)

let rec unmarshal ~loc (p : Dynload.addr) : Form.t =
  let str () =
    let b = Dynload.peek_ptr p ptr_off in
    let n = Int64.to_int (Dynload.peek_i64 p len_off) in
    if n = 0 then "" else Dynload.peek_bytes b 0 n
  in
  let seq () =
    let b = Dynload.peek_ptr p ptr_off in
    let n = Int64.to_int (Dynload.peek_i64 p len_off) in
    List.init n (fun i ->
      unmarshal ~loc (Nativeint.add b (Nativeint.of_int (i * form_size))))
  in
  let v =
    match tag_of_int (Dynload.peek_i32 p 0) with
    | TSym -> Form.Sym (str ())
    | TKw -> Form.Kw (str ())
    | TStr -> Form.Str (str ())
    | TInt -> Form.Int (Dynload.peek_i64 p payload)
    | TFloat -> Form.Float (Dynload.peek_f64 p payload)
    | TByte -> Form.Byte (Int32.to_int (Dynload.peek_i32 p payload) land 0xff)
    | TList -> Form.List (seq ())
    | TVec -> Form.Vec (seq ())
    | TMap -> Form.Map (seq ())
  in
  Form.make v loc

(* ── One call ──────────────────────────────────────────────────────
   The arguments are one contiguous run of Forms, not an array of pointers,
   because the macro's parameter is [[Form]] and a Flan slice is { ptr, len }
   over elements. *)

let call ~loc (fn : Dynload.addr) (args : Form.t list) : Form.t =
  let n = List.length args in
  let a = Dynload.take (max (n * form_size) 1) in
  List.iteri
    (fun i x -> write (Nativeint.add a (Nativeint.of_int (i * form_size))) x)
    args;
  let out = Dynload.take form_size in
  Dynload.call fn a (Int64.of_int n) out;
  unmarshal ~loc out

(* ── Quasiquote ────────────────────────────────────────────────────
   A desugaring over [Form], and nothing more: a quasiquoted (if ~t ~b) becomes
   calls to the prelude's form-building surface, which the checker then sees as
   ordinary code. There is no quasiquote left in the language after this runs,
   which is why the expander's own walk needs no idea that quoting exists: by
   the time it looks for macro calls, a [cond] written inside a quasiquote is a
   (Form.Sym {.s "cond"}) and there is no head there to mistake for a call the
   compiler should make now.

   The reader stays dumb and produces (quasiquote x), (unquote x) and
   (unquote-splicing x) with no idea whether one is inside another. Counting
   levels is this file's job, and it does not: a quasiquote inside a quasiquote
   is refused by name. A macro that writes a macro is the only thing that wants
   one, nothing in the corpus does, and CL's level arithmetic has a real cost
   that no use case has asked for. *)

let sym loc s = Form.make (Form.Sym s) loc
let lst loc xs = Form.make (Form.List xs) loc

(* (Form.Case {.field value}) — a node of the image, written as the Flan
   constructor the prelude declares. *)
let node loc case field v =
  lst loc [ sym loc ("Form." ^ case);
            Form.make (Form.Map [ sym loc ("." ^ field); Form.make v loc ]) loc ]

let unquote_of (f : Form.t) =
  match f.Form.v with
  | Form.List [ { Form.v = Form.Sym "unquote"; _ }; x ] -> Some x
  | _ -> None

let splice_of (f : Form.t) =
  match f.Form.v with
  | Form.List [ { Form.v = Form.Sym "unquote-splicing"; _ }; x ] -> Some x
  | _ -> None

let rec quote (f : Form.t) : Form.t =
  let loc = f.Form.loc in
  match unquote_of f with
  (* The escape: whatever the program wrote, evaluated. It is already a Form,
     because a Form is what a macro body deals in. *)
  | Some x -> x
  | None ->
    match splice_of f with
    | Some _ ->
      Loc.fail loc
        "~@x splices into a list or a vector, and there is nothing here for it \
         to splice into"
    | None ->
      match f.Form.v with
      | Form.List ({ Form.v = Form.Sym "quasiquote"; _ } :: _) ->
        Loc.fail loc
          "a quasiquote inside a quasiquote is not implemented: the reader does \
           not count nesting levels and neither does this, so the inner one has \
           no meaning to give. Build the inner form with form-cons"
      | Form.Sym s -> node loc "Sym" "s" (Form.Str s)
      | Form.Kw s -> node loc "Kw" "s" (Form.Str s)
      | Form.Int i -> node loc "Int" "i" (Form.Int i)
      | Form.Float x -> node loc "Float" "x" (Form.Float x)
      | Form.Str s -> node loc "Str" "s" (Form.Str s)
      | Form.Byte b -> node loc "Byte" "b" (Form.Int (Int64.of_int b))
      | Form.List xs -> node loc "List" "xs" (seq loc xs).Form.v
      | Form.Vec xs -> node loc "Vec" "xs" (seq loc xs).Form.v
      | Form.Map xs -> node loc "Map" "xs" (seq loc xs).Form.v

(* The [Form] slice one bracket's worth of items comes to. Built right to left,
   so each item is consed onto what follows it and a splice is an append — the
   three prelude functions and no fourth. *)
and seq loc items =
  List.fold_left
    (fun acc (item : Form.t) ->
       match splice_of item with
       | Some x -> lst item.Form.loc [ sym item.Form.loc "form-append"; x; acc ]
       | None -> lst item.Form.loc [ sym item.Form.loc "form-cons"; quote item; acc ])
    (lst loc [ sym loc "form-nil" ])
    (List.rev items)

(* Every quasiquote in a form, outermost first. Pure, total, and dependent on
   nothing but Form, which is what lets [Parse] run it on the way in rather
   than needing the whole expander wired up first. *)
let rec quasiquote (f : Form.t) : Form.t =
  match f.Form.v with
  | Form.List [ { Form.v = Form.Sym "quasiquote"; _ }; x ] -> quote x
  | Form.List xs -> Form.make (Form.List (List.map quasiquote xs)) f.Form.loc
  | Form.Vec xs -> Form.make (Form.Vec (List.map quasiquote xs)) f.Form.loc
  | Form.Map xs -> Form.make (Form.Map (List.map quasiquote xs)) f.Form.loc
  | _ -> f