(** What the indented printer needs to know of macros: which take a body run
    in order and at which argument it starts, and which names their
    expansions spell.

    A body is read off the 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, provided
    nothing else of it depends on how the body is split into arguments: out
    of its templates the rest parameter may only be counted against the
    body's start ([(< (length r) 1)]), read before the body ([(at r 0)] when
    the body starts at 1), or have the body's first form tested with a
    predicate ([(form-empty-list? (at r 0))]), since a [let] taking in the
    forms after it changes how many there are and nothing about the first.
    [comment] counts too: nothing in it runs. The printer lets a [let] in such
    a body take in the statements after it, as it does in a [do].

    The names a template spells outside its unquotes are what its expansion
    can refer to without its call spelling them: a [let] of one of those
    names is not given a longer scope over a call of that macro. *)

type t = {
  bodies : (string, int) Hashtbl.t;  (** called name -> argument the body starts at *)
  names : (string, string list) Hashtbl.t;  (** called name -> names its templates spell *)
}

let create () = { bodies = Hashtbl.create 32; names = Hashtbl.create 64 }

(* 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 body_start (t : t) h =
  match List.assoc_opt h core with Some k -> Some k | None -> Hashtbl.find_opt t.bodies h

(* The names a template spells outside its unquotes. *)
let rec template_names (f : Form.t) acc =
  match f.v with
  | Form.List ({ v = Form.Sym ("unquote" | "unquote-splicing"); _ } :: _) -> acc
  | Form.Sym s -> s :: acc
  | Form.List l | Form.Vec l | Form.Map l -> List.fold_left (fun a x -> template_names x a) acc l
  | _ -> acc

let rec templates (f : Form.t) acc =
  match f.v with
  | Form.List [ { v = Form.Sym "quasiquote"; _ }; x ] -> x :: acc
  | Form.List l | Form.Vec l | Form.Map l -> List.fold_left (fun a x -> templates x a) acc l
  | _ -> acc

(* The argument index the body of [(defmacro name [p ... & r] body ...)]
   starts at, or [None]. [known] gives another macro's. *)
let body_of ~(known : string -> int option) name ps body : int option =
  if name = "comment" then Some 0
  else
    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 is_r (f : Form.t) = f.v = Form.Sym r in
      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
      let int (f : Form.t) = match f.v with Form.Int i -> Some (Int64.to_int i) | _ -> None in
      let at_r (f : Form.t) =
        match f.v with
        | Form.List [ { v = Form.Sym "at"; _ }; x; i ] when is_r x -> int i
        | _ -> None
      in
      let bodies = ref [] and reads = ref [] and bad = ref false in
      (* Splices of [r] in a template, each where it lands. *)
      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"; _ }; x; k ] when is_r x -> int k
        | _ -> None
      in
      let lookup h = match List.assoc_opt h core with Some k -> Some k | None -> known h in
      let rec template (f : Form.t) =
        match f.v with
        | Form.List [ { v = Form.Sym "unquote"; _ }; e ] ->
          (match at_r e with
           | Some i -> reads := i :: !reads
           | None -> if mentions e then bad := true)
        | Form.List [ { v = Form.Sym "unquote-splicing"; _ }; e ] ->
          (* Anywhere but in a list's items: a vector, a map. *)
          if mentions e then bad := true
        | Form.List l ->
          let head = match l with { v = Form.Sym h; _ } :: _ -> Some h | _ -> None in
          (* A label after while, until or dotimes comes before the test. *)
          let label =
            match head, l with
            | Some ("while" | "until" | "dotimes"), _ :: { v = Form.Kw _; _ } :: _ -> 1
            | _ -> 0
          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 with
                  | Some k, Some b when i >= b + 1 + label -> bodies := k :: !bodies
                  | _ -> if mentions e then bad := true)
               | _ -> template x)
            l
        | Form.Vec l | Form.Map l ->
          List.iter
            (fun (x : Form.t) ->
               match x.v with
               | Form.List [ { v = Form.Sym "unquote-splicing"; _ }; e ] ->
                 if mentions e then bad := true
               | _ -> template x)
            l
        | _ -> ()
      in
      (* The macro's own code, out of its templates: [r] only counted, read
         before the body, or its first form tested. [k] is the body's start
         within [r], known once the templates are read. *)
      let rec code k (f : Form.t) =
        match f.v with
        | Form.List [ { v = Form.Sym "quasiquote"; _ }; x ] -> template x
        | Form.Sym s when s = r -> bad := true
        | Form.List [ { v = Form.Sym ("<" | ">=" | "=" | "<=" | ">"); _ }; a; b ] ->
          (match a.v, b.v with
           | Form.List [ { v = Form.Sym "length"; _ }; x ], _ when is_r x ->
             (match int b with Some c when c <= k + 1 -> () | _ -> code k b; bad := true)
           | _, Form.List [ { v = Form.Sym "length"; _ }; x ] when is_r x ->
             (match int a with Some c when c <= k + 1 -> () | _ -> code k a; bad := true)
           | _ -> code k a; code k b)
        | Form.List [ { v = Form.Sym p; _ }; e ]
          when String.length p > 1 && p.[String.length p - 1] = '?' && at_r e <> None ->
          (match at_r e with Some i when i <= k -> () | _ -> bad := true)
        | Form.List _ when at_r f <> None ->
          (match at_r f with Some i when i < k -> () | _ -> bad := true)
        | Form.List l | Form.Vec l | Form.Map l -> List.iter (code k) l
        | _ -> ()
      in
      (* The templates first, for [k]; then the rest of the code against it. *)
      List.iter (fun t -> template t) (List.fold_left (fun a x -> templates x a) [] body);
      match !bodies with
      | k :: more when (not !bad) && List.for_all (( = ) k) more
                       && List.for_all (fun i -> i < k) !reads ->
        List.iter (code k) body;
        if !bad then None else Some (fixed + k)
      | _ -> None

let add (t : t) ~qualify forms =
  let local = Hashtbl.create 8 in
  List.iter
    (fun (f : Form.t) ->
       match f.v with
       | Form.List ({ v = Form.Sym "defmacro"; _ } :: { v = Form.Sym name; _ }
                    :: { v = Form.Vec ps; _ } :: body) ->
         let known h =
           match Hashtbl.find_opt local h with
           | Some k -> Some k
           | None -> Hashtbl.find_opt t.bodies h
         in
         Hashtbl.replace t.names (qualify name)
           (List.fold_left (fun a x -> template_names x a) []
              (List.fold_left (fun a x -> templates x a) [] body));
         (match body_of ~known name ps body with
          | Some k -> Hashtbl.replace local name k; Hashtbl.replace t.bodies (qualify name) k
          (* A definition of the same name as a prelude macro replaces it. *)
          | None -> Hashtbl.remove local name; Hashtbl.remove t.bodies (qualify name))
       | _ -> ())
    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 t = create () in
  let prelude = try Reader.read_all ~file:Prelude.file Prelude.source with _ -> [] in
  add t ~qualify:Fun.id prelude;
  (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 t ~qualify:(fun n -> alias ^ "/" ^ n) fs
          with _ -> ())
       (Load.imports_of forms));
  add t ~qualify:Fun.id forms;
  t