Source file ast_io.ml

(* Copyright (C) 2023 Antonio Nuno Monteiro
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Lesser General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * In addition to the permissions granted to you by the LGPL, you may combine
 * or link a "work that uses the Library" with a publicly distributed version
 * of this file to produce a combined library or application, then distribute
 * that combined work under the terms of your choosing, with no requirement
 * to comply with the obligations normally placed on you by section 4 of the
 * LGPL version 3 (or the corresponding section of a later version of the LGPL
 * should you choose to use a later version).
 *
 * 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 Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser 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. *)

module Ast = Ppxlib_ast.Ast
module Compiler_version = Ppxlib_ast.Compiler_version

module type OCaml_version = Ppxlib_ast.OCaml_version

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

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 Ppxlib_ast.Location_error.t * input_version
  | System_error of Ppxlib_ast.Location_error.t * input_version

type input_source = Stdin | File of string

type input_kind =
  | Possibly_source of {
      filename : string;
      parse_fun : Lexing.lexbuf -> Intf_or_impl.t;
    }
  | Necessarily_binary

let read_error_to_string (error : read_error) =
  match error with
  | Not_a_binary_ast -> "Error: Not a binary ast"
  | Unknown_version (s, _) -> "Error: Unknown version " ^ s
  | Source_parse_error (loc, _) ->
      "Source parse error:" ^ Ppxlib_ast.Location_error.message loc
  | System_error (loc, _) ->
      "System error: " ^ Ppxlib_ast.Location_error.message loc

let magic_length = String.length Astlib.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 0 len in
  if len = magic_length then Ok s else Error s

let from_channel ch ~input_kind : (t, read_error) result =
  let input_version = (module Compiler_version : OCaml_version) in
  let module Convert = Ppxlib_ast.Convert in
  let module Find_version = Ppxlib_ast__.Versions.Find_version in
  let module Js = Ppxlib_ast__.Import.Js in
  let handle_non_binary () =
    match input_kind with
    | Possibly_source { filename; parse_fun } ->
        seek_in ch 0;
        let lexbuf = Lexing.from_channel ch in
        Location.init lexbuf filename;
        let ast = parse_fun lexbuf in
        Ok { input_name = filename; input_version; ast }
    | Necessarily_binary -> Error Not_a_binary_ast
  in
  match read_magic ch with
  | Error _ -> handle_non_binary ()
  | Ok s -> (
      match Find_version.from_magic s with
      | Intf (module Input_version : OCaml_version) ->
          let input_name : string = input_value ch in
          Location.set_input_name input_name;
          let ast = input_value ch in
          let module Input_to_ppxlib = Convert (Input_version) (Js) in
          Ok
            {
              input_name;
              input_version = (module Input_version : OCaml_version);
              ast =
                Intf
                  (Input_to_ppxlib.copy_signature ast
                  |> Melange_ppxlib_ast.Of_ppxlib.copy_signature);
            }
      | Impl (module Input_version : OCaml_version) ->
          let input_name : string = input_value ch in
          Location.set_input_name input_name;
          let ast = input_value ch in
          let module Input_to_ppxlib = Convert (Input_version) (Js) in
          Ok
            {
              input_name;
              input_version = (module Input_version : OCaml_version);
              ast =
                Impl
                  (Input_to_ppxlib.copy_structure ast
                  |> Melange_ppxlib_ast.Of_ppxlib.copy_structure);
            }
      | Unknown ->
          if
            String.equal (String.sub s 0 9)
              (String.sub Astlib.Config.ast_impl_magic_number 0 9)
            || String.equal (String.sub s 0 9)
                 (String.sub Astlib.Config.ast_intf_magic_number 0 9)
          then Error (Unknown_version (s, fall_back_input_version))
          else handle_non_binary ())

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

let read_exn input_source ~input_kind =
  match read input_source ~input_kind with
  | Ok ret -> ret
  | Error e -> raise (Arg.Bad (read_error_to_string e))