open! Core

(* The [generated_build_info ()] string can be:

   - The empty string, from the default C function. This indicates that the build system
     does not support build info for this executable, e.g. because it's a test runner.

   - A printed sexp, from the C function generated by the build system, which is parsable
     by the [Build_info.t] type below. *)
external generated_build_info : unit -> string = "generated_build_info"

(* The [generated_hg_version ()] can be:

   - The empty string, from the default C function. This indicates that the build system
     does not support version info for this executable, e.g. because it's a test runner.

   - "NO_VERSION_UTIL" from the C function generated by the build system; or

   - "repo1 rev40\nrepo2 rev40" from the C function generated by the build system.

   The last two are prefixed by the [Version_util.start_marker]. When the prefix is
   present, there is also enough padding for ~4K worth of data. This allows us to rewrite
   binaries to insert the version util after linking.
*)
external generated_hg_version : unit -> string = "generated_hg_version"

(** Make sure to update [bin/generate_static_string_c_code.sh] too if you are changing
    these constants. *)
module Version_util_section = Section.Make (struct
  let name = "version_util"
  let length_including_start_marker = 4096

  (* BEFORE CHANGING: Please note the bidirectional version compatibility guarantee
       granted in the mli file for [Expert.get_version_util]. If we ever need to change the
       version util format, we should update the code to be able to read both the old and
       new formats, then wait a month, then change the write function to only write the new
       format. The old and new formats can be distinguished by minting a new [start_marker]
       for the new format.
    *)
  let start_marker =
    (* This trick is to prevent the marker from occurring verbatim in the binary that uses
         [Expert.insert_version_util], so that we don't by accident rewrite our own code.

         [opaque_identity] is used to prevent the compiler from converting this computation
         into a literal, thus undoing this trick. We could split the marker in half instead,
         but that would make grepping hard for humans.
         Grep in the tree to see the place that generates this. *)
    (Sys.opaque_identity ( ^ )) "rUb71QgfHXXwnBWBoJfb0Sa3R60vihdV" ":"
  ;;
end)

(** Make sure to update [bin/generate_static_string_c_code.sh] too if you are changing
    these constants. *)
module Build_info_section = Section.Make (struct
  let name = "build info"
  let length_including_start_marker = 4096

  let start_marker =
    (* Same trick as in [Version_util_section]. *)
    (Sys.opaque_identity ( ^ )) "vNxXpiccvPI9MHVFJuNwNxj8eu9W5KCB" ":"
  ;;
end)

(* BEFORE CHANGING: Note version compatibility guarantee above. *)
let parse_generated_hg_version = function
  | "" -> [ "NO_VERSION_UTIL" ]
  | generated_hg_version ->
    generated_hg_version
    |> String.chop_suffix_if_exists ~suffix:"\n"
    |> Version_util_section.chop_start_marker_if_exists
    |> String.split ~on:'\n'
    |> List.map ~f:(fun line ->
         match String.rsplit2 line ~on:' ' with
         | None -> line (* no version util *)
         | Some (repo, rev_status) ->
           (* For compability with downstream tools that might rely on this output format,
           and with [Version.parse].*)
           String.concat
             [ repo
             ; "_"
             ; String.prefix rev_status 12
             ; (* The revision can have a one-character '+' suffix. Keep it. *)
               (if String.length rev_status mod 2 = 1
                then String.suffix rev_status 1
                else "")
             ])
;;

let version_list = parse_generated_hg_version (generated_hg_version ())
let version = String.concat version_list ~sep:" "

