(******************************************************************************)
(*                                                                            *)
(*                                    Menhir                                  *)
(*                                                                            *)
(*   Copyright Inria. All rights reserved. This file is distributed under     *)
(*   the terms of the GNU General Public License version 2, as described in   *)
(*   the file LICENSE.                                                        *)
(*                                                                            *)
(******************************************************************************)

open Range
open Located
open Attribute
open Syntax
open Sort

(* When a parameterized symbol is expanded (specialized), if it carries a
   %merge function, then this function is copied. This can cause some code
   duplication. If necessary, one could try to avoid this duplication, but the
   current approach is simpler. (Currently, type inference is performed after
   duplication. There is no requirement for the copies to share a common type
   or even a common type scheme.) *)

(* -------------------------------------------------------------------------- *)

(* Expansion modes. *)

type mode =
  | ExpandHigherSort
  | ExpandAll

(* -------------------------------------------------------------------------- *)

(* Expansion can be understood as traversing a graph where every vertex is
   labeled with a pair of a nonterminal symbol [nt] and an instantiation of
   the formal parameters of [nt]. *)

(* We allow partial instantiations, where some of the formal parameters of
   [nt] are instantiated, while others remain present. For this reason, we
   represent an instantation as a list of *optional* actual parameters. *)

(* The actual parameters that appear in an instantiation make sense *in the
   source namespace* (at the toplevel). That is, they refer to (terminal and
   nonterminal) symbols that exist (at the toplevel) in the original
   grammar. *)

type instantiation =
  parameter option list

type label =
  nonterminal * instantiation

(* Equality and hashing for labels. *)

module Label = struct
  type t = label
  let equal (nt1, inst1) (nt2, inst2) =
    nt1 = nt2 &&
    List.for_all2 (Option.equal Parameter.equal) inst1 inst2
  let hash (nt, inst) =
    Hashtbl.hash (nt, MList.hash (MOption.hash Parameter.hash) inst)
end

(* -------------------------------------------------------------------------- *)

(* [mangle label] chooses a concrete name for the new nonterminal symbol that
   corresponds to the label [label]. *)

(* We include parentheses and commas in this name, because that is readable
   and acceptable in many situations. We replace them with underscores in
   situations where these characters are not valid; see [MString.normalize]. *)

let mangle_po (po : parameter option) =
  match po with
  | None ->
      (* When a parameter remains uninstantiated,
         we put an underscore in its place. *)
      "_"
  | Some p ->
      Parameter.print "," p

let mangle_inst (pos : instantiation) =
  MString.separated_list mangle_po "," pos

let mangle ((nt, pos) : label) : nonterminal =
  if pos = [] then nt else
    Printf.sprintf "%s(%s)" nt (mangle_inst pos)

(* -------------------------------------------------------------------------- *)

(* As a general policy, the attributes associated with a branch (that is, a
   production) are not affected when this production is specialized; they are
   just transmitted to every specialized copy. *)

(* As an exception to this general policy, the [@name] attribute (if there is
   one) is specialized, in the hope of maintaining the property that no two
   productions have the same name. The name is transformed as follows: to the
   original name, we append the instantiation, where each actual parameter
   is preceded with an underscore character. *)

let normalize_mangle_po po =
  MString.normalize (mangle_po po)

let subst_name_attribute inst (name : string) : string =
  name ^ MString.preceded_list normalize_mangle_po "_" inst

let subst_attributes inst (attrs : attributes) : attributes =
  transform_attribute "name" (subst_name_attribute inst) attrs

(* -------------------------------------------------------------------------- *)

(* An environment maps all of the formal parameters of a rule to actual
   parameters, which make sense in the source namespace. *)

module Env =
  StringMap

type env =
  parameter Env.t

let subst_symbol env sym : parameter =
  try
    Env.find (value sym) env
  with Not_found ->
    (* [x] is not a formal parameter. It is a toplevel symbol. *)
    ParamVar sym

let subst_parameter env param : parameter =
  Parameter.subst (subst_symbol env) param

(* -------------------------------------------------------------------------- *)
(* -------------------------------------------------------------------------- *)

(* For syntactic convenience, the rest of this file is a functor. *)

module Run (G : sig
  (* Expansion mode. *)
  val mode: mode
  (* Sort information. *)
  val sorts: sorts
  (* The grammar [g] whose expansion is desired. *)
  val g : grammar
end) = struct

open G

(* -------------------------------------------------------------------------- *)

