(*********************************************************************************)
(*                Chamo                                                          *)
(*                                                                               *)
(*    Copyright (C) 2003-2021 Institut National de Recherche en Informatique     *)
(*    et en Automatique. All rights reserved.                                    *)
(*                                                                               *)
(*    This program is free software; you can redistribute it and/or modify       *)
(*    it under the terms of the GNU Lesser General Public License version        *)
(*    3 as published by the Free Software Foundation.                            *)
(*                                                                               *)
(*    This program is distributed in the hope that it will be useful,            *)
(*    but WITHOUT ANY WARRANTY; without even the implied warranty of             *)
(*    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the              *)
(*    GNU General Public License for more details.                               *)
(*                                                                               *)
(*    You should have received a copy of the GNU General Public License          *)
(*    along with this program; if not, write to the Free Software                *)
(*    Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA                   *)
(*    02111-1307  USA                                                            *)
(*                                                                               *)
(*    Contact: Maxence.Guesdon@inria.fr                                          *)
(*                                                                               *)
(*********************************************************************************)

(** Getting information from ocaml-generated .cmt files. *)

let () = Compmisc.init_path ()

let load_cmt file =
  try Result.Ok (Cmt_format.read_cmt file)
  with e ->
    let msg = Printexc.to_string e in
    Result.Error msg

type lookup_res = {
    loc: Location.t ;
    path: Path.t option ;
    typ: Types.type_expr ;
    def: Location.t option ; }

let pos_in_loc (fname, p) loc =
  Misc.safe_same_files loc.Location.loc_start.Lexing.pos_fname fname
  && loc.loc_start.pos_cnum <= p
  && loc.loc_end.pos_cnum >= p


let apply_mapper mapper = function
| Cmt_format.Partial_structure str -> ignore(Untypeast.untype_structure ~mapper str)
| Partial_structure_item stri -> ignore(mapper.Untypeast.structure_item mapper stri)
| Partial_expression e -> ignore(Untypeast.untype_expression ~mapper e)
| Partial_pattern (_, p) -> ignore(Untypeast.untype_pattern ~mapper p)
| Partial_class_expr c -> ignore(mapper.Untypeast.class_expr mapper c)
| Partial_signature s -> ignore(Untypeast.untype_signature ~mapper s)
| Partial_signature_item si -> ignore(mapper.Untypeast.signature_item mapper si)
| Partial_module_type mt -> ignore(mapper.Untypeast.module_type mapper mt)

let lookup_in_binary_part pos binary_part =
  let best = ref None in
  let set_best res =
    match !best with
    | None -> best := Some res
    | Some r ->
      let wr = r.loc.loc_end.pos_cnum - r.loc.loc_start.pos_cnum in
      let wres = res.loc.loc_end.pos_cnum - res.loc.loc_start.pos_cnum in
      if wres < wr then
        best := Some res
  in
  let open Typedtree in
  let defmap = Untypeast.default_mapper in
  let expr mapper exp =
    if pos_in_loc pos exp.exp_loc then
      (
       let (def,path) =
         match exp.exp_desc with
         | Texp_ident (path, { loc }, tval_desc) ->
             let def =
               try
                 (* from ocaml/typing/cmt2annot.ml *)
                 let env =
                   try Env.env_of_only_summary
                     Envaux.env_from_summary exp.exp_env
                   with
                     Envaux.Error e ->
                       let b = Buffer.create 256 in
                       let fmt = Format.formatter_of_buffer b in
                       Envaux.report_error fmt e;
                       Format.pp_print_flush fmt () ;
                       Log.err (fun m -> m "%s" (Buffer.contents b));
                       exp.exp_env
                 in
                 let desc = Env.find_value path env in
                 let dloc = desc.Types.val_loc in
                 if dloc.Location.loc_ghost then None
                  else Some dloc
               with
                 Not_found -> None
             in
             (def, Some path)
         | _ -> (None, None)
       in
       set_best { loc = exp.exp_loc ;
         path ;
         typ = exp.exp_type; def ;
       } ;
       defmap.Untypeast.expr mapper exp
      )
    else
      defmap.Untypeast.expr mapper exp
  in
  let pat mapper p =
    let loc = p.pat_loc in
    if pos_in_loc pos loc then
      (
       let typ = p.pat_type in
       set_best { loc ; path = None ; typ ; def = None }
      );
    defmap.Untypeast.pat mapper p
  in
  let mapper = { defmap with
      expr ; pat
    }
  in
  apply_mapper mapper binary_part;
  !best

let cmt_parts cmt =
  match cmt.Cmt_format.cmt_annots with
  | Implementation str -> [Cmt_format.Partial_structure str]
  | Partial_implementation arr -> Array.to_list arr
  | _ -> []

let lookup_by_pos pos cmt =
  Envaux.reset_cache ();
  List.iter Load_path.add_dir cmt.Cmt_format.cmt_loadpath;
  match cmt_parts cmt with
  | [] -> None
  | l ->
      let rec iter = function
      | [] -> None
      | h :: q ->
          match lookup_in_binary_part pos h with
          | None -> iter q
          | x -> x
      in
      iter l

let fold_exp_idents f acc cmt =
  Envaux.reset_cache ();
  List.iter Load_path.add_dir cmt.Cmt_format.cmt_loadpath;
  let parts = cmt_parts cmt in
  let idents = ref [] in
  let open Typedtree in
  let defmap = Untypeast.default_mapper in
  let expr mapper exp =
    let () =
      match exp.exp_desc with
      | Texp_ident (path, { loc }, tval_desc) ->
          let defloc =
            try
              let env =
                try Env.env_of_only_summary
                  Envaux.env_from_summary exp.exp_env
                with
                  Envaux.Error e ->
                    let b = Buffer.create 256 in
                    let fmt = Format.formatter_of_buffer b in
                    Envaux.report_error fmt e;
                    Format.pp_print_flush fmt () ;
                    Log.err (fun m -> m "%s" (Buffer.contents b));
                    exp.exp_env
              in
              let desc = Env.find_value path env in
              let dloc = desc.Types.val_loc in
              if dloc.Location.loc_ghost then None
              else Some dloc
            with
              Not_found -> None
          in
          idents := (path, loc, defloc) :: !idents
      | _ ->
          ()
    in
    defmap.Untypeast.expr mapper exp
  in
  let mapper = { defmap with expr } in
  List.iter (apply_mapper mapper) parts ;
  let l = List.rev !idents in
  List.fold_left f acc l
