(*
 * Copyright (c) Facebook, Inc. and its affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 *)

(** Global namespace opened in each source file by the build system *)

include (
  Base :
    sig
      include
        (module type of Base
          (* prematurely deprecated, remove and use Stdlib instead *)
          with module Filename := Base.Filename
           and module Format := Base.Format
           and module Marshal := Base.Marshal
           and module Scanf := Base.Scanf
           and type ('ok, 'err) result := ('ok, 'err) Base.result
         [@warning "-3"])
    end )

(* undeprecate *)
external ( == ) : 'a -> 'a -> bool = "%eq"
external ( != ) : 'a -> 'a -> bool = "%noteq"

exception Not_found = Caml.Not_found

include Stdio
module Command = Core.Command
module Hash_queue = Core_kernel.Hash_queue

(** Tuple operations *)

let fst3 (x, _, _) = x
let snd3 (_, y, _) = y
let trd3 (_, _, z) = z

(** Function combinators *)

let ( >> ) f g x = g (f x)
let ( << ) f g x = f (g x)
let ( $ ) f g x = f x ; g x
let ( $> ) x f = f x ; x
let ( <$ ) f x = f x ; x

(** Pretty-printing *)

type 'a pp = Formatter.t -> 'a -> unit
type ('a, 'b) fmt = ('a, 'b) Trace.fmt

(** Failures *)

let fail = Trace.fail

exception Unimplemented of string

let todo fmt = Trace.raisef (fun msg -> Unimplemented msg) fmt

let warn fmt =
  let fs = Format.std_formatter in
  Format.pp_open_box fs 2 ;
  Format.pp_print_string fs "Warning: " ;
  Format.kfprintf
    (fun fs () ->
      Format.pp_close_box fs () ;
      Format.pp_force_newline fs () )
    fs fmt

(** Assertions *)

let assertf cnd fmt =
  if not cnd then fail fmt
  else Format.ikfprintf (fun _ () -> ()) Format.str_formatter fmt

let checkf cnd fmt =
  if not cnd then fail fmt
  else Format.ikfprintf (fun _ () -> true) Format.str_formatter fmt

let check f x =
  assert (f x ; true) ;
  x

let violates f x =
  assert (f x ; true) ;
  assert false

type 'a or_error = ('a, exn * Caml.Printexc.raw_backtrace) result

let or_error f x () =
  try Ok (f x) with exn -> Error (exn, Caml.Printexc.get_raw_backtrace ())

(** Extensions *)

module Invariant = struct
  include Base.Invariant

  let invariant here t sexp_of_t f =
    assert (
      ( try f ()
        with exn ->
          let bt = Caml.Printexc.get_raw_backtrace () in
          let exn =
            Error.to_exn
              (Error.create_s
                 (Base.Sexp.message "invariant failed"
                    [ ("", sexp_of_exn exn)
                    ; ("", Source_code_position.sexp_of_t here)
                    ; ("", sexp_of_t t) ]))
          in
          Caml.Printexc.raise_with_backtrace exn bt ) ;
      true )
end

