(* Js_of_ocaml compiler
 * http://www.ocsigen.org/js_of_ocaml/
 * Copyright (C) 2017 Hugo Heuzard
 *
 * 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, with linking exception;
 * either version 2.1 of the License, or (at your option) any 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.
 *)
open! Stdlib

let times = Debug.find "times"

let debug = Debug.find "link"

let info_prefix = "//#"

let sourceMappingURL = "//# sourceMappingURL="

let sourceMappingURL_base64 = "//# sourceMappingURL=data:application/json;base64,"

module Line_reader : sig
  type t

  val open_ : string -> t

  val next : t -> string

  val peek : t -> string option

  val drop : t -> unit

  val reset : t -> unit

  val lnum : t -> int

  val fname : t -> string
end = struct
  type t =
    { fname : string
    ; lines : string list
    ; mutable current : string list
    ; mutable lnum : int
    }

  let reset t =
    t.current <- t.lines;
    t.lnum <- 0

  let open_ fname =
    let lines = file_lines_bin fname in
    { lines; lnum = 0; fname; current = lines }

  let next t =
    let lnum = t.lnum in
    match t.current with
    | [] -> raise End_of_file
    | s :: rem ->
        t.current <- rem;
        t.lnum <- lnum + 1;
        s

  let peek t =
    match t.current with
    | x :: _ -> Some x
    | [] -> None

  let drop t =
    match t.current with
    | [] -> ()
    | _ :: rem ->
        t.current <- rem;
        t.lnum <- t.lnum + 1

  let lnum t = t.lnum

  let fname t = t.fname
end

module Line_writer : sig
  type t

  val of_channel : out_channel -> t

  val write : t -> string -> unit

  val write_lines : t -> string -> unit

  val lnum : t -> int
end = struct
  type t =
    { oc : out_channel
    ; mutable lnum : int
    }

  let of_channel oc = { oc; lnum = 0 }

  let write t s =
    output_string t.oc s;
    output_string t.oc "\n";
    t.lnum <- t.lnum + 1

  let write_lines t lines =
    let l = String.split_on_char ~sep:'\n' lines in
    let rec w = function
      | [ "" ] | [] -> ()
      | s :: xs ->
          write t s;
          w xs
    in
    w l

  let lnum t = t.lnum
end

type action =
  | Keep
  | Drop
  | Unit
  | Build_info of Build_info.t
  | Source_map of Source_map.t

let prefix_kind line =
  match String.starts_with ~prefix:info_prefix line with
  | false -> `Other
  | true -> (
      match String.starts_with ~prefix:sourceMappingURL line with
      | false -> (
          match Build_info.parse_comment line with
          | Some bi -> `Build_info bi
          | None -> (
              match Unit_info.parse Unit_info.empty line with
              | Some _ -> `Unit
              | None -> `Other))
      | true -> (
          match String.starts_with ~prefix:sourceMappingURL_base64 line with
          | true -> `Json_base64 (String.length sourceMappingURL_base64)
          | false -> `Url (String.length sourceMappingURL)))

let action ~resolve_sourcemap_url ~drop_source_map file line =
  match prefix_kind line, drop_source_map with
  | `Other, (true | false) -> (
      match line with
      | "" -> Drop
      | s when String.equal s Global_constant.header -> Drop
      | _ -> Keep)
  | `Unit, (true | false) -> Unit
  | `Build_info bi, _ -> Build_info bi
  | (`Json_base64 _ | `Url _), true -> Drop
  | `Json_base64 offset, false ->
      let raw = Base64.decode_exn ~off:offset line in
      let sm = Source_map.of_string raw in
      Source_map sm
  | `Url _, false when not resolve_sourcemap_url -> Drop
  | `Url offset, false ->
      let url = String.sub line ~pos:offset ~len:(String.length line - offset) in
      let base = Filename.dirname file in
      let ic = open_in_bin (Filename.concat base url) in
      let l = in_channel_length ic in
      let content = really_input_string ic l in
      close_in ic;
      Source_map (Source_map.of_string content)

module Units : sig
  val read : Line_reader.t -> Unit_info.t -> Unit_info.t

  val scan_file : Line_reader.t -> Build_info.t option * Unit_info.t list
