(** Which macros take a body run in order, and at which argument it starts.

    Read off each macro's definition: a macro whose rest parameter is spliced,
    whole or from a fixed index on, only into places whose forms run in order
    — a [do], the body of a [let], [fn], [when], [while] or [loop], or the body
    of another such macro — takes a body there. [comment] counts too: nothing
    in it runs. The indented printer lets a [let] in such a body take in the
    statements after it, as it does in a [do].

    What a macro does with its arguments outside its templates (a guard that
    counts them, say) is not looked at. *)

type t = (string, int) Hashtbl.t

(* Core forms whose trailing arguments are a body run in order, and how many
   arguments come before it. *)
let core = [ ("do", 0); ("let", 1); ("fn", 1); ("when", 1); ("while", 1); ("loop", 1);
             ("defer", 0); ("with-allocator", 1) ]

let lookup (known : string -> int option) h =
  match List.assoc_opt h core with Some k -> Some k | None -> known h

(* The argument index the body starts at, from [(defmacro name [p ... & r]
   body ...)]; [None] when it does not take one. *)
let of_defmacro ~(known : string -> int option) (f : Form.t) : (string * int) option =
  match f.v with
  | Form.List ({ v = Form.Sym "defmacro"; _ } :: { v = Form.Sym "comment"; _ } :: _) ->
    Some ("comment", 0)
  | Form.List ({ v = Form.Sym "defmacro"; _ } :: { v = Form.Sym name; _ }
               :: { v = Form.Vec ps; _ } :: body) ->
    let rec split fixed = function
      | { Form.v = Form.Sym "&"; _ } :: [ { Form.v = Form.Sym r; _ } ] -> Some (fixed, r)
      | { Form.v = Form.Sym "&"; _ } :: _ -> None
      | _ :: rest -> split (fixed + 1) rest
      | [] -> None
    in
    (match split 0 ps with
     | None -> None
     | Some (fixed, r) ->
       let rec mentions (f : Form.t) =
         match f.v with
         | Form.Sym s -> s = r
         | Form.List l | Form.Vec l | Form.Map l -> List.exists mentions l
         | _ -> false
       in
       (* Each splice of the rest parameter: [Some k] when it splices from
          argument [k] of the rest on into a place run in order. *)
       let starts = ref [] and bad = ref false in
       let splice_from (e : Form.t) =
         match e.v with
         | Form.Sym s when s = r -> Some 0
         | Form.List [ { v = Form.Sym "form-rest"; _ }; { v = Form.Sym s; _ };
                       { v = Form.Int k; _ } ] when s = r -> Some (Int64.to_int k)
         | _ -> None
       in
       let rec template (f : Form.t) =
         match f.v with
         | Form.List ({ v = Form.Sym "unquote"; _ } :: e) ->
           (* [~(at r i)] reads one argument, which is fine before the body
              and not in it; the check is made once the body's start is
              known. Anything else of [r] under an unquote is not followed. *)
           List.iter
             (fun (e : Form.t) ->
                match e.v with
                | Form.List [ { v = Form.Sym "at"; _ }; { v = Form.Sym s; _ };
                              { v = Form.Int i; _ } ] when s = r ->
                  starts := `Read (Int64.to_int i) :: !starts
                | _ -> if mentions e then bad := true)
             e
         | Form.List l ->
           let head = match l with { v = Form.Sym h; _ } :: _ -> Some h | _ -> None in
           List.iteri
             (fun i (x : Form.t) ->
                match x.v with
                | Form.List [ { v = Form.Sym "unquote-splicing"; _ }; e ] ->
                  (match splice_from e, Option.bind head (lookup known) with
                   | Some k, Some b when i >= b + 1 -> starts := `Body k :: !starts
                   | Some _, _ -> bad := true
                   | None, _ -> if mentions e then bad := true)
                | _ -> template x)
             l
         | Form.Vec l | Form.Map l -> List.iter template l
         | _ -> ()
       in
       let rec code (f : Form.t) =
         match f.v with
         | Form.List [ { v = Form.Sym "quasiquote"; _ }; x ] -> template x
         | Form.List l | Form.Vec l | Form.Map l -> List.iter code l
         | _ -> ()
       in
       List.iter code body;
       let bodies = List.filter_map (function `Body k -> Some k | _ -> None) !starts in
       let reads = List.filter_map (function `Read i -> Some i | _ -> None) !starts in
       match bodies with
       | k :: more when (not !bad) && List.for_all (( = ) k) more
                        && List.for_all (fun i -> i < k) reads ->
         Some (name, fixed + k)
       | _ -> None)
  | _ -> None

let add (tbl : t) ~qualify ~known forms =
  let local = Hashtbl.create 8 in
  List.iter
    (fun f ->
       match of_defmacro ~known:(fun h ->
           match Hashtbl.find_opt local h with Some k -> Some k | None -> known h) f with
       | Some (name, k) -> Hashtbl.replace local name k; Hashtbl.replace tbl (qualify name) k
       | None -> ())
    forms

(** The prelude's macros, those of the packages [forms] imports (qualified
    by their alias) and [forms]' own. An import that cannot be found or read
    adds nothing. *)
let table ?file (forms : Form.t list) : t =
  let tbl = Hashtbl.create 32 in
  let prelude = try Reader.read_all ~file:Prelude.file Prelude.source with _ -> [] in
  add tbl ~qualify:Fun.id ~known:(fun _ -> None) prelude;
  let known h = Hashtbl.find_opt tbl h in
  (match file with
   | None -> ()
   | Some file ->
     List.iter
       (fun (alias, path, loc) ->
          try
            let d = Load.resolve_dir ~file loc path in
            let files = if Sys.is_directory d then Load.source_entries d else [ d ] in
            let fs = List.concat_map Source.read_file files in
            add tbl ~qualify:(fun n -> alias ^ "/" ^ n) ~known fs
          with _ -> ())
       (Load.imports_of forms));
  add tbl ~qualify:Fun.id ~known forms;
  tbl