(************************************************************************)
(*         *   The Coq Proof Assistant / The Coq Development Team       *)
(*  v      *   INRIA, CNRS and contributors - Copyright 1999-2019       *)
(* <O___,, *       (see CREDITS file for the list of authors)           *)
(*   \VV/  **************************************************************)
(*    //   *    This file is distributed under the terms of the         *)
(*         *     GNU Lesser General Public License Version 2.1          *)
(*         *     (see LICENSE file for the text of the license)         *)
(************************************************************************)

open Util
open Names
open Univ

module NamedDecl = Context.Named.Declaration

type section_entry =
| SecDefinition of Constant.t
| SecInductive of MutInd.t

type 'a entry_map = 'a Cmap.t * 'a Mindmap.t

type 'a section = {
  sec_context : int;
  (** Length of the named context suffix that has been introduced locally *)
  sec_mono_universes : ContextSet.t;
  sec_poly_universes : Name.t array * UContext.t;
  (** Universes local to the section *)
  has_poly_univs : bool;
  (** Are there polymorphic universes or constraints, including in previous sections. *)
  sec_entries : section_entry list;
  (** Definitions introduced in the section *)
  sec_data : (Instance.t * AUContext.t) entry_map;
  (** Additional data synchronized with the section *)
  sec_custom : 'a;
}

(** Sections can be nested with the proviso that no monomorphic section can be
    opened inside a polymorphic one. The reverse is allowed. *)
type 'a t = 'a section list

let empty = []

let is_empty = List.is_empty

let has_poly_univs = function
  | [] -> false
  | sec :: _ -> sec.has_poly_univs

let find_emap e (cmap, imap) = match e with
| SecDefinition con -> Cmap.find con cmap
| SecInductive ind -> Mindmap.find ind imap

let add_emap e v (cmap, imap) = match e with
| SecDefinition con -> (Cmap.add con v cmap, imap)
| SecInductive ind -> (cmap, Mindmap.add ind v imap)

let on_last_section f sections = match sections with
| [] -> CErrors.user_err (Pp.str "No opened section")
| sec :: rem -> f sec :: rem

let with_last_section f sections = match sections with
| [] -> f None
| sec :: _ -> f (Some sec)

let push_local s =
  let on_sec sec = { sec with sec_context = sec.sec_context + 1 } in
  on_last_section on_sec s

let push_context (nas, ctx) s =
    let on_sec sec =
      if UContext.is_empty ctx then sec
      else
        let (snas, sctx) = sec.sec_poly_universes in
        let sec_poly_universes = (Array.append snas nas, UContext.union sctx ctx) in
        { sec with sec_poly_universes; has_poly_univs = true }
    in
    on_last_section on_sec s

let is_polymorphic_univ u s =
  let check sec =
    let (_, uctx) = sec.sec_poly_universes in
    Array.exists (fun u' -> Level.equal u u') (Instance.to_array (UContext.instance uctx))
  in
  List.exists check s