module Version = struct
  type t =
    { repo : string
    ; version : string
    }
  [@@deriving compare, sexp_of]

  let parse1 version =
    match String.rsplit2 version ~on:'_' with
    | None -> error_s [%message "Could not parse version" version]
    | Some (repo, version) -> Ok { repo; version }
  ;;

  let parse_list l =
    (* We might get multiple such lines if we have multiple repos *)
    if List.exists l ~f:(String.( = ) "NO_VERSION_UTIL")
    then Ok None
    else
      List.map l ~f:parse1 |> Or_error.combine_errors |> Or_error.map ~f:(fun x -> Some x)
  ;;

  let parse_lines versions = parse_list (String.split_lines versions)
  let current_version () = ok_exn (parse_list version_list)

  let present = function
    | None -> error_s [%sexp "executable built without version util"]
    | Some x -> Ok x
  ;;

  let parse_list_present x = Or_error.bind ~f:present (parse_list x)
  let parse_lines_present x = Or_error.bind ~f:present (parse_lines x)
  let current_version_present () = present (current_version ())
end

module Expert = struct
  (* BEFORE CHANGING: Please note the bidirectional version compatibility guarantee
     granted in the mli file. *)
  let get_version_util ~contents_of_exe =
    let%map.Option section = Version_util_section.get ~contents_of_exe in
    String.concat ~sep:" " (parse_generated_hg_version section)
  ;;

  let text (versions_opt : Version.t list option) =
    match versions_opt with
    | None -> "NO_VERSION_UTIL"
    | Some versions ->
      if List.is_empty versions
      then failwith "version_util must include at least one repository";
      if List.contains_dup
           ~compare:String.compare
           (List.map versions ~f:(fun v -> v.repo))
      then failwith "version_util must not contain duplicate repositories";
      versions
      |> List.sort ~compare:Version.compare
      |> List.map ~f:(fun { repo; version } ->
           if not (String.mem repo '/')
           then failwith [%string "%{repo} doesn't look like a repo url"];
           (let version' = String.chop_suffix_if_exists version ~suffix:"+" in
            if (String.length version' = 40 || String.length version' = 64)
               && String.for_all version' ~f:Char.is_hex_digit_lower
            then ()
            else failwith [%string "%{version} doesn't look like a full hg version"]);
           repo ^ " " ^ version ^ "\n")
      |> String.concat
  ;;

  let raw_text v = Version_util_section.Expert.pad_with_at_least_one_nul_byte_exn (text v)

  let replace_version_util ~contents_of_exe versions_opt =
    (* Expert because we don't really want people to casually use this, so its contents
       can be trusted. *)
    Version_util_section.replace ~contents_of_exe ~data:(text versions_opt)
  ;;

  let parse_generated_hg_version = parse_generated_hg_version

  module Experimental = struct
    let get_build_info = Build_info_section.get

    let remove_build_info ~contents_of_exe =
      Build_info_section.replace ~contents_of_exe ~data:"NO_BUILD_INFO"
    ;;
  end

  module For_tests = struct
    let count_section_occurrences ~contents_of_exe =
      Version_util_section.count_occurrences ~contents_of_exe
      + Build_info_section.count_occurrences ~contents_of_exe
    ;;
  end
end