let map_preserving_phys_equal map t ~f =
  let change = ref false in
  let t' =
    map t ~f:(fun x ->
        let x' = f x in
        if not (x' == x) then change := true ;
        x' )
  in
  if !change then t' else t

let filter_map_preserving_phys_equal filter_map t ~f =
  let change = ref false in
  let t' =
    filter_map t ~f:(fun x ->
        let x'_opt = f x in
        ( match x'_opt with
        | Some x' when x' == x -> ()
        | _ -> change := true ) ;
        x'_opt )
  in
  if !change then t' else t

module type Applicative_syntax = sig
  type 'a t

  val ( let+ ) : 'a t -> ('a -> 'b) -> 'b t
  val ( and+ ) : 'a t -> 'b t -> ('a * 'b) t
end

module type Monad_syntax = sig
  include Applicative_syntax

  val ( let* ) : 'a t -> ('a -> 'b t) -> 'b t
  val ( and* ) : 'a t -> 'b t -> ('a * 'b) t
end

module Option = struct
  include Base.Option

  let pp fmt pp_elt fs = function
    | Some x -> Format.fprintf fs fmt pp_elt x
    | None -> ()

  let cons xo xs = match xo with Some x -> x :: xs | None -> xs

  module Monad_syntax = struct
    type nonrec 'a t = 'a t

    let ( let+ ) x f = map ~f x
    let ( and+ ) x y = both x y
    let ( let* ) x f = bind ~f x
    let ( and* ) x y = both x y
  end
end

include Option.Monad_infix
include Option.Monad_syntax

module List = struct
  include Base.List

  let rec pp ?pre ?suf sep pp_elt fs = function
    | [] -> ()
    | x :: xs ->
        Option.iter pre ~f:(Format.fprintf fs) ;
        pp_elt fs x ;
        ( match xs with
        | [] -> ()
        | xs -> Format.fprintf fs "%( %)%a" sep (pp sep pp_elt) xs ) ;
        Option.iter suf ~f:(Format.fprintf fs)

  let pop_exn = function x :: xs -> (x, xs) | [] -> raise Not_found

  let find_map_remove xs ~f =
    let rec find_map_remove_ ys = function
      | [] -> None
      | x :: xs -> (
        match f x with
        | Some x' -> Some (x', rev_append ys xs)
        | None -> find_map_remove_ (x :: ys) xs )
    in
    find_map_remove_ [] xs

  let fold_option xs ~init ~f =
    With_return.with_return
    @@ fun {return} ->
    Some
      (fold xs ~init ~f:(fun acc elt ->
           match f acc elt with Some res -> res | None -> return None ))

  let filter_map_preserving_phys_equal t ~f =
    filter_map_preserving_phys_equal filter_map t ~f

  let map_preserving_phys_equal t ~f = map_preserving_phys_equal map t ~f

  let rev_map_unzip xs ~f =
    fold xs ~init:([], []) ~f:(fun (ys, zs) x ->
        let y, z = f x in
        (y :: ys, z :: zs) )

  let remove_exn ?(equal = phys_equal) xs x =
    let rec remove_ ys = function
      | [] -> raise Not_found
      | z :: xs ->
          if equal x z then rev_append ys xs else remove_ (z :: ys) xs
    in
    remove_ [] xs

  let remove ?equal xs x =
    try Some (remove_exn ?equal xs x) with Not_found -> None

  let rec rev_init n ~f =
    if n = 0 then []
    else
      let n = n - 1 in
      let xs = rev_init n ~f in
      f n :: xs

  let symmetric_diff ~compare xs ys =
    let rec symmetric_diff_ xxs yys =
      match (xxs, yys) with
      | x :: xs, y :: ys ->
          let ord = compare x y in
          if ord = 0 then symmetric_diff_ xs ys
          else if ord < 0 then Either.First x :: symmetric_diff_ xs yys
          else Either.Second y :: symmetric_diff_ xxs ys
      | xs, [] -> map ~f:Either.first xs
      | [], ys -> map ~f:Either.second ys
    in
    symmetric_diff_ (sort ~compare xs) (sort ~compare ys)

  let pp_diff ~compare sep pp_elt fs (xs, ys) =
    let pp_diff_elt fs elt =
      match (elt : _ Either.t) with
      | First x -> Format.fprintf fs "-- %a" pp_elt x
      | Second y -> Format.fprintf fs "++ %a" pp_elt y
    in
    pp sep pp_diff_elt fs (symmetric_diff ~compare xs ys)
end

module Map = struct
  include Base.Map

  let pp pp_k pp_v fs m =
    Format.fprintf fs "@[<1>[%a]@]"
      (List.pp ",@ " (fun fs (k, v) ->
           Format.fprintf fs "@[%a @<2>↦ %a@]" pp_k k pp_v v ))
      (to_alist m)

  let pp_diff ~data_equal pp_key pp_val pp_diff_val fs (x, y) =
    let pp_diff_elt fs = function
      | k, `Left v ->
          Format.fprintf fs "-- [@[%a@ @<2>↦ %a@]]" pp_key k pp_val v
      | k, `Right v ->
          Format.fprintf fs "++ [@[%a@ @<2>↦ %a@]]" pp_key k pp_val v
      | k, `Unequal vv ->
          Format.fprintf fs "[@[%a@ @<2>↦ %a@]]" pp_key k pp_diff_val vv
    in
    let sd = Sequence.to_list (symmetric_diff ~data_equal x y) in
    if not (List.is_empty sd) then
      Format.fprintf fs "[@[<hv>%a@]];@ " (List.pp ";@ " pp_diff_elt) sd

  let equal_m__t (module Elt : Compare_m) equal_v = equal equal_v

  let find_and_remove m k =
    let found = ref None in
    let m =
      change m k ~f:(fun v ->
          found := v ;
          None )
    in
    let+ v = !found in
    (v, m)

  let find_or_add (type data) map key ~(default : data) ~if_found ~if_added
      =
    let exception Found of data in
    match
      update map key ~f:(function
        | Some old_data -> Exn.raise_without_backtrace (Found old_data)
        | None -> default )
    with
    | exception Found old_data -> if_found old_data
    | map -> if_added map

  let map_preserving_phys_equal t ~f = map_preserving_phys_equal map t ~f
end

module Result = struct
  include Base.Result

  let pp fmt pp_elt fs = function
    | Ok x -> Format.fprintf fs fmt pp_elt x
    | Error _ -> ()
end

module Vector = struct
  include Vector

  let pp sep pp_elt fs v = List.pp sep pp_elt fs (to_list v)
end

include Vector.Infix

module Set = struct
  include Base.Set

  type ('elt, 'cmp) tree = ('elt, 'cmp) Using_comparator.Tree.t

  let equal_m__t (module Elt : Compare_m) = equal
  let pp pp_elt fs x = List.pp ",@ " pp_elt fs (to_list x)

  let pp_diff pp_elt fs (xs, ys) =
    let lose = diff xs ys and gain = diff ys xs in
    if not (is_empty lose) then Format.fprintf fs "-- %a" (pp pp_elt) lose ;
    if not (is_empty gain) then Format.fprintf fs "++ %a" (pp pp_elt) gain

  let disjoint x y = is_empty (inter x y)
  let add_option yo x = Option.fold ~f:add ~init:x yo
  let add_list ys x = List.fold ~f:add ~init:x ys
  let diff_inter x y = (diff x y, inter x y)
  let diff_inter_diff x y = (diff x y, inter x y, diff y x)
  let of_vector cmp x = of_array cmp (Vector.to_array x)
  let to_tree = Using_comparator.to_tree

  let union x y =
    let xy = union x y in
    let xy_tree = to_tree xy in
    if xy_tree == to_tree x then x
    else if xy_tree == to_tree y then y
    else xy
end

module Qset = struct
  include Qset

  let pp sep pp_elt fs s = List.pp sep pp_elt fs (to_list s)
end

module Array = struct
  include Base.Array

  let pp sep pp_elt fs a = List.pp sep pp_elt fs (to_list a)
end

module Q = struct
  let pp = Q.pp_print
  let hash = Hashtbl.hash
  let hash_fold_t s q = Int.hash_fold_t s (hash q)
  let sexp_of_t q = Sexp.Atom (Q.to_string q)

  let t_of_sexp = function
    | Sexp.Atom s -> Q.of_string s
    | _ -> assert false

  let of_z = Q.of_bigint

  include Q
end

module Z = struct
  let pp = Z.pp_print
  let hash = [%hash: Z.t]
  let hash_fold_t s z = Int.hash_fold_t s (hash z)
  let sexp_of_t z = Sexp.Atom (Z.to_string z)

  let t_of_sexp = function
    | Sexp.Atom s -> Z.of_string s
    | _ -> assert false

  (* the signed 1-bit integers are -1 and 0 *)
  let true_ = Z.minus_one
  let false_ = Z.zero
  let of_bool = function true -> true_ | false -> false_
  let is_true = Z.equal true_
  let is_false = Z.equal false_

  include Z
end