(* A utility to gather information from caml compiled interface files

   Copyright (C) 2007 Eric Stokes

   This library is free software; you can redistribute it and/or
   modify it under the terms of the GNU General Public License as
   published by the Free Software Foundation; either version 2.1 of
   the License, or (at your option) any later version.
   
   This library 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           
   Lesser General Public License for more details.                             
   
   You should have received a copy of the GNU General Public License
   along with this library; if not, write to the Free Software
   Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA 02111-1307
   USA
*)

open Types

exception Break

module Pathset = Set.Make(struct
  type t = string
  let trailing_slash = Pcre.regexp "/\\s*$|\\\\\\s*$"
  let compare p1 p2 = 
    let p1' = Pcre.replace ~rex:trailing_slash ~templ:"" p1 in
    let p2' = Pcre.replace ~rex:trailing_slash ~templ:"" p2 in
    String.compare p1' p2'
end)

(* the standard library should not be so deficient *)
module List = struct
  include List

  let filter_map f l = 
    List.fold_left
      (fun acc item -> 
         match f item with
         | Some x -> x :: acc
         | None -> acc)
      []
      l

  let find_map f l = 
    let res = ref None in
    try
      List.iter
        (fun x -> 
           match f x with
           | None -> ()
           | Some y -> res := Some y; raise Break)
        l;
      raise Not_found
    with Break -> 
      begin match !res with
      | Some y -> y
      | None -> raise Break
      end

  let map f l = 
    let r = rev_map f l in
    rev r
end

module Unix = struct
  include Unix

  let fold_path ~f ~init path = 
    let dir = Unix.opendir path in
    let acc = ref init in
    try
      while true do
        let file = Unix.readdir dir in
        acc := f file !acc
      done;
      !acc
    with 
    | End_of_file -> Unix.closedir dir; !acc
    | exn -> Unix.closedir dir; raise exn
end

module Match = struct
  let comma = Pcre.regexp "\\s*,\\s*"
end

module Module_expression = struct
  type t = 
  | Exact of string
  | Begins_with of string * Pcre.regexp
  | Ends_with of string * Pcre.regexp
  | Begins_and_ends of string * Pcre.regexp
  | Contains of string * Pcre.regexp
  | Any

  let to_string = function
    | Exact m -> m
    | Begins_with (s, _) 
    | Ends_with (s, _)
    | Begins_and_ends (s, _)
    | Contains (s, _) -> s
    | Any -> "*"

  (* ModA,ModB,Foo*.Make *)
  let parse = 
    let dot = Pcre.regexp "\\." in
    let capname = Pcre.regexp "^[A-Z][A-Za-z_'0-9]*$" in
    let starend = Pcre.regexp "^[A-Z][A-Za-z0-9_']*\\*$" in
    let starbegin = Pcre.regexp "^\\*[A-Za-z0-9_']+$" in
    let starboth = Pcre.regexp "^\\*[A-Za-z0-9_']+\\*$" in
    let starmiddle = Pcre.regexp "^([A-Z][A-Za-z0-9_']*)\\*([A-Za-z0-9_']+)$" in
    let star = Pcre.regexp "\\*" in
    fun exp ->
      List.map
        (fun token ->
          let token_no_star = Pcre.replace ~rex:star ~templ:"" token in
          if token = "*" then
            Any
          else if Pcre.pmatch ~rex:capname token then
            Exact token
          else if Pcre.pmatch ~rex:starboth token then
            Contains (token, Pcre.regexp ("^.*" ^ token_no_star ^ ".*$"))
          else if Pcre.pmatch ~rex:starmiddle token then
            begin match Pcre.extract ~rex:starmiddle token with
            | [|_whole; begins; ends |] -> 
                let rex = Pcre.regexp (Printf.sprintf "^%s.*%s$" begins ends) in
                Begins_and_ends (token, rex)
            | _ -> failwith "invalid begins and ends with match"
            end
          else if Pcre.pmatch ~rex:starbegin token then
            Ends_with (token, Pcre.regexp (Printf.sprintf "%s$" token_no_star))
          else if Pcre.pmatch ~rex:starend token then
            Begins_with (token, Pcre.regexp (Printf.sprintf "^%s" token_no_star))
          else
            failwith "invalid module expression")
        (Pcre.split ~rex:dot exp)

  let is_exact t = 
    List.for_all 
      (function 
        | Exact _ -> true 
        | Begins_with _
        | Ends_with _
        | Begins_and_ends _
        | Contains _
        | Any -> false)
      t

  let parse_exact exp = 
    let t = parse exp in
    if is_exact t then
      t
    else
      failwith "Module_expression.parse_exact: expression is not exact!"