end = struct
  let rec read ic uinfo =
    match Line_reader.peek ic with
    | None -> uinfo
    | Some line -> (
        match Unit_info.parse uinfo line with
        | None -> uinfo
        | Some uinfo ->
            Line_reader.drop ic;
            read ic uinfo)

  let find_unit_info ic =
    let rec find_next ic =
      match Line_reader.peek ic with
      | None -> None
      | Some line -> (
          match prefix_kind line with
          | `Json_base64 _ | `Url _ | `Other | `Build_info _ ->
              Line_reader.drop ic;
              find_next ic
          | `Unit -> Some (read ic Unit_info.empty))
    in
    find_next ic

  let find_build_info ic =
    let rec find_next ic =
      match Line_reader.peek ic with
      | None -> None
      | Some line -> (
          match prefix_kind line with
          | `Other ->
              Line_reader.drop ic;
              find_next ic
          | `Build_info bi -> Some bi
          | `Unit | `Json_base64 _ | `Url _ -> None)
    in
    find_next ic

  let scan_file ic =
    let rec scan_all ic acc =
      match find_unit_info ic with
      | None -> List.rev acc
      | Some x -> scan_all ic (x :: acc)
    in
    Line_reader.reset ic;
    let build_info = find_build_info ic in
    let units = scan_all ic [] in
    build_info, units
end

