(************************************************************************)
(*         *      The Rocq Prover / The Rocq Development Team           *)
(*  v      *         Copyright INRIA, CNRS and contributors             *)
(* <O___,, * (see version control and CREDITS file for authors & dates) *)
(*   \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 Miniml
open Constr
open Declarations
open Mod_declarations
open Names
open ModPath
open Libnames
open Pp
open CErrors
open Util
open Table
open Extraction
open Modutil
open Common

(****************************************)
(*S Part I: computing Rocq environment. *)
(****************************************)

let toplevel_env () =
  let mp, struc = Safe_typing.flatten_env (Global.safe_env ()) in
  mp, List.rev struc

let environment_until dir_opt =
  let rec parse = function
    | [] when Option.is_empty dir_opt -> [toplevel_env ()]
    | [] -> []
    | d :: l ->
      let meb =
        Modops.destr_nofunctor (MPfile d) (mod_type @@ Global.lookup_module (MPfile d))
      in
      match dir_opt with
      | Some d' when DirPath.equal d d' -> [MPfile d, meb]
      | _ -> (MPfile d, meb) :: (parse l)
  in parse (Library.loaded_libraries ())


(*s Visit:
  a structure recording the needed dependencies for the current extraction *)

module type VISIT = sig
  type t
  (* Environment of visited data *)
  val make : unit -> t

  (* Add the module_path and all its prefixes to the mp visit list.
     We'll keep all fields of these modules. *)
  val add_mp_all : t -> ModPath.t -> unit

  (* Add reference / ... in the visit lists.
     These functions silently add the mp of their arg in the mp list *)
  val add_ref : t -> global -> unit
  val add_kn : t -> KerName.t -> InfvInst.t -> unit
  val add_decl_deps : t -> ml_decl -> unit
  val add_spec_deps : t -> ml_spec -> unit

  (* Test functions:
     is a particular object a needed dependency for the current extraction ? *)
  val needed_ind : t -> MutInd.t -> InfvInst.t -> bool
  val needed_cst : t -> Constant.t -> InfvInst.t -> bool
  val needed_mp : t -> ModPath.t -> bool
  val needed_mp_all : t -> ModPath.t -> bool
end

module Visit : VISIT = struct
  module KNOrd =
  struct
    type t = KerName.t * InfvInst.t
    let compare (kn1, i1) (kn2, i2) =
      let c = KerName.compare kn1 kn2 in
      if Int.equal c 0 then InfvInst.compare i1 i2 else c
  end
  module KNset = Set.Make(KNOrd)

  type t =
      { mutable kn : KNset.t;
        mutable mp : MPset.t;
        mutable mp_all : MPset.t }
  (* the imperative internal visit lists *)
  let make () = {
    kn = KNset.empty;
    mp = MPset.empty;
    mp_all = MPset.empty;
  }
  (* the accessor functions *)
  let needed_ind v i inst = KNset.mem (MutInd.user i, inst) v.kn
  let needed_cst v c inst = KNset.mem (Constant.user c, inst) v.kn
  let needed_mp v mp = MPset.mem mp v.mp || MPset.mem mp v.mp_all
  let needed_mp_all v mp = MPset.mem mp v.mp_all
  let add_mp v mp =
    check_loaded_modfile mp; v.mp <- MPset.union (prefixes_mp mp) v.mp
  let add_mp_all v mp =
    check_loaded_modfile mp;
    v.mp <- MPset.union (prefixes_mp mp) v.mp;
    v.mp_all <- MPset.add mp v.mp_all
  let add_kn v kn inst = v.kn <- KNset.add (kn, inst) v.kn; add_mp v (KerName.modpath kn)
  let add_ref v r = let open GlobRef in match r.glob with
    | ConstRef c -> add_kn v (Constant.user c) r.inst
    | IndRef (ind,_) | ConstructRef ((ind,_),_) -> add_kn v (MutInd.user ind) r.inst
    | VarRef _ -> assert false
  let add_decl_deps v decl =
    decl_iter_references (fun kn -> add_ref v kn) (fun r -> add_ref v r) (fun r -> add_ref v r) decl
  let add_spec_deps v spec =
    spec_iter_references (fun r -> add_ref v r) (fun r -> add_ref v r) (fun r -> add_ref v r) spec
end

let get_mono_inst_univs = function
| Monomorphic -> [InfvInst.empty]
| Polymorphic uctx -> InfvInst.generate uctx

let get_mono_inst = function
| SFBconst cb -> get_mono_inst_univs cb.const_universes
| SFBmind mib -> get_mono_inst_univs mib.mind_universes
| SFBrules _ -> [InfvInst.empty]
| SFBmodule _ | SFBmodtype _ -> assert false

let add_field_label venv mp = function
  | (lab, (SFBconst _|SFBmind _ | SFBrules _  as f)) ->
    let insts = get_mono_inst f in
    List.iter (fun inst -> Visit.add_kn venv (KerName.make mp lab) inst) insts
  | (lab, (SFBmodule _|SFBmodtype _)) -> Visit.add_mp_all venv (MPdot (mp,lab))

let rec add_labels venv mp = function
  | MoreFunctor (_,_,m) -> add_labels venv mp m
  | NoFunctor sign -> List.iter (fun f -> add_field_label venv mp f) sign

exception Impossible

let check_arity env cb =
  let t = cb.const_type in
  if Reduction.is_arity env t then raise Impossible

let get_body lbody =
  EConstr.of_constr lbody

let check_fix env sg cb i =
  match cb.const_body with
    | Def lbody ->
        (match EConstr.kind sg (get_body lbody) with
          | Fix ((_,j),recd) when Int.equal i j -> check_arity env cb; (true,recd)
          | CoFix (j,recd) when Int.equal i j -> check_arity env cb; (false,recd)
          | _ -> raise Impossible)
    | Undef _ | OpaqueDef _ | Primitive _ | Symbol _ -> raise Impossible

let prec_declaration_equal sg (na1, ca1, ta1) (na2, ca2, ta2) =
  Array.equal (Context.eq_annot Name.equal (EConstr.ERelevance.equal sg)) na1 na2 &&
  Array.equal (EConstr.eq_constr sg) ca1 ca2 &&
  Array.equal (EConstr.eq_constr sg) ta1 ta2

let factor_fix env sg l cb msb =
  let _,recd as check = check_fix env sg cb 0 in
  let n = Array.length (let fi,_,_ = recd in fi) in
  if Int.equal n 1 then [|l|], recd, msb
  else begin
    if List.length msb < n-1 then raise Impossible;
    let msb', msb'' = List.chop (n-1) msb in
    let labels = Array.make n l in
    List.iteri
      (fun j ->
         function
           | (l,SFBconst cb') ->
              let check' = check_fix env sg cb' (j+1) in
              if not ((fst check : bool) == (fst check') &&
                        prec_declaration_equal sg (snd check) (snd check'))
               then raise Impossible;
               labels.(j+1) <- l;
           | _ -> raise Impossible) msb';
    labels, recd, msb''
  end

(** Expanding a [module_alg_expr] into a version without abbreviations
    or functor applications. This is done via a detour to entries
    (hack proposed by Elie)
*)

let vm_state =
  (* VM bytecode is not needed here *)
  let vm_handler _ _ _ () = (), None in
  ((), { Mod_typing.vm_handler })

let expand_mexpr env mp me =
  let inl = Some (Flags.get_inline_level()) in
  let state = ((Environ.universes env, Univ.Constraints.empty), Reductionops.inferred_universes) in
  let mb, (_, cst), _ = Mod_typing.translate_module state vm_state env mp inl (MExpr ([], me, None)) in
  mod_type mb, mod_delta mb

let expand_modtype env mp me =
  let inl = Some (Flags.get_inline_level()) in
  let state = ((Environ.universes env, Univ.Constraints.empty), Reductionops.inferred_universes) in
  let mtb, _cst, _ = Mod_typing.translate_modtype state vm_state env mp inl ([],me) in
  mtb

let no_delta = Mod_subst.empty_delta_resolver

let flatten_modtype env mp me_alg struc_opt =
  match struc_opt with
  | Some me -> me, no_delta mp
  | None ->
     let mtb = expand_modtype env mp me_alg in
     mod_type mtb, mod_delta mtb

(** Ad-hoc update of environment, inspired by [Mod_typing.check_with_aux_def].
*)

let env_for_mtb_with_def env mp me reso idl =
  let struc = Modops.destr_nofunctor mp me in
  let l = Label.of_id (List.hd idl) in
  let spot = function (l',SFBconst _) -> Label.equal l l' | _ -> false in
  let before = fst (List.split_when spot struc) in
  Modops.add_structure mp before reso env

let make_cst resolver mp l =
  Mod_subst.constant_of_delta_kn resolver (KerName.make mp l)

let make_mind resolver mp l =
  Mod_subst.mind_of_delta_kn resolver (KerName.make mp l)

(* From a [structure_body] (i.e. a list of [structure_field_body])
   to specifications. *)

let rec extract_structure_spec table venv env mp reso = function
  | [] -> []
  | (l, SFBconst cb) :: msig ->
    let insts = get_mono_inst_univs cb.const_universes in
    let c = make_cst reso mp l in
    let map inst = extract_constant_spec table env c inst cb in
    let consts = List.map map insts in
    let specs = extract_structure_spec table venv env mp reso msig in
    let fold s specs =
      if logical_spec s then specs
      else
        let () = Visit.add_spec_deps venv s in
        (l, Spec s) :: specs
    in
    List.fold_right fold consts specs
  | (l, SFBmind mib) :: msig ->
    let insts = get_mono_inst_univs mib.mind_universes in
    let mind = make_mind reso mp l in
    let map inst = Sind (extract_inductive table env mind inst) in
    let minds = List.map map insts in
    let specs = extract_structure_spec table venv env mp reso msig in
    let fold s specs =
      if logical_spec s then specs
      else
        let () = Visit.add_spec_deps venv s in
        (l, Spec s) :: specs
    in
    List.fold_right fold minds specs
  | (l, SFBrules _) :: msig ->
      let specs = extract_structure_spec table venv env mp reso msig in
      specs
  | (l,SFBmodule mb) :: msig ->
      let specs = extract_structure_spec table venv env mp reso msig in
      let mp = MPdot (mp, l) in
      let spec = extract_mbody_spec table venv env mp mb in
      (l,Smodule spec) :: specs
  | (l,SFBmodtype mtb) :: msig ->
      let specs = extract_structure_spec table venv env mp reso msig in
      let mp = MPdot (mp, l) in
      let spec = extract_mbody_spec table venv env mp mtb in
      (l,Smodtype spec) :: specs

(* From [module_expression] to specifications *)

(* Invariant: the [me_alg] given to [extract_mexpr_spec] and
   [extract_mexpression_spec] should come from a [mod_type_alg] field.
   This way, any encountered [MEident] should be a true module type. *)

and extract_mexpr_spec table venv env mp1 (me_struct_o,me_alg) = match me_alg with
  | MEident mp ->
    let () = Visit.add_mp_all venv mp in
    MTident mp
  | MEwith(me',WithDef(idl,(c,ctx)))->
      let () = match ctx with
      | None -> ()
      | Some auctx ->
        (* XXX *)
        if Array.is_empty (UVars.AbstractContext.names auctx).quals then ()
        else user_err Pp.(str "Extraction of \"with Definition\" clauses not supported for sort polymorphic definitions.")
      in
      let me_struct,delta = flatten_modtype env mp1 me' me_struct_o in
      let env' = env_for_mtb_with_def env mp1 me_struct delta idl in
      let mt = extract_mexpr_spec table venv env mp1 (None,me') in
      let sg = Evd.from_env env in
      (match extract_with_type table env' sg (EConstr.of_constr c) with
       (* cb may contain some kn *)
         | None -> mt
         | Some (vl,typ) ->
            let () = type_iter_references (fun r -> Visit.add_ref venv r) typ in
            MTwith (mt, ML_With_type (InfvInst.empty, idl, vl, typ)))
  | MEwith(me',WithMod(idl,mp))->
      let () = Visit.add_mp_all venv mp in
      MTwith (extract_mexpr_spec table venv env mp1 (None, me'), ML_With_module(idl, mp))
  | MEapply _ ->
     (* No higher-order module type in OCaml : we use the expanded version *)
     let me_struct,delta = flatten_modtype env mp1 me_alg me_struct_o in
     extract_msignature_spec table venv env mp1 delta me_struct

and extract_mexpression_spec table venv env mp1 (me_struct,me_alg) = match me_alg with
  | MEMoreFunctor me_alg' ->
      let mbid, mtb, me_struct' = match me_struct with
      | MoreFunctor (mbid, mtb, me') -> (mbid, mtb, me')
      | _ -> assert false
      in
      let mp = MPbound mbid in
      let env' = Modops.add_module_parameter mbid mtb env in
      MTfunsig (mbid, extract_mbody_spec table venv env mp mtb,
                extract_mexpression_spec table venv env' mp1 (me_struct', me_alg'))
  | MENoFunctor m -> extract_mexpr_spec table venv env mp1 (Some me_struct, m)

and extract_msignature_spec table venv env mp1 reso = function
  | NoFunctor struc ->
      let env' = Modops.add_structure mp1 struc reso env in
      MTsig (mp1, extract_structure_spec table venv env' mp1 reso struc)
  | MoreFunctor (mbid, mtb, me) ->
      let mp = MPbound mbid in
      let env' = Modops.add_module_parameter mbid mtb env in
      MTfunsig (mbid, extract_mbody_spec table venv env mp mtb,
                extract_msignature_spec table venv env' mp1 reso me)

and extract_mbody_spec : 'a. State.t -> _ -> _ -> _ -> 'a generic_module_body -> _ =
  fun table venv env mp mb -> match mod_type_alg mb with
  | Some ty -> extract_mexpression_spec table venv env mp (mod_type mb, ty)
  | None -> extract_msignature_spec table venv env mp (mod_delta mb) (mod_type mb)

(* From a [structure_body] (i.e. a list of [structure_field_body])
   to implementations.

   NB: when [all=false], the evaluation order of the list is
   important: last to first ensures correct dependencies.
*)

let rec extract_structure table access venv env mp reso ~all = function
  | [] -> []
  | (l, SFBconst cb) :: struc ->
    let sg = Evd.from_env env in
    let fix, struc = match factor_fix env sg l cb struc with
    | (vl, recd, struc) -> Some (vl, recd), struc
    | exception Impossible -> None, struc
    in
    let ms = extract_structure table access venv env mp reso ~all struc in
    let insts = get_mono_inst_univs cb.const_universes in
    let c = make_cst reso mp l in
    let map inst = match fix with
    | None ->
      let b = Visit.needed_cst venv c inst in
      if all || b then
        let d = extract_constant table access env c inst cb in
        if (not b) && (logical_decl d) then None
        else
        let () = Visit.add_decl_deps venv d in
        Some (l, SEdecl d)
      else None
    | Some (vl, recd) ->
      let vc = Array.map (make_cst reso mp) vl in
      let b = Array.exists (fun vf -> Visit.needed_cst venv vf inst) vc in
      if all || b then
        let d = extract_fixpoint table env sg vc inst recd in
        if (not b) && (logical_decl d) then None
        else
          let () = Visit.add_decl_deps venv d in
          Some (l, SEdecl d)
      else None
    in
    let consts = List.map_filter map insts in
    consts @ ms
  | (l, SFBmind mib) :: struc ->
    let ms = extract_structure table access venv env mp reso ~all struc in
    let insts = get_mono_inst_univs mib.mind_universes in
    let mind = make_mind reso mp l in
    let map inst =
      let b = Visit.needed_ind venv mind inst in
      if all || b then
        let d = Dind (extract_inductive table env mind inst) in
        if (not b) && (logical_decl d) then None
        else
          let () = Visit.add_decl_deps venv d in
          Some (l, SEdecl d)
      else None
    in
    let inds = List.map_filter map insts in
    inds @ ms
  | (l, SFBrules rrb) :: struc ->
      let inst = InfvInst.empty in (* FIXME ? *)
      let b = List.exists (fun (cst, _) -> Visit.needed_cst venv cst inst) rrb.rewrules_rules in
      let ms = extract_structure table access venv env mp reso ~all struc in
      if all || b then begin
        List.iter (fun (cst, _) -> Table.add_symbol_rule (State.get_table table) { glob = ConstRef cst; inst } l) rrb.rewrules_rules;
        ms
      end else ms
  | (l,SFBmodule mb) :: struc ->
      let ms = extract_structure table access venv env mp reso ~all struc in
      let mp = MPdot (mp,l) in
      let all' = all || Visit.needed_mp_all venv mp in
      if all' || Visit.needed_mp venv mp then
        (l, SEmodule (extract_module table access venv env mp ~all:all' mb)) :: ms
      else ms
  | (l,SFBmodtype mtb) :: struc ->
      let ms = extract_structure table access venv env mp reso ~all struc in
      let mp = MPdot (mp,l) in
      if all || Visit.needed_mp venv mp then
        (l, SEmodtype (extract_mbody_spec table venv env mp mtb)) :: ms
      else ms

(* From [module_expr] and [module_expression] to implementations *)

and extract_mexpr table access venv env mp = function
  | MEwith _ -> assert false (* no 'with' syntax for modules *)
  | me when lang () != Ocaml ->
      (* In Haskell/Scheme, we expand everything.
         For now, we also extract everything, dead code will be removed later
         (see [Modutil.optimize_struct]. *)
      let sign, delta = expand_mexpr env mp me in
      extract_msignature table access venv env mp delta ~all:true sign
  | MEident mp ->
      if is_modfile mp && not (State.get_modular table) then error_MPfile_as_mod mp false;
      Visit.add_mp_all venv mp; Miniml.MEident mp
  | MEapply (me, arg) ->
      Miniml.MEapply (extract_mexpr table access venv env mp me,
                      extract_mexpr table access venv env mp (MEident arg))

and extract_mexpression table access venv env mp mty = function
  | MENoFunctor me -> extract_mexpr table access venv env mp me
  | MEMoreFunctor me ->
      let (mbid, mtb, mty) = match mty with
      | MoreFunctor (mbid, mtb, mty) -> (mbid, mtb, mty)
      | NoFunctor _ -> assert false
      in
      let mp1 = MPbound mbid in
      let env' = Modops.add_module_parameter mbid mtb env in
      Miniml.MEfunctor
        (mbid,
         extract_mbody_spec table venv env mp1 mtb,
         extract_mexpression table access venv env' mp mty me)

and extract_msignature table access venv env mp reso ~all = function
  | NoFunctor struc ->
      let env' = Modops.add_structure mp struc reso env in
      Miniml.MEstruct (mp,extract_structure table access venv env' mp reso ~all struc)
  | MoreFunctor (mbid, mtb, me) ->
      let mp1 = MPbound mbid in
      let env' = Modops.add_module_parameter mbid mtb env in
      Miniml.MEfunctor
        (mbid,
         extract_mbody_spec table venv env mp1 mtb,
         extract_msignature table access venv env' mp reso ~all me)

and extract_module table access venv env mp ~all mb =
  (* A module has an empty [mod_expr] when :
     - it is a module variable (for instance X inside a Module F [X:SIG])
     - it is a module assumption (Declare Module).
     Since we look at modules from outside, we shouldn't have variables.
     But a Declare Module at toplevel seems legal (cf #2525). For the
     moment we don't support this situation. *)
  let impl = match Mod_declarations.mod_expr mb with
    | Abstract -> error_no_module_expr mp
    | Algebraic me -> extract_mexpression table access venv env mp (mod_type mb) me
    | Struct sign ->
      (* This module has a signature, otherwise it would be FullStruct.
         We extract just the elements required by this signature. *)
      let () = add_labels venv mp (mod_type mb) in
      let sign = Modops.annotate_struct_body sign (mod_type mb) in
      extract_msignature table access venv env mp (mod_delta mb) ~all:false sign
    | FullStruct -> extract_msignature table access venv env mp (mod_delta mb) ~all (mod_type mb)
  in
  (* Slight optimization: for modules without explicit signatures
     ([FullStruct] case), we build the type out of the extracted
     implementation *)
  let typ = match Mod_declarations.mod_expr mb with
    | FullStruct ->
      assert (Option.is_empty @@ mod_type_alg mb);
      mtyp_of_mexpr impl
    | _ -> extract_mbody_spec table venv env mp mb
  in
  { ml_mod_expr = impl;
    ml_mod_type = typ }

let mono_environment table ~opaque_access refs mpl =
  let venv = Visit.make () in
  let () = List.iter (fun r -> Visit.add_ref venv r) refs in
  let () = List.iter (fun mp -> Visit.add_mp_all venv mp) mpl in
  let env = Global.env () in
  let l = List.rev (environment_until None) in
  List.rev_map
    (fun (mp,struc) ->
      mp, extract_structure table opaque_access venv env mp (no_delta mp) ~all:(Visit.needed_mp_all venv mp) struc)
    l

(**************************************)
(*S Part II : Input/Output primitives *)
(**************************************)

let descr () = match lang () with
  | Ocaml -> Ocaml.ocaml_descr
  | Haskell -> Haskell.haskell_descr
  | Scheme -> Scheme.scheme_descr
  | JSON -> Json.json_descr

(* From a filename string "foo.ml" or "foo", builds "foo.ml" and "foo.mli"
   Works similarly for the other languages. *)

let default_id = Id.of_string "Main"

let mono_filename f =
  let d = descr () in
  match f with
    | None -> None, None, default_id
    | Some f ->
        let f =
          if Filename.check_suffix f d.file_suffix then
            Filename.chop_suffix f d.file_suffix
          else f
        in
        let id =
          if lang () != Haskell then default_id
          else
            try Id.of_string (Filename.basename f)
            with UserError _ ->
              user_err Pp.(str "Extraction: provided filename is not a valid identifier")
        in
        let f =
          if Filename.is_relative f then
            Filename.concat (output_directory ()) f
          else f
        in
        Some (f^d.file_suffix), Option.map ((^) f) d.sig_suffix, id

(* Builds a suitable filename from a module id *)

let module_filename table mp =
  let f = file_of_modfile (State.get_table table) mp in
  let id = Id.of_string f in
  let f = Filename.concat (output_directory ()) f in
  let d = descr () in
  let fimpl_base = d.file_naming table mp ^ d.file_suffix in
  let fimpl = Filename.concat (output_directory ()) fimpl_base in
  Some fimpl, Option.map ((^) f) d.sig_suffix, id

(*s Extraction of one decl to stdout. *)

let print_one_decl table struc mp decl =
  let d = descr () in
  let () = State.reset table in
  let table = State.set_phase table Pre in
  ignore (d.pp_struct table struc);
  let table = State.set_phase table Impl in
  let ans = State.with_visibility table mp [] begin fun table ->
    d.pp_decl table decl
  end in
  v 0 ans

(*s Extraction of a ml struct to a file. *)

(** For Recursive Extraction, writing directly on stdout
    won't work with rocqide, we use a buffer instead *)

let formatter buf dry file =
  let ft =
    if dry then Format.make_formatter (fun _ _ _ -> ()) (fun _ -> ())
    else
      match file with
        | Some f -> Topfmt.with_output_to f
        | None -> Format.formatter_of_buffer buf
  in
  (* XXX: Fixme, this shouldn't depend on Topfmt *)
  (* We never want to see ellipsis ... in extracted code *)
  Format.pp_set_max_boxes ft max_int;
  (* We reuse the width information given via "Set Printing Width" *)
  (match Topfmt.get_margin () with
    | None -> ()
    | Some i ->
      Format.pp_set_margin ft i;
      Format.pp_set_max_indent ft (i-10));
      (* note: max_indent should be < margin above, otherwise it's ignored *)
  ft

let get_comment () =
  let s = file_comment () in
  if String.is_empty s then None
  else
    let split_comment = Str.split (Str.regexp "[ \t\n]+") s in
    Some (prlist_with_sep spc str split_comment)

let print_structure_to_file table (fn,si,mo) dry struc =
  let buf = Buffer.create 1000 in
  let d = descr () in
  let () = State.reset table in
  let unsafe_needs = {
    mldummy = struct_ast_search Mlutil.isMLdummy struc;
    tdummy = struct_type_search Mlutil.isTdummy struc;
    tunknown = struct_type_search ((==) Tunknown) struc;
    magic =
      if lang () != Haskell then false
      else struct_ast_search (function MLmagic _ -> true | _ -> false) struc }
  in
  (* First, a dry run, for computing objects to rename or duplicate *)
  let table = State.set_phase table Pre in
  ignore (d.pp_struct table struc);
  let opened = opened_libraries table in
  (* Print the implementation *)
  let cout = if dry then None else Option.map open_out fn in
  let ft = formatter buf dry cout in
  let comment = get_comment () in
  begin try
    (* The real printing of the implementation *)
    let table = State.set_phase table Impl in
    pp_with ft (d.preamble table mo comment opened unsafe_needs);
    pp_with ft (d.pp_struct table struc);
    Format.pp_print_flush ft ();
    Option.iter close_out cout;
  with reraise ->
    Format.pp_print_flush ft ();
    Option.iter close_out cout; raise reraise
  end;
  if not dry then Option.iter info_file fn;
  (* Now, let's print the signature *)
  Option.iter
    (fun si ->
       let cout = open_out si in
       let ft = formatter buf false (Some cout) in
       begin try
         let table = State.set_phase table Intf in
         pp_with ft (d.sig_preamble table mo comment opened unsafe_needs);
         pp_with ft (d.pp_sig table (signature_of_structure struc));
         Format.pp_print_flush ft ();
         close_out cout;
       with reraise ->
         Format.pp_print_flush ft ();
         close_out cout; raise reraise
       end;
       info_file si)
    (if dry then None else si);
  (* Print the buffer content via Rocq standard formatter (ok with rocqide). *)
  if not (Int.equal (Buffer.length buf) 0) then begin
    Feedback.msg_notice (str (Buffer.contents buf));
  end


(*********************************************)
(*s Part III: the actual extraction commands *)
(*********************************************)

let init ?(inner=false) modular library =
  if not inner then check_inside_section ();
  let keywords = (descr ()).keywords in
  let state = State.make ~modular ~library ~keywords () in
  if modular && lang () == Scheme then error_scheme ();
  state

let warns table =
  let table = State.get_table table in
  warning_opaques table (access_opaque ());
  warning_axioms table

(* From a list of [reference], let's retrieve whether they correspond
   to modules or [global_reference]. Warn the user if both is possible. *)

let rec locate_ref = function
  | [] -> [],[]
  | qid::l ->
      let mpo = try Some (Nametab.locate_module qid) with Not_found -> None
      and ro =
        try
          let gr = Smartlocate.global_with_alias qid in
          let inst = Environ.universes_of_global (Global.env ()) gr in
          Some (List.map (fun inst -> { glob = gr; inst }) (InfvInst.generate inst))
        with Nametab.GlobalizationError _ | UserError _ -> None
      in
      match mpo, ro with
        | None, None -> Nametab.error_global_not_found ~info:Exninfo.null qid
        | None, Some r ->
          let refs, mps = locate_ref l in r @ refs,mps
        | Some mp, None -> let refs,mps = locate_ref l in refs,mp::mps
        | Some mp, Some r ->
          let () = warning_ambiguous_name ?loc:qid.CAst.loc (qid, mp, (List.hd r).glob) in
          let refs,mps = locate_ref l in refs,mp::mps

(*s Recursive extraction in the Rocq toplevel. The vernacular command is
    \verb!Recursive Extraction! [qualid1] ... [qualidn]. Also used when
    extracting to a file with the command:
    \verb!Extraction "file"! [qualid1] ... [qualidn]. *)

let full_extr opaque_access f (refs,mps) =
  let table = init false false in
  List.iter (fun mp -> if is_modfile mp then error_MPfile_as_mod mp true) mps;
  let struc = optimize_struct table (refs,mps) (mono_environment table ~opaque_access refs mps) in
  let () = warns table in
  print_structure_to_file table (mono_filename f) false struc

let full_extraction ~opaque_access f lr =
  full_extr opaque_access f (locate_ref lr)

(*s Separate extraction is similar to recursive extraction, with the output
   decomposed in many files, one per Rocq .v file *)

let separate_extraction ~opaque_access lr =
  let table = init true false in
  let refs,mps = locate_ref lr in
  let struc = optimize_struct table (refs,mps) (mono_environment table ~opaque_access refs mps) in
  let () = List.iter (function
    | MPfile _, _ -> ()
    | (MPdot _ | MPbound _), _ ->
      user_err (str "Separate Extraction from inside a module is not supported."))
      struc
  in
  let () = warns table in
  let print = function
    | (MPfile dir as mp, sel) as e ->
        print_structure_to_file table (module_filename table mp) false [e]
    | (MPdot _ | MPbound _), _ -> assert false
  in
  let () = List.iter print struc in
  ()

(*s Simple extraction in the Rocq toplevel. The vernacular command
    is \verb!Extraction! [qualid]. *)

let simple_extraction ~opaque_access r =
  match locate_ref [r] with
  | ([], [mp]) as p -> full_extr opaque_access None p
  | [r],[] ->
      let table = init false false in
      let struc = optimize_struct table ([r],[]) (mono_environment table ~opaque_access [r] []) in
      let d = get_decl_in_structure r struc in
      let () = warns table in
      let flag =
        if is_custom r then str "(** User defined extraction *)" ++ fnl()
        else mt ()
      in
      let ans = flag ++ print_one_decl table struc (modpath_of_r r) d in
      Feedback.msg_notice ans
  | _ -> assert false


(*s (Recursive) Extraction of a library. The vernacular command is
  \verb!(Recursive) Extraction Library! [M]. *)

let extraction_library ~opaque_access is_rec CAst.{loc;v=m} =
  let table = init true true in
  let dir_m =
    (* XXX WTF is going on here? *)
    let q = qualid_of_ident m in
    try Nametab.full_name_module q with Not_found -> error_unknown_module ?loc q
  in
  let dir_m = dirpath_of_path dir_m in
  let venv = Visit.make () in
  let () = Visit.add_mp_all venv (MPfile dir_m) in
  let env = Global.env () in
  let l = List.rev (environment_until (Some dir_m)) in
  let select l (mp,struc) =
    if Visit.needed_mp venv mp
    then (mp, extract_structure table opaque_access venv env mp (no_delta mp) ~all:true struc) :: l
    else l
  in
  let struc = List.fold_left select [] l in
  let struc = optimize_struct table ([],[]) struc in
  let () = warns table in
  let print = function
    | (MPfile dir as mp, sel) as e ->
        let dry = not is_rec && not (DirPath.equal dir dir_m) in
        print_structure_to_file table (module_filename table mp) dry [e]
    | _ -> assert false
  in
  let () = List.iter print struc in
  ()

(* For the test-suite :
   extraction to a temporary file + run ocamlc on it *)

let compile f =
  try
    let args = [ "ocamlc"
               ; "-package";"zarith"
               ; "-I"; Filename.dirname f
               ; "-c"; f^"i"
               ; f
               ] in
    let res = CUnix.sys_command (Boot.Env.ocamlfind ()) args in
    match res with
    | Unix.WEXITED 0 -> ()
    | Unix.WEXITED n | Unix.WSIGNALED n | Unix.WSTOPPED n ->
       CErrors.user_err
         Pp.(str "Compilation of file " ++ str f ++
             str " failed with exit code " ++ int n)
  with Unix.Unix_error (e,_,_) ->
    CErrors.user_err
      Pp.(str "Compilation of file " ++ str f ++
          str " failed with error " ++ str (Unix.error_message e))

let remove f =
  if Sys.file_exists f then Sys.remove f

let extract_and_compile ~opaque_access l =
  if lang () != Ocaml then
    CErrors.user_err (Pp.str "This command only works with OCaml extraction");
  let f = Filename.temp_file "testextraction" ".ml" in
  let () = full_extraction ~opaque_access (Some f) l in
  let () = compile f in
  let () = remove f; remove (f^"i") in
  let base = Filename.chop_suffix f ".ml" in
  let () = remove (base^".cmo"); remove (base^".cmi") in
  Feedback.msg_notice (str "Extracted code successfully compiled")

(* Show the extraction of the current ongoing proof *)
let show_extraction ~pstate =
  let table = init ~inner:true false false in
  let prf = Declare.Proof.get pstate in
  let sigma, env = Declare.Proof.get_current_context pstate in
  let trms = Proof.partial_proof prf in
  let extr_term t =
    (* FIXME: substitute relevances with ground ones *)
    let ast, ty = extract_constr table env sigma t in
    let mp = Lib.current_mp () in
    let l = Label.of_id (Declare.Proof.get_name pstate) in
    let fake_ref = { glob = GlobRef.ConstRef (Constant.make2 mp l); inst = InfvInst.empty } in
    let decl = Dterm (fake_ref, ast, ty) in
    print_one_decl table [] mp decl
  in
  Feedback.msg_notice (Pp.prlist_with_sep Pp.fnl extr_term trms)