end

type mode = 
  | Find_type of Pcre.regexp
  | Find_constructor of Pcre.regexp
  | Find_polymorphic_variant of Pcre.regexp
  | Find_record_label of Pcre.regexp
  | Find_value of Pcre.regexp
  | Find_exception of Pcre.regexp
  | Find_module
  | Find_class of Pcre.regexp
  | Find_all of Pcre.regexp

type module_tree = 
  | Leaf of string * signature
  | Node of string * signature * module_tree list

type args = {
  mode: mode;
  path: Pathset.t;
  context: Module_expression.t list list; (* open modules *)
  modname: Module_expression.t list list;
}

let parse_args () =
  let module Parse = struct
    open Arg
    let mode = ref None
    let path = ref (Pathset.add "." (Pathset.singleton Config.standard_library))
    let context = ref ["Pervasives"]
    let modname = ref []

    let set_mode m =
      match !mode with
      | None -> mode := Some m;
      | Some _ -> raise (Invalid_argument "the mode is already set")

    let add_packages p = 
      Findlib.init ();
      let packages = Pcre.split ~rex:Match.comma p in
      List.iter 
        (fun package ->
          try
            let dir = Findlib.package_directory package in
            path := Pathset.add dir !path
          with exn -> 
            Printf.eprintf "warning, error finding package dir: %s\n" (Printexc.to_string exn))
        packages

    let add_opens s = context := Pcre.split ~rex:Match.comma s

    let args = 
      Arg.align
        [("-t", String (fun s -> set_mode (Find_type (Pcre.regexp s))),
          "      (regexp) print types with matching names");
         ("-r", String (fun s -> set_mode (Find_record_label (Pcre.regexp s))),
          "      (regexp) print record field labels with matching names");
         ("-c", String (fun s -> set_mode (Find_constructor (Pcre.regexp s))), 
          "      (regexp) print constructors with matching names");
         ("-p", String (fun s -> 
                         set_mode 
                           (Find_polymorphic_variant 
                             (Pcre.regexp s))),
          "      (regexp) print polymorphic variants with matching names");
         ("-m", Unit (fun () -> set_mode Find_module),
          "      (regexp) print all matching module names in the path");
         ("-v", String (fun s -> set_mode (Find_value (Pcre.regexp s))), 
          "      (regexp) print values with matching names");
         ("-e", String (fun s -> set_mode (Find_exception (Pcre.regexp s))), 
          "      (regexp) print exceptions with matching constructors");
         ("-o", String (fun s -> set_mode (Find_class (Pcre.regexp s))),
          "      (regexp) print all classes with matching names");
         ("-a", String (fun s -> set_mode (Find_all (Pcre.regexp s))),
          "      (regexp) print all names which match the given expression");
         ("-I", String (fun s -> path := Pathset.add s !path), 
          "      (directory) add additional directory to the search path");
         ("-package", String (fun s -> add_packages s),
          "      (packages) comma seperated list of findlib packages to search");
         ("-open", String (fun s -> add_opens s), 
          "      (modules) comma seperated list of open modules (in order!)")]
    let usage = 
      Printf.sprintf
        ("%s: <args> <module-expr> \n" ^^
           "extract information from caml compiled interface files\n" ^^
           " <module-expr> can be an exact module name, " ^^
           " or a shell wildcard. Multiple modules can be specified " ^^
           "E.G. \"ModA ModB Foo*.Make\" means to search ModA, ModB, and " ^^
           "any submodule Make of a module that starts with Foo.")
        Sys.argv.(0)

    let parse () = 
      Arg.parse args
        (fun anon -> modname := (Module_expression.parse anon) :: !modname)
        usage

    let error msg = 
      prerr_endline msg;
      Arg.usage args usage;
      exit 1
  end
  in
  Parse.parse ();
  let mode = 
    match !Parse.mode with
    | Some m -> m
    | None -> Parse.error "you must specify a search mode"
  in
  {mode = mode;
   path = 
      if Pathset.is_empty !Parse.path then Parse.error "you must specify a search path"
      else !Parse.path;
   context = 
      List.map
        Module_expression.parse_exact
        !Parse.context;
   modname = 
      (match !Parse.modname with
       | [] -> 
           if !Parse.context = [] then
             Parse.error "you must specify a module expression, or a list of open modules"
           else 
             []
       | name -> name)}