module Build_info = struct
  module Application_specific_fields = struct
    type t = Sexp.t String.Map.t [@@deriving sexp]
  end

  module Time_with_limited_parsing = struct
    type t = Time_float.t * Sexp.t

    let t_of_sexp sexp =
      let str = string_of_sexp sexp in
      try
        match String.chop_suffix str ~suffix:"Z" with
        | None -> failwith "zone must be Z"
        | Some rest ->
          (match String.lsplit2 rest ~on:' ' with
           | None -> failwith "time must contain one space between date and ofday"
           | Some (date, ofday) ->
             let date = Date.t_of_sexp (sexp_of_string date) in
             let ofday = Time_float.Ofday.t_of_sexp (sexp_of_string ofday) in
             Time_float.of_date_ofday date ofday ~zone:Time_float.Zone.utc, sexp)
      with
      | Sexplib.Conv.Of_sexp_error (e, _) | e ->
        raise (Sexplib.Conv.Of_sexp_error (e, sexp))
    ;;

    let sexp_of_t_ref = ref (fun (_, sexp) -> sexp)
    let sexp_of_t time = !sexp_of_t_ref time
    let epoch : t = Time_float.epoch, Atom "1970-01-01 00:00:00Z"
  end

  type t =
    { username : string option [@sexp.option]
    ; hostname : string option [@sexp.option]
    ; kernel : string option [@sexp.option]
    ; build_time : Time_with_limited_parsing.t option [@sexp.option]
    ; x_library_inlining : bool
    ; portable_int63 : bool
    ; dynlinkable_code : bool
    ; risk_system : bool [@sexp.default false]
    ; ocaml_version : string
    ; executable_path : string
    ; build_system : string
    ; allowed_projections : string list option [@sexp.option]
    ; with_fdo : (string * Md5.t option) option [@sexp.option]
    ; application_specific_fields : Application_specific_fields.t option [@sexp.option]
    }
  [@@deriving sexp]

  module Structured = struct
    type nonrec t =
      | Not_supported
      | Unset
      | Set of (string * Sexp.t * t)

    let t =
      Exn.handle_uncaught_and_exit (fun () ->
        match generated_build_info () with
        | "" -> Not_supported
        | non_empty_string ->
          (match Build_info_section.chop_start_marker_if_exists non_empty_string with
           | "NO_BUILD_INFO" -> Unset
           | str ->
             let sexp = Sexp.of_string str in
             let t = t_of_sexp sexp in
             Set (str, sexp, t)))
    ;;
  end

  let dummy : t =
    { username = None
    ; hostname = None
    ; kernel = None
    ; build_time = Some Time_with_limited_parsing.epoch
    ; x_library_inlining = false
    ; portable_int63 = true
    ; dynlinkable_code = false
    ; risk_system = false
    ; ocaml_version = ""
    ; executable_path = ""
    ; build_system = ""
    ; allowed_projections = None
    ; with_fdo = None
    ; application_specific_fields = None
    }
  ;;

  let build_system_supports_version_util =
    match Structured.t with
    | Not_supported -> false
    | Unset | Set _ -> true
  ;;

  let build_info_status =
    match Structured.t with
    | Not_supported -> `Not_supported
    | Unset -> `Unset
    | Set _ -> `Set
  ;;

  let build_info, build_info_as_sexp, t =
    match Structured.t with
    | Not_supported | Unset ->
      let t = dummy in
      let sexp = sexp_of_t t in
      let str = Sexp.to_string_mach sexp in
      str, sexp, t
    | Set tuple -> tuple
  ;;

  let { username
      ; hostname
      ; kernel
      ; build_time = build_time_and_sexp
      ; x_library_inlining
      ; portable_int63 = _
      ; dynlinkable_code
      ; risk_system = _
      ; ocaml_version
      ; executable_path
      ; build_system
      ; allowed_projections
      ; with_fdo
      ; application_specific_fields
      }
    =
    t
  ;;

  let build_time =
    match build_time_and_sexp with
    | None -> None
    | Some (time, _sexp) -> Some time
  ;;

  let reprint_build_info sexp_of_time =
    Ref.set_temporarily
      Time_with_limited_parsing.sexp_of_t_ref
      (fun (time, _) -> sexp_of_time time)
      ~f:(fun () -> Sexp.to_string (sexp_of_t t))
  ;;
end

include Build_info

let compiled_for_speed = x_library_inlining && not dynlinkable_code

module For_tests = struct
  let build_info_status = Build_info.build_info_status
  let parse_generated_hg_version = parse_generated_hg_version
end

let arg_spec =
  [ ( "-version"
    , Arg.Unit
        (fun () ->
          List.iter version_list ~f:print_endline;
          exit 0)
    , " Print the hg revision of this build and exit" )
  ; ( "-build_info"
    , Arg.Unit
        (fun () ->
          print_endline build_info;
          exit 0)
    , " Print build info as sexp and exit" )
  ]
;;

module Private__For_version_util_async = struct
  let version_util_start_marker = Version_util_section.Expert.start_marker
  let parse_generated_hg_version = parse_generated_hg_version
  let raw_text = Expert.raw_text
end