let push_constraints uctx s =
  let on_sec sec =
    if sec.has_poly_univs && Constraint.exists (fun (l,_,r) -> is_polymorphic_univ l s || is_polymorphic_univ r s) (snd uctx)
    then CErrors.user_err Pp.(str "Cannot add monomorphic constraints which refer to section polymorphic universes.");
    let uctx' = sec.sec_mono_universes in
    let sec_mono_universes =  (ContextSet.union uctx uctx') in
    { sec with sec_mono_universes }
  in
  on_last_section on_sec s

let open_section ~custom sections =
  let sec = {
    sec_context = 0;
    sec_mono_universes = ContextSet.empty;
    sec_poly_universes = ([||], UContext.empty);
    has_poly_univs = has_poly_univs sections;
    sec_entries = [];
    sec_data = (Cmap.empty, Mindmap.empty);
    sec_custom = custom;
  } in
  sec :: sections

let close_section sections =
  match sections with
  | sec :: sections ->
    sections, sec.sec_entries, sec.sec_mono_universes, sec.sec_custom
  | [] ->
    CErrors.user_err (Pp.str "No opened section.")

let make_decl_univs (nas,uctx) = abstract_universes nas uctx

let push_global ~poly e s =
  if is_empty s then s
  else if has_poly_univs s && not poly
  then CErrors.user_err
      Pp.(str "Cannot add a universe monomorphic declaration when \
               section polymorphic universes are present.")
  else
    let on_sec sec =
      { sec with
        sec_entries = e :: sec.sec_entries;
        sec_data = add_emap e (make_decl_univs sec.sec_poly_universes) sec.sec_data;
      }
    in
    on_last_section on_sec s

let push_constant ~poly con s = push_global ~poly (SecDefinition con) s

let push_inductive ~poly ind s = push_global ~poly (SecInductive ind) s

type abstr_info = {
  abstr_ctx : Constr.named_context;
  abstr_subst : Instance.t;
  abstr_uctx : AUContext.t;
}

let empty_segment = {
  abstr_ctx = [];
  abstr_subst = Instance.empty;
  abstr_uctx = AUContext.empty;
}

let extract_hyps sec vars hyps =
  (* FIXME: this code is fishy. It is supposed to check that declared section
     variables are an ordered subset of the ambient ones, but it doesn't check
     e.g. uniqueness of naming nor convertibility of the section data. *)
  let rec aux ids hyps = match ids, hyps with
  | id :: ids, decl :: hyps when Names.Id.equal id (NamedDecl.get_id decl) ->
    decl :: aux ids hyps
  | _ :: ids, hyps ->
    aux ids hyps
  | [], _ -> []
  in
  let ids = List.map NamedDecl.get_id @@ List.firstn sec.sec_context vars in
  aux ids hyps

let section_segment_of_entry vars e hyps sections =
  (* [vars] are the named hypotheses, [hyps] the subset that is declared by the
    global *)
  let with_sec s = match s with
  | None ->
    CErrors.user_err (Pp.str "No opened section.")
  | Some sec ->
    let hyps = extract_hyps sec vars hyps in
    let inst, auctx = find_emap e sec.sec_data in
    {
      abstr_ctx = hyps;
      abstr_subst = inst;
      abstr_uctx = auctx;
    }
  in
  with_last_section with_sec sections

let segment_of_constant env con s =
  let body = Environ.lookup_constant con env in
  let vars = Environ.named_context env in
  section_segment_of_entry vars (SecDefinition con) body.Declarations.const_hyps s

let segment_of_inductive env mind s =
  let mib = Environ.lookup_mind mind env in
  let vars = Environ.named_context env in
  section_segment_of_entry vars (SecInductive mind) mib.Declarations.mind_hyps s

let instance_from_variable_context =
  List.rev %> List.filter NamedDecl.is_local_assum %> List.map NamedDecl.get_id %> Array.of_list

let extract_worklist info =
  let args = instance_from_variable_context info.abstr_ctx in
  info.abstr_subst, args

let replacement_context env s =
  let with_sec sec = match sec with
  | None -> CErrors.user_err (Pp.str "No opened section.")
  | Some sec ->
    let cmap, imap = sec.sec_data in
    let cmap = Cmap.mapi (fun con _ -> extract_worklist @@ segment_of_constant env con s) cmap in
    let imap = Mindmap.mapi (fun ind _ -> extract_worklist @@ segment_of_inductive env ind s) imap in
    (cmap, imap)
  in
  with_last_section with_sec s

let is_in_section env gr s =
  let with_sec sec = match sec with
  | None -> false
  | Some sec ->
    let open GlobRef in
    match gr with
    | VarRef id ->
      let vars = List.firstn sec.sec_context (Environ.named_context env) in
      List.exists (fun decl -> Id.equal id (NamedDecl.get_id decl)) vars
    | ConstRef con ->
      Cmap.mem con (fst sec.sec_data)
    | IndRef (ind, _) | ConstructRef ((ind, _), _) ->
      Mindmap.mem ind (snd sec.sec_data)
  in
  with_last_section with_sec s