let match_ident exp id = Pcre.pmatch ~rex:exp (Ident.name id)

let whsp = Pcre.regexp ~study:true "\\s+|$"

let print_type print_path path s exp =
  List.iter
    (function 
       | Tsig_type (id, type_decl, rec_status) ->
           if match_ident exp id then begin
             Printtyp.type_declaration id Format.std_formatter type_decl;
             if print_path then
               Format.print_string (Printf.sprintf " (* %s *)" path);
             Format.print_newline ()
           end
       | _ -> ())
    s

let print_constructor print_path path s exp =
  let type_expr_to_string exp = 
    Printtyp.type_expr Format.str_formatter exp;
    Format.flush_str_formatter ();
  in
  List.iter
    (function
       | Tsig_type (id, type_decl, _rec_status) ->
           begin match type_decl.type_kind with
           | Type_variant (constructors) ->
               List.iter
                 (fun (name, type_exprs) ->
                    if Pcre.pmatch ~rex:exp name then begin
                      Format.print_string name;
                      if type_exprs <> [] then begin
                        Format.print_string " of ";
                        Format.print_string
                          (String.concat " * "
                             (List.map
                                (fun e -> type_expr_to_string e)
                                type_exprs))
                      end;
                      Format.print_string " (* ";
                      if print_path then
                        Format.print_string (path ^ ".");
                      Format.print_string (Ident.name id);
                      Format.print_string " *)";
                      Format.print_newline ()
                    end)
                 constructors
           | _ -> ()
           end
       | _ -> ())
    s