(* Determining the sort of a symbol or parameter. *)

(* Be careful: these functions use the toplevel sort environment [sorts],
   so they must not be used within a rule. (The sort environment would have
   to be extended with information about the formal parameters.) *)

let symbol_sort symbol =
  try
    StringMap.find (value symbol) sorts
  with Not_found ->
    assert false

let sort param =
  match param with
  | ParamVar sym ->
      symbol_sort sym
  | ParamApp (_, params) ->
      assert (params <> []);
      (* An application always has sort [*]. *)
      star
  | ParamAnonymous _ ->
      assert false

(* -------------------------------------------------------------------------- *)

(* Looking up the [%attribute] declarations, looking for attributes attached
   with a nonterminal symbol [nt]. This is used when we create a specialized
   version of this symbol. *)

(* We use an inefficient linear search, but that shouldn't be a problem. *)

let global_attributes (nt : symbol) : attribute list =
  let param = Parameter.var nt in
  List.concat (List.map (fun (params, attrs) ->
    if List.exists (Parameter.equal param) params then attrs else []
  ) g.p_symbol_attributes)

(* -------------------------------------------------------------------------- *)

(* A queue keeps track of the graph vertices that have been discovered but
   not yet visited. *)

let enqueue, repeatedly =
  let queue = Queue.create() in
  let enqueue label =
    Queue.add label queue
  and repeatedly visit =
    MQueue.repeat queue visit
  in
  enqueue, repeatedly

(* -------------------------------------------------------------------------- *)

(* A hash table is used to mark the graph vertices that have been discovered. *)

let mark, marked =
  let module H = Hashtbl.Make(Label) in
  let table = H.create 128 in
  let mark label =
    H.add table label ()
  and marked label =
    H.mem table label
  in
  mark, marked

(* -------------------------------------------------------------------------- *)

(* The rules of the expanded grammar are gradually collected. *)

let emit, rules =
  let rules = ref StringMap.empty in
  let emit rule =
    assert (not (StringMap.mem rule.pr_nt !rules));
    rules := StringMap.add rule.pr_nt rule !rules
  and rules() =
    !rules
  in
  emit, rules

(* -------------------------------------------------------------------------- *)

(* On top of the function [mangle], we set up a mechanism that checks that
   every (normalized) mangled name is unique. (Indeed, in principle, there
   could be clashes, although in practice this is unlikely.) We must check
   that every application of [mangle] to a *new* argument yields a *new*
   (normalized) result. This is succinctly expressed by combining a claim
   and a memoizer. *)

(* [fail name] causes a fatal failure. This failure has never been observed
   in practice. *)

let fail name =
  Printf.eprintf "Error: internal name clash over %s.\n" name;
  assert false

(* [new_claim()] creates a new service for claiming names. It returns a
   function [claim] of type [int -> unit] such that the call [claim x]
   succeeds if and only if [claim x] has never been called before. *)

let new_claim () : (string -> unit) =
  let names = ref StringSet.empty in
  let claim name =
    if StringSet.mem name !names then
      fail name;
    names := StringSet.add name !names
  in
  claim

let mangle : label -> nonterminal =
  let ensure_fresh = new_claim() in
  let module M = Fix.Memoize.ForHashedType(Label) in
  M.memoize (fun label ->
    let name = mangle label in
    ensure_fresh (MString.normalize name);
    name
  )

(* -------------------------------------------------------------------------- *)

(* [recognize] receives an actual parameter [param] that makes sense in the
   source namespace and transforms it into a parameter that makes sense in the
   target namespace. This involves examining each application and
   "recognizing" it as an application of a label to a sequence of residual
   actual parameters, as explained next. All labels thus recognized are
   enqueued. *)

(* [recognize] governs how much specialization is performed. For instance,
   [F(X, Y, Z)] could be recognized as:

   - an application of the symbol [F] to the residual arguments [X, Y, Z].
     Then, no specialization at all takes place.

   - an application of the symbol [F(X,Y,Z)] to no residual arguments.
     Then, [F] is fully specialized for [X, Y, Z].

   - in between these extremes, say,
     an application of the symbol [F(X,_,Z)] to the residual argument [Y].
     Then, [F] is partially specialized.

   If there are any residual arguments, then they must be recursively
   recognized. For instance, [F(X,G(Y),Z)] could be recognized as an
   application of the symbol [F(X,_,Z)] to [G(Y)], which itself could
   be recognized as an application of the symbol [G(Y)] to no residual
   arguments. *)

let rec recognize (param : parameter) : parameter =
  (* [param] must have sort [star], in an appropriate sort environment. *)
  match param with
  | ParamAnonymous _ ->
      assert false
  | ParamVar _ ->
      param
  | ParamApp (sym, ps) ->
      assert (ps <> []);
      let x = value sym in
      (* This symbol is applied to at least one argument, so cannot be
         a terminal symbol. It must be either a nonterminal symbol or
         an (uninstantiated) formal parameter of the current rule. *)
      (* Actually, in both modes, formal parameters of higher sort are
         expanded away, so [sym] cannot be an uninstantiated parameter
         of the current rule. It must be a nonterminal symbol. We can
         therefore look up its sort in the toplevel environment [sorts]. *)
      let (inst : instantiation), (residuals : parameters) =
        match mode with
        | ExpandAll ->
            (* Expansion of all parameters. *)
            let inst = List.map (fun p -> Some p) ps
            and residuals = [] in
            inst, residuals
        | ExpandHigherSort ->
            (* Expansion of only the parameters of higher sort. *)
            let ss : sort list = domain (symbol_sort sym) in
            assert (List.length ps = List.length ss);
            let pss = List.combine ps ss in
            let inst =
              pss
              |> List.map (fun (param, sort) ->
                   if sort = star then None else Some param)
            in
            let residuals =
              pss
              |> List.filter (fun (_, sort) -> sort = star)
              |> List.map (fun (param, _) -> recognize param)
            in
            inst, residuals
      in
      let label : label = (x, inst) in
      enqueue label;
      let sym = mangle label in
      Parameter.apply (locate dummy sym) residuals

(* -------------------------------------------------------------------------- *)

(* The following functions take arguments in the source namespace and produce
   results in the target namespace. *)

(* The environment [env] alone is normally sufficient to determine the
   substitution that must be performed. However, [subst_attributes]
   requires the instantiation [inst] in order to transform [@name]
   attributes. [inst] contains roughly slightly more information than
   [env], as it retains the order of the formal parameters. *)

let subst_parameter env param =
  (* [param] must have sort [star], in an appropriate sort environment. *)
  recognize (subst_parameter env param)

let subst_producer env (id, param, attrs) =
  let param = subst_parameter env param in
  (id, param, attrs)

let subst_producers env producers =
  List.map (subst_producer env) producers

let subst_branch inst env branch =
  { branch with
    pb_producers = subst_producers env branch.pb_producers;
    pb_attributes = subst_attributes inst branch.pb_attributes }

let subst_branches inst env branches =
  List.map (subst_branch inst env) branches

(* -------------------------------------------------------------------------- *)

(* A quick and dirty way of mapping a name to a fresh name. *)

let freshen : string -> string =
  let c = ref 0 in
  fun x ->
    Printf.sprintf "%s__menhir__%d" x (MInt.postincrement c)

(* -------------------------------------------------------------------------- *)

(* [instantiation_env] expects the formal parameters of a rule, [formals], and
   an instantiation [inst] that dictates how this rule must be specialized. It
   returns an environment [env] that can be used to perform specialization and
   a list of residual formal parameters (those that are not specialized). *)

let instantiation_env (formals : symbols) (inst : instantiation) : env * symbols =
  assert (List.length formals = List.length inst);
  List.fold_right2 (fun formal po (env, residuals) ->
    let param, residuals =
      match po with
      | Some param ->
          (* This formal parameter is instantiated. *)
          param, residuals
      | None ->
          (* This formal parameter is not instantiated. *)
          (* We would like to map it to itself. *)
          (* However, we must in principle be a bit careful: if a toplevel
             symbol by the same name as [formal] appears free in the codomain
             of the environment that we are building, then we will run intro
             trouble. We avoid this problem by systematically renaming every
             formal parameter to a fresh unlikely name. *)
          let formal = freshen formal in
          Parameter.var formal,
          formal :: residuals
    in
    Env.add formal param env, residuals
  ) formals inst (Env.empty, [])

(* -------------------------------------------------------------------------- *)

(* [visit label] visits a vertex labeled [label] in the graph. This label is a
   pair of a nonterminal symbol [nt] and an instantiation [inst]. Unless this
   vertex has been visited already, we create a specialized copy of [nt] for
   this instantiation. This involves a call to [subst_branches], which can
   cause more vertices to be discovered and enqueued. *)

(* The specialized symbol retains any attributes carried by the original
   parameterized symbol. These attributes could be either attached with this
   rule ([rule.pr_attributes]) or specified via an [%attribute] declaration.
   We have to look up [%attribute] declarations now (as opposed to letting
   [Drop] handle them) if this is a parameterized symbol, as the connection
   between the original parameterized symbol and its specialized version is
   evident here but is lost afterwards. *)

let visit label =
  if not (marked label) then
    let () = mark label in
    let (nt, inst) = label in
    let rule = StringMap.find nt g.p_rules in
    let formals = rule.pr_parameters in
    let env, residuals = instantiation_env formals inst in
    emit {
      rule with
      pr_nt = mangle label;
      pr_parameters = residuals;
      pr_branches = subst_branches inst env rule.pr_branches;
      pr_attributes =
        (if formals = [] then [] else global_attributes nt) @
        rule.pr_attributes
    }

(* -------------------------------------------------------------------------- *)

(* The entry points of the graph traversal include the nonterminal symbols of
   sort [*]. (Not just the start symbols, as we haven't run the reachability
   analysis, and the grammar may contain unreachable parts, which we still
   want to expand.) Because a start symbol must have sort [*], this includes
   the start symbols. *)

let () =
  g.p_rules |> StringMap.iter @@ fun nt prule ->
  if prule.pr_parameters = [] then
    let label = (nt, []) in
    enqueue label

(* -------------------------------------------------------------------------- *)

(* The parameters that appear in [%type] declarations and [%on_error_reduce]
   declarations are also considered entry points. They have sort [*]. *)

let subst_parameter param =
  subst_parameter Env.empty param

let subst_declaration (param, info) =
  assert (sort param = star);
  (subst_parameter param, info)

let subst_declarations decls =
  List.map subst_declaration decls

(* -------------------------------------------------------------------------- *)

(* An [%attribute] declaration for a parameter of sort [*] is treated as an
   entry point. An [%attribute] declaration for a symbol of higher sort is not
   regarded as an entry point, and at the end, is kept only if this symbol
   still appears in the expanded grammar. *)

(* This is done in two passes over the list of [%attribute] declarations,
   named [thingify] and [unthingify]. The first pass runs as part of the
   discovery of entry points, before the graph traversal. The second pass runs
   after the graph traversal is complete. *)

type thing =
  | TargetParameterOfSortStar of parameter
  | SourceParameterOfHigherSort of parameter

let thingify_parameter param : thing =
  if sort param = star then
    TargetParameterOfSortStar (subst_parameter param)
  else
    SourceParameterOfHigherSort param

let thingify_attribute_declaration (params, attrs) =
  (List.map thingify_parameter params, attrs)

let thingify_attribute_declarations decls =
  List.map thingify_attribute_declaration decls

let unthingify_parameter rules thing =
  match thing with
  | TargetParameterOfSortStar param ->
      (* This parameter has sort [star]. Keep it. *)
      Some param
  | SourceParameterOfHigherSort param ->
      (* This parameter has higher sort. It must be a symbol.
         Keep it if it still appears in the expanded grammar. *)
      let symbol, params = Parameter.destruct param in
      assert (params = []);
      let symbol = value symbol in
      if StringMap.mem symbol rules then Some param else None

let unthingify_attribute_declaration rules (params, attrs) =
  (MList.filter_map (unthingify_parameter rules) params, attrs)

let unthingify_attribute_declarations rules decls =
  List.map (unthingify_attribute_declaration rules) decls

(* -------------------------------------------------------------------------- *)

(* Put everything together a construct a new grammar. *)

let g =
  (* Discovery of entry points. *)
  let p_types = subst_declarations g.p_types
  and p_on_error_reduce = subst_declarations g.p_on_error_reduce
  and things = thingify_attribute_declarations g.p_symbol_attributes in
  (* Graph traversal. *)
  repeatedly visit;
  (* Construction of the new grammar. *)
  let p_rules = rules() in
  let p_symbol_attributes = unthingify_attribute_declarations p_rules things in
  { g with
    p_types;
    p_on_error_reduce;
    p_symbol_attributes;
    p_rules }

end (* Run *)

(* -------------------------------------------------------------------------- *)

(* Re-package the above functor as a function. *)

let expand mode sorts g =
  let module G = Run(struct
    let mode = mode
    let sorts = sorts
    let g = g
  end) in
  G.g
