Source file utils.ml

open Import

let with_output fn ~binary ~f =
  match fn with
  | None | Some "-" -> f stdout
  | Some fn -> Out_channel.with_file fn ~binary ~f
;;

module Kind = struct
  type t = Intf | Impl

  let of_filename fn : t option =
    if Caml.Filename.check_suffix fn ".ml" then
      Some Impl
    else if Caml.Filename.check_suffix fn ".mli" then
      Some Intf
    else
      None
  ;;

  let describe = function
    | Impl -> "implementation"
    | Intf -> "interface"
  ;;

  let equal : t -> t -> bool = Poly.equal
end

module Intf_or_impl = struct
  type t =
    | Intf of signature
    | Impl of structure

  let map t (map : Ast_traverse.map) =
    match t with
    | Impl x -> Impl (map#structure x)
    | Intf x -> Intf (map#signature x)
  ;;

  let map_with_context t (map : _ Ast_traverse.map_with_context) ctx =
    match t with
    | Impl x -> Impl (map#structure ctx x)
    | Intf x -> Intf (map#signature ctx x)
  ;;

  let kind : _ -> Kind.t = function
    | Intf _ -> Intf
    | Impl _ -> Impl
end

module Ast_io = struct
  type input_version = (module OCaml_version)

  let fall_back_input_version = (module Compiler_version : OCaml_version)
  (* This should only be used when the input version can't be determined due to
      loading or preprocessing errors *)

  type t = {
    input_name : string;
    input_version : input_version;
    ast : Intf_or_impl.t;
  }

  type read_error =
    | Not_a_binary_ast
    | Unknown_version of string * input_version
    | Source_parse_error of Location.Error.t * input_version
    | System_error of Location.Error.t * input_version

  type input_source = Stdin | File of string

  type input_kind =
    | Possibly_source of Kind.t * string
    | Necessarily_binary

  let parse_source_code ~(kind : Kind.t) ~input_name ~prefix_read_from_source ic =
    (* The input version is determined by the fact that the input will get parsed by
       the current compiler Parse module *)
    let input_version = (module Compiler_version : OCaml_version) in
    try
    (* To test if a file is an AST file, we have to read the first few bytes of the
        file. If it is not, we have to parse these bytes and the rest of the file as
        source code.

        The compiler just does [seek_on 0] in this case, however this doesn't work when
        the input is a pipe.

        What we do instead is create a lexing buffer from the input channel and pre-fill
        it with what we read to do the test. *)
      let lexbuf = Lexing.from_channel ic in
      let len = String.length prefix_read_from_source in
      Bytes.blit_string ~src:prefix_read_from_source ~src_pos:0 ~dst:lexbuf.lex_buffer ~dst_pos:0
        ~len;
      lexbuf.lex_buffer_len <- len;
      lexbuf.lex_curr_p <-
        { pos_fname = input_name
        ; pos_lnum  = 1
        ; pos_bol   = 0
        ; pos_cnum  = 0
        };
      Lexer.skip_hash_bang lexbuf;
      let ast : Intf_or_impl.t =
        match kind with
        | Intf -> Intf (Parse.interface      lexbuf)
        | Impl -> Impl (Parse.implementation lexbuf)
      in Ok {input_name; input_version; ast}
    with exn -> (
      match Location.Error.of_exn exn with
      | None -> raise exn
      | Some error -> Error (Source_parse_error (error, input_version)))

  let magic_length = String.length Ocaml_common.Config.ast_impl_magic_number

  let read_magic ic =
    let buf = Bytes.create magic_length in
    let len = input ic buf 0 magic_length in
    let s = Bytes.sub_string buf ~pos:0 ~len in
    if len = magic_length then
      Ok s
    else
      Error s

  let from_channel ch ~input_kind =
    let handle_non_binary prefix_read_from_source =
      match input_kind with
      | Possibly_source (kind, input_name) ->
          parse_source_code ~kind ~input_name ~prefix_read_from_source ch
      | Necessarily_binary -> Error Not_a_binary_ast
    in
    match read_magic ch with Error s -> handle_non_binary s
    | Ok s -> (
        match Find_version.from_magic s with
        | Intf (module Input_version : OCaml_version) ->
            let input_name : string = input_value ch in
            let ast = input_value ch in
            let module Input_to_ppxlib = Convert (Input_version) (Js) in
            let ast = Intf_or_impl.Intf (Input_to_ppxlib.copy_signature ast) in
            Ok
              {
                input_name;
                input_version = (module Input_version : OCaml_version);
                ast;
              }
        | Impl (module Input_version : OCaml_version) ->
            let input_name : string = input_value ch in
            let ast = input_value ch in
            let module Input_to_ppxlib = Convert (Input_version) (Js) in
            let ast = Intf_or_impl.Impl (Input_to_ppxlib.copy_structure ast) in
            Ok
              {
                input_name;
                input_version = (module Input_version : OCaml_version);
                ast;
              }
        | Unknown ->
          if
            String.equal
              (String.sub s ~pos:0 ~len:9)
              (String.sub Ocaml_common.Config.ast_impl_magic_number ~pos:0
                ~len:9)
            || String.equal
                (String.sub s ~pos:0 ~len:9)
                (String.sub Ocaml_common.Config.ast_intf_magic_number ~pos:0
                    ~len:9)
            then Error (Unknown_version (s, fall_back_input_version))
            else (handle_non_binary s))

  let read input_source ~input_kind =
    try
      match input_source with
        | Stdin -> from_channel stdin ~input_kind
        | File fn -> In_channel.with_file fn ~f:(from_channel ~input_kind)
    with exn ->
      match Location.Error.of_exn exn with
        | None -> raise exn
        | Some error ->
      Error (System_error (error, fall_back_input_version))



  let write oc { input_name; input_version = (module Input_version); ast }
      ~add_ppx_context =
    let module Ppxlib_to_input = Convert (Js) (Input_version) in
    let module Ocaml_to_input = Convert (Compiler_version) (Input_version) in
    match ast with
    | Intf sg ->
        let sg =
          if add_ppx_context then
            Selected_ast.To_ocaml.copy_signature sg
            |> Ocaml_common.Ast_mapper.add_ppx_context_sig ~tool_name:"ppx_driver"
            |> Ocaml_to_input.copy_signature
          else Ppxlib_to_input.copy_signature sg
        in
        output_string oc Input_version.Ast.Config.ast_intf_magic_number;
        output_value oc input_name;
        output_value oc sg
    | Impl st ->
        let st =
          if add_ppx_context then
            Selected_ast.To_ocaml.copy_structure st
            |> Ocaml_common.Ast_mapper.add_ppx_context_str ~tool_name:"ppx_driver"
            |> Ocaml_to_input.copy_structure
          else Ppxlib_to_input.copy_structure st
        in
        output_string oc Input_version.Ast.Config.ast_impl_magic_number;
        output_value oc input_name;
        output_value oc st
end

module System = struct
  let run_preprocessor ~pp ~input ~output =
    let command =
      Printf.sprintf "%s %s > %s"
        pp (if String.equal input "-" then "" else Caml.Filename.quote input)
        (Caml.Filename.quote output)
    in
  if Caml.Sys.command command = 0 then
    Ok ()
  else Error (command, Ast_io.fall_back_input_version)
end