let print_polymorphic_variant print_path path s expr = 
  let print_if_polymorphic_variant id type_decl =
    begin match type_decl.type_manifest with
    | None -> ()
    | Some {desc = type_descr} ->
        begin match type_descr with
        | Tvariant variant_descr ->
            List.iter
              (fun (name, param) ->
                 let src_name = "`" ^ name in
                 if Pcre.pmatch ~rex:expr src_name then begin
                   Format.print_string src_name;
                   begin match param with
                   | Rpresent None -> ()
                   | Rabsent -> ()
                   | Reither _ -> () (* this can't happen in a type *)
                   | Rpresent (Some type_expr) ->
                       Format.print_string " of ";
                       Printtyp.type_expr 
                         Format.str_formatter type_expr;
                       let s = 
                         Pcre.replace ~rex:whsp ~templ:" "
                           (Format.flush_str_formatter ()) 
                       in
                       Format.print_string s;
                   end;
                   Format.print_string 
                     (Printf.sprintf " (* %s%s *)"
                        (if print_path then (path ^ ".") else "")
                        (Ident.name id));
                   Format.print_newline ()
                 end)
              variant_descr.row_fields
        | _ -> ()
        end
    end
  in
  List.iter
    (function 
       | Tsig_type (id, type_decl, _rec_status) ->
           begin match type_decl.type_kind with
           | Type_abstract -> print_if_polymorphic_variant id type_decl
           | _ -> ()
           end
       | _ -> ())
    s

let print_record_label print_path path s exp =
  List.iter
    (function
       | Tsig_type (id, type_decl, _rec_status) ->
           begin match type_decl.type_kind with
           | Type_record (labels, _) ->
               List.iter
                 (fun (name, mutable_flag, type_expr) ->
                    if Pcre.pmatch ~rex:exp name then begin
                      begin match mutable_flag with
                      | Asttypes.Mutable -> Format.print_string "mutable "
                      | Asttypes.Immutable -> ()
                      end;
                      Format.print_string name;
                      Format.print_string ": ";
                      Printtyp.type_expr Format.std_formatter type_expr;
                      Format.print_string " (* ";
                      if print_path then
                        Format.print_string (path ^ ".");
                      Format.print_string (Ident.name id);
                      Format.print_string " *)";
                      Format.print_newline ()
                    end)
                 labels
           | _ -> ()
           end
       | _ -> ())
    s

let print_value print_path path s exp =
  List.iter
    (function
       | Tsig_value (id, desc) ->
           if match_ident exp id then begin
             Printtyp.value_description id Format.str_formatter desc;
             let s = 
               Pcre.replace ~rex:whsp ~templ:" "
                 (Format.flush_str_formatter ()) 
             in
             if print_path then
               print_endline (s ^ (Printf.sprintf " (* %s *)" path))
             else
               print_endline s
           end
       | _ -> ())
    s

let print_class print_path path s exp =
  List.iter
    (function
       | Tsig_class (id, cd, _) when match_ident exp id ->
           Printtyp.class_declaration id Format.std_formatter cd;
           if print_path then
             Format.print_string (Printf.sprintf " (* %s *)" path);
           Format.print_newline ()
       | Tsig_cltype (id, ct, _) when match_ident exp id ->
           Printtyp.cltype_declaration id Format.std_formatter ct;
           if print_path then
             Format.print_string (Printf.sprintf " (* %s *)" path);
           Format.print_newline ()
       | _ -> ())
    s

let print_all print_path path s exp = 
  let new_s = 
    List.filter 
      (function
         | Tsig_value (id, _)
         | Tsig_type (id, _, _)
         | Tsig_exception (id, _)
         | Tsig_module (id, _, _)
         | Tsig_modtype (id, _)
         | Tsig_class (id, _, _)
         | Tsig_cltype (id, _, _) ->
             match_ident exp id)
      s
  in
  Printtyp.signature Format.std_formatter new_s;
  if print_path then
    Format.print_string (Printf.sprintf " (* %s *)" path);
  Format.print_newline ()

let print_exception print_path path s exp = 
  List.iter
    (function 
       | Tsig_exception (id, exn) ->
           if match_ident exp id then begin
             Printtyp.exception_declaration id Format.std_formatter exn;
             if print_path then
               Format.print_string (Printf.sprintf " (* %s *)" path);
             Format.print_newline ()
           end
       | _ -> ())
    s

let warn_env_error e =
  Env.report_error Format.str_formatter e;
  let e = Format.flush_str_formatter () in
  Printf.eprintf "%s\n%!" e

let match_mod_expr expr mod_name = 
  let module E = Module_expression in
  match expr with
  | E.Exact name -> mod_name = name
  | E.Begins_with (_, rex)
  | E.Ends_with (_, rex)
  | E.Begins_and_ends (_, rex)
  | E.Contains (_, rex) -> Pcre.pmatch ~rex mod_name
  | E.Any -> true

let cmi_file = Pcre.regexp "\\.cmi$"
let modname_of_cmi f = 
  String.capitalize (Pcre.replace ~templ:"" ~rex:cmi_file f)

let cmi_of_modname n = (String.lowercase n) ^ ".cmi"

let cmi_files args mod_expr = 
  let module E = Module_expression in
  match mod_expr with
  | E.Exact mod_name ->
      let cmi_name = cmi_of_modname mod_name in
      Pathset.fold
        (fun path acc -> 
           if Sys.file_exists (Filename.concat path cmi_name) then
             (mod_name, Filename.concat path cmi_name) :: acc
           else
             acc)
        args.path
        []
  | _ ->
      Pathset.fold
        (fun path cmi_files ->
           Unix.fold_path
             ~f:(fun file cmi_files -> 
                   if Pcre.pmatch ~rex:cmi_file file then begin
                     let mod_name = modname_of_cmi file in
                     if match_mod_expr mod_expr mod_name then 
                       (mod_name, Filename.concat path file) :: cmi_files
                     else
                       cmi_files
                   end else
                     cmi_files)
             ~init:cmi_files
             path)
        args.path
        []

let rec matching_submods mod_expr s =
  match s with
  | Tsig_module (id, mt, _) :: tl when match_mod_expr mod_expr (Ident.name id) ->
      begin match mt with
      | Tmty_signature sg -> (Ident.name id, sg) :: matching_submods mod_expr tl
      | Tmty_functor (_, mt, _) ->
          begin match mt with
          | Tmty_signature sg -> (Ident.name id, sg) :: matching_submods mod_expr tl
          | _ -> matching_submods mod_expr tl
          end
      | Tmty_ident _ -> matching_submods mod_expr tl
      end
  | _ :: tl -> matching_submods mod_expr tl
  | [] -> []

let rec build_module_tree name root mod_expr = 
  match mod_expr with
  | [] -> Leaf (name, root)
  | mod_expr :: tl ->
      begin match matching_submods mod_expr root with
      | [] -> Leaf (name, root)
      | mods -> 
          let children = 
            List.map 
              (fun (name, sg) -> build_module_tree name sg tl)
              mods
          in
          Node (name, root, children)
      end

let rec extract_nodes depth path modtree =
  let concatpath path name = 
    if path = "" then name
    else (path ^ "." ^ name)
  in
  match modtree with
  | Leaf (name, sg) -> [(concatpath path name, depth, sg)]
  | Node (name, sg, children) ->
      (concatpath path name, depth, sg) ::
        (List.flatten
           (List.map
              (fun submod -> extract_nodes (depth + 1) (concatpath path name) submod)
              children))

let print_requested_stuff print_path name s args =
  match args.mode with
  | Find_type e -> print_type print_path name s e
  | Find_constructor e -> print_constructor print_path name s e
  | Find_polymorphic_variant e -> print_polymorphic_variant print_path name s e
  | Find_record_label e -> print_record_label print_path name s e
  | Find_value e -> print_value print_path name s e
  | Find_exception e -> print_exception print_path name s e
  | Find_class e -> print_class print_path name s e
  | Find_all e -> print_all print_path name s e
  | Find_module -> 
      Format.print_string name;
      Format.print_newline ()

let read_cmi_file filename =
  let ic = open_in_bin filename in
  try
    let buffer = String.create (String.length Config.cmi_magic_number) in
    really_input ic buffer 0 (String.length Config.cmi_magic_number);
    if buffer <> Config.cmi_magic_number then begin
      close_in ic;
      failwith (Printf.sprintf "not an interface: %s" filename)
    end;
    let (name, sg) = input_value ic in
    close_in ic;
    sg
  with exn ->
    close_in ic;
    failwith 
      (Printf.sprintf 
         "bad cmi file: %s, error: %s"
         filename
         (Printexc.to_string exn))

let module_exists args mod_exp = 
  let expr_len = List.length mod_exp in
  let mod_name = List.hd mod_exp in
  let submods = List.tl mod_exp in
  match cmi_files args mod_name with
  | [] -> false
  | cmi_files ->
      List.exists
        (fun (name, cmi_file) ->
           let s = read_cmi_file cmi_file in
           let sgs = 
             List.filter
               (fun (_, depth, _) -> depth = expr_len)
               (extract_nodes 1 "" (build_module_tree name s submods))
           in
           List.length sgs > 0)
        cmi_files

let gen_qualified args context = 
  let context = Array.of_list context in
  for i = 0 to Array.length context - 1 do
    try
      for j = i downto 0 do
        let maybe_parent = context.(j) in
        let child = context.(i) in
        let qualified = maybe_parent @ child in
        if module_exists args qualified then begin
          context.(i) <- qualified;
          raise Break
        end
      done
    with Break -> ()
  done;
  Array.to_list context

let () = 
  let args = parse_args () in
  let qualified_context = gen_qualified args args.context in
  let mod_exprs = 
    (* combine the list of fully qualified open modules with the
       module expressions that the user has specified following the
       compiler's rules about module opens for exact expressions, and
       combining everything for non exact expressions. *)
    match args.modname with
    | [] -> qualified_context
    | exps ->
        List.flatten
          (List.map
             (fun exp -> 
                if Module_expression.is_exact exp then
                  [try 
                     List.find_map
                       (fun qual -> 
                          let exp' = qual @ exp in
                          if module_exists args exp' then Some exp'
                          else None)
                       (List.rev qualified_context) (* look from the bottom up *)
                   with Not_found -> exp]
                else
                  exp ::
                    (List.rev_map
                       (fun qual -> qual @ exp)
                       qualified_context))
             exps)
  in
  List.iter
    (fun mod_expr ->
       try
         let expr_len = List.length mod_expr in
         let mod_name = List.hd mod_expr in
         let submods = List.tl mod_expr in
         let cmi_files = cmi_files args mod_name in
         List.iter
           (fun (name, cmi_file) ->
              let s = read_cmi_file cmi_file in
              let sgs = 
                List.filter
                  (fun (_, depth, _) -> depth = expr_len)
                  (extract_nodes 1 "" (build_module_tree name s submods))
              in
              let print_path = 
                List.length sgs > 1 || 
                  List.length cmi_files > 1 ||
                  List.length args.modname > 1
              in
              List.iter
                (fun (name, _,  sg) -> print_requested_stuff print_path name sg args)
                sgs)
           cmi_files
       with exn -> 
         Printf.eprintf
           "failed to operate on: \"%s\", %s\n%!"
           (String.concat " " (List.map Module_expression.to_string mod_expr))
           (Printexc.to_string exn))
    mod_exprs