let link
    ~output
    ~linkall
    ~mklib
    ~files
    ~resolve_sourcemap_url
    ~(source_map : Source_map.Encoding_spec.t option) =
  let t = Timer.make () in
  let oc = Line_writer.of_channel output in
  let warn_effects = ref false in
  let files =
    List.map files ~f:(fun file ->
        let lr = Line_reader.open_ file in
        file, lr, Units.scan_file lr)
  in
  let missing, to_link, all =
    List.fold_right
      files
      ~init:
        ( Global_name.Compunit_set.empty
        , Global_name.Compunit_set.empty
        , Global_name.Compunit_set.empty )
      ~f:(fun (_file, _lr, (build_info, units)) acc ->
        let cmo_file =
          match build_info with
          | Some bi -> (
              match Build_info.kind bi with
              | `Cmo -> true
              | `Cma | `Exe | `Runtime | `Unknown -> false)
          | None -> false
        in
        List.fold_right
          units
          ~init:acc
          ~f:(fun (info : Unit_info.t) (requires, to_link, all) ->
            let all = Global_name.Compunit_set.union all info.provides in
            if
              (not (Config.Flag.auto_link ()))
              || mklib
              || cmo_file
              || linkall
              || info.force_link
              || not
                   (Global_name.Compunit_set.is_empty
                      (Global_name.Compunit_set.inter requires info.provides))
            then
              ( Global_name.Compunit_set.diff
                  (Global_name.Compunit_set.union info.requires requires)
                  info.provides
              , Global_name.Compunit_set.union to_link info.provides
              , all )
            else requires, to_link, all))
  in
  let _skip = Global_name.Compunit_set.diff all to_link in
  if (not (Global_name.Compunit_set.is_empty missing)) && not mklib
  then
    failwith
      (Printf.sprintf
         "Could not find compilation unit for %s"
         (String.concat
            ~sep:", "
            (List.map
               ~f:(fun (Global_name.Compunit name) -> name)
               (Global_name.Compunit_set.elements missing))));
  if times () then Format.eprintf "  scan: %a@." Timer.print t;
  let sm = ref [] in
  let build_info = ref None in
  let t = Timer.make () in
  let sym = ref Ocaml_compiler.Symtable.GlobalMap.empty in
  Array.iter Runtimedef.builtin_exceptions ~f:(fun name ->
      ignore
        (Ocaml_compiler.Symtable.GlobalMap.enter
           sym
           (Global_name.Glob_predef (Predef name))
          : int));
  List.iter files ~f:(fun (_, _, (_, units)) ->
      List.iter units ~f:(fun (u : Unit_info.t) ->
          Global_name.Compunit_set.iter
            (fun (Global_name.Compunit name) ->
              ignore
                (Ocaml_compiler.Symtable.GlobalMap.enter
                   sym
                   (Global_name.Glob_compunit (Compunit name))
                  : int))
            u.Unit_info.provides));

  let build_info_emitted = ref false in
  List.iter files ~f:(fun (file, ic, (build_info_for_file, units)) ->
      Line_reader.reset ic;
      let is_runtime =
        match build_info_for_file with
        | Some bi -> (
            match Build_info.kind bi with
            | `Runtime -> Some bi
            | `Cma | `Exe | `Cmo | `Unknown -> None)
        | None -> None
      in
      let sm_for_file = ref None in
      let skip ic = Line_reader.drop ic in
      let line_offset = Line_writer.lnum oc in
      let reloc = ref [] in
      let copy ic oc =
        let line = Line_reader.next ic in
        Line_writer.write oc line
      in
      let rec read () =
        match Line_reader.peek ic with
        | None -> ()
        | Some line ->
            let drop_source_map =
              match source_map with
              | None | Some { keep_empty = true; _ } -> true
              | Some { keep_empty = false; _ } -> false
            in
            (match action ~resolve_sourcemap_url ~drop_source_map file line with
            | Keep -> copy ic oc
            | Build_info bi ->
                skip ic;
                if not !build_info_emitted
                then (
                  let bi = Build_info.with_kind bi (if mklib then `Cma else `Unknown) in
                  Line_writer.write oc Global_constant.header;
                  Line_writer.write_lines oc (Build_info.to_comment bi);
                  build_info_emitted := true)
            | Drop -> skip ic
            | Unit ->
                let u = Units.read ic Unit_info.empty in
                if
                  u.Unit_info.force_link
                  || not (Global_name.Compunit_set.disjoint u.Unit_info.provides to_link)
                then (
                  if u.effects_without_cps && not !warn_effects
                  then (
                    warn_effects := true;
                    Warning.warn
                      `Effect_handlers_without_effect_backend
                      "your program contains effect handlers; you should probably run \
                       js_of_ocaml with option '--effects=cps'@.");
                  (if mklib
                   then
                     let u = if linkall then { u with force_link = true } else u in
                     Line_writer.write_lines oc (Unit_info.to_string u));
                  Line_writer.write
                    oc
                    (Printf.sprintf
                       "//# %d %S"
                       (Line_reader.lnum ic)
                       (Line_reader.fname ic));
                  let read_loffset = Line_reader.lnum ic in
                  let write_loffset = Line_writer.lnum oc in
                  let bsize = ref 0 in
                  let lsize = ref 0 in
                  while
                    match Line_reader.peek ic with
                    | None -> false
                    | Some line -> (
                        match prefix_kind line with
                        | `Other -> true
                        | `Json_base64 _ | `Url _ | `Build_info _ | `Unit -> false)
                  do
                    bsize := !bsize + String.length line + 1;
                    incr lsize;
                    copy ic oc
                  done;
                  assert (read_loffset + !lsize = Line_reader.lnum ic);
                  assert (write_loffset + !lsize = Line_writer.lnum oc);
                  reloc := `Copy (read_loffset, write_loffset, !lsize) :: !reloc;
                  if debug ()
                  then
                    Format.eprintf
                      "Copy %d bytes for %s@."
                      !bsize
                      (match is_runtime with
                      | None ->
                          String.concat
                            ~sep:", "
                            (List.map
                               ~f:(fun (Global_name.Compunit name) -> name)
                               (Global_name.Compunit_set.elements u.provides))
                      | Some _ -> "the js runtime"))
                else (
                  if debug ()
                  then
                    Format.eprintf
                      "Skip %s@."
                      (String.concat
                         ~sep:","
                         (List.map
                            ~f:(fun (Global_name.Compunit name) -> name)
                            (Global_name.Compunit_set.elements u.provides)));
                  let lnum = ref 0 in
                  let read_loffset = Line_reader.lnum ic in
                  while
                    match Line_reader.peek ic with
                    | None -> false
                    | Some line -> (
                        match prefix_kind line with
                        | `Other -> true
                        | `Json_base64 _ | `Url _ | `Build_info _ | `Unit -> false)
                  do
                    skip ic;
                    incr lnum
                  done;
                  assert (read_loffset + !lnum = Line_reader.lnum ic);
                  reloc := `Drop (read_loffset, !lnum) :: !reloc)
            | Source_map x ->
                skip ic;
                sm_for_file := Some x);
            read ()
      in
      read ();
      (match is_runtime with
      | None -> ()
      | Some bi ->
          Build_info.configure bi;
          let primitives =
            List.fold_left units ~init:StringSet.empty ~f:(fun acc (u : Unit_info.t) ->
                List.iter u.aliases ~f:(fun (a, b) -> Primitive.alias a b);
                StringSet.union acc (StringSet.of_list u.primitives))
          in
          let num_globals =
            Ocaml_compiler.Symtable.GlobalMap.fold (fun _ n m -> max n m) !sym 0 + 1
          in
          let code =
            Parse_bytecode.link_info ~symbols:!sym ~primitives ~crcs:[] ~num_globals
          in
          let b = Buffer.create 100 in
          let fmt = Pretty_print.to_buffer b in
          Driver.configure fmt;
          Driver.f' ~standalone:false ~link:`No ~wrap_with_fun:`Iife fmt code;
          let content = Buffer.contents b in
          Line_writer.write_lines oc content;
          Line_writer.write oc "");
      (match !sm_for_file with
      | None -> ()
      | Some x -> sm := (x, List.rev !reloc, line_offset) :: !sm);
      match !build_info, build_info_for_file with
      | None, None -> ()
      | Some _, None -> ()
      | None, Some build_info_for_file -> build_info := Some (file, build_info_for_file)
      | Some (first_file, bi), Some build_info_for_file ->
          build_info :=
            Some
              ( first_file
              , Build_info.merge `JavaScript first_file bi file build_info_for_file ));
  if times () then Format.eprintf "  emit: %a@." Timer.print t;
  let t = Timer.make () in
  match source_map with
  | None -> ()
  | Some { output_file = file; source_map = init_sm; keep_empty = _ } ->
      let sections =
        List.rev_map !sm ~f:(fun (sm, reloc, _offset) ->
            let sm =
              match (sm : Source_map.t) with
              | Standard sm ->
                  [ ( Source_map.Mappings.first_line sm.mappings
                    , Source_map.Mappings.number_of_lines sm.mappings
                    , 0
                    , 0
                    , sm )
                  ]
              | Index sm ->
                  List.map
                    sm.Source_map.Index.sections
                    ~f:(fun { offset = { gen_line; gen_column }; map } ->
                      ( gen_line + Source_map.Mappings.first_line map.mappings
                      , gen_line + Source_map.Mappings.number_of_lines map.mappings
                      , gen_line
                      , gen_column
                      , map ))
            in
            (* select sourcemaps that cover copied section *)
            let maps =
              List.concat_map reloc ~f:(function
                | `Drop _ -> []
                | `Copy (src, dst, len) ->
                    List.filter_map sm ~f:(fun (first, last, gen_line, gen_column, sm) ->
                        if first > src + len || last < src
                        then None
                        else (
                          (* We don't want to deal with overlapping but not included
                               sourcemap, but we could in theory filter out part of it. *)
                          assert (src <= first && last <= src + len);
                          Some (first, last, gen_line + dst - src, gen_column, sm))))
            in
            (* Make sure dropped sections are not overlapping selected sourcemap. *)
            List.iter reloc ~f:(function
              | `Copy _ -> ()
              | `Drop (src, len) ->
                  List.iter maps ~f:(fun (first, last, _, _, _) ->
                      if first > src + len || last < src then () else assert false));
            maps)
      in
      let sections = List.concat sections in
      let sm =
        { Source_map.Index.version = init_sm.Source_map.Standard.version
        ; file = init_sm.file
        ; sections =
            (* preserve some info from [init_sm] *)
            List.map sections ~f:(fun (_, _, gen_line, gen_column, sm) ->
                { Source_map.Index.offset = { gen_line; gen_column }
                ; map = { sm with sourceroot = init_sm.sourceroot }
                })
        }
      in
      let sm = Source_map.Index sm in
      (match file with
      | None ->
          let data = Source_map.to_string sm in
          let s = sourceMappingURL_base64 ^ Base64.encode_exn data in
          Line_writer.write oc s
      | Some file ->
          Source_map.to_file sm file;
          let s = sourceMappingURL ^ Filename.basename file in
          Line_writer.write oc s);
      if times () then Format.eprintf "  sourcemap: %a@." Timer.print t

let link ~output ~linkall ~mklib ~files ~resolve_sourcemap_url ~source_map =
  try link ~output ~linkall ~mklib ~files ~resolve_sourcemap_url ~source_map
  with Build_info.Incompatible_build_info { key; first = f1, v1; second = f2, v2 } ->
    let string_of_v = function
      | None -> "<empty>"
      | Some v -> v
    in
    failwith
      (Printf.sprintf
         "Incompatible build info detected while linking.\n - %s: %s=%s\n - %s: %s=%s"
         f1
         key
         (string_of_v v1)
         f2
         key
         (string_of_v v2))
