(**************************************************************************)
(*     Sail                                                               *)
(*                                                                        *)
(*  Copyright (c) 2013-2017                                               *)
(*    Kathyrn Gray                                                        *)
(*    Shaked Flur                                                         *)
(*    Stephen Kell                                                        *)
(*    Gabriel Kerneis                                                     *)
(*    Robert Norton-Wright                                                *)
(*    Christopher Pulte                                                   *)
(*    Peter Sewell                                                        *)
(*    Alasdair Armstrong                                                  *)
(*    Brian Campbell                                                      *)
(*    Thomas Bauereiss                                                    *)
(*    Anthony Fox                                                         *)
(*    Jon French                                                          *)
(*    Dominic Mulligan                                                    *)
(*    Stephen Kell                                                        *)
(*    Mark Wassell                                                        *)
(*                                                                        *)
(*  All rights reserved.                                                  *)
(*                                                                        *)
(*  This software was developed by the University of Cambridge Computer   *)
(*  Laboratory as part of the Rigorous Engineering of Mainstream Systems  *)
(*  (REMS) project, funded by EPSRC grant EP/K008528/1.                   *)
(*                                                                        *)
(*  Redistribution and use in source and binary forms, with or without    *)
(*  modification, are permitted provided that the following conditions    *)
(*  are met:                                                              *)
(*  1. Redistributions of source code must retain the above copyright     *)
(*     notice, this list of conditions and the following disclaimer.      *)
(*  2. Redistributions in binary form must reproduce the above copyright  *)
(*     notice, this list of conditions and the following disclaimer in    *)
(*     the documentation and/or other materials provided with the         *)
(*     distribution.                                                      *)
(*                                                                        *)
(*  THIS SOFTWARE IS PROVIDED BY THE AUTHOR AND CONTRIBUTORS ``AS IS''    *)
(*  AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED     *)
(*  TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A       *)
(*  PARTICULAR PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE AUTHOR OR   *)
(*  CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,          *)
(*  SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT      *)
(*  LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF      *)
(*  USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND   *)
(*  ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,    *)
(*  OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT    *)
(*  OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF    *)
(*  SUCH DAMAGE.                                                          *)
(**************************************************************************)

module Big_int = Nat_big_num
open Ast
open Ast_util
open Util

let opt_smt_verbose = ref false

type solver = {
    command : string;
    header : string;
    footer : string;
    negative_literals : bool;
    uninterpret_power : bool
  }

let cvc4_solver = {
    command = "cvc4 -L smtlib2 --tlimit=2000";
    header = "(set-logic QF_UFNIA)\n";
    footer = "";
    negative_literals = false;
    uninterpret_power = true
  }

let mathsat_solver = {
    command = "mathsat";
    header = "(set-logic QF_UFLIA)\n";
    footer = "";
    negative_literals = false;
    uninterpret_power = true
  }

let z3_solver = {
    command = "z3 -t:1000 -T:10";
    (* Using push and pop is much faster, I believe because
       incremental mode uses a different solver. *)
    header = "(push)\n";
    footer = "(pop)\n";
    negative_literals = true;
    uninterpret_power = false;
  }

let yices_solver = {
    command = "yices-smt2 --timeout=2";
    header = "(set-logic QF_UFLIA)\n";
    footer = "";
    negative_literals = false;
    uninterpret_power = true
  }

let vampire_solver = {
    (* vampire sometimes likes to ignore its time limit *)
    command = "timeout -s SIGKILL 3s vampire --time_limit 2s --input_syntax smtlib2 --mode smtcomp";
    header = "";
    footer = "";
    negative_literals = false;
    uninterpret_power = true
  }

let alt_ergo_solver ={
    command = "alt-ergo";
    header = "";
    footer = "";
    negative_literals = false;
    uninterpret_power = true
  }

let opt_solver = ref z3_solver

let set_solver = function
  | "z3" -> opt_solver := z3_solver
  | "alt-ergo" -> opt_solver := alt_ergo_solver
  | "cvc4" -> opt_solver := cvc4_solver
  | "mathsat" -> opt_solver := mathsat_solver
  | "vampire" -> opt_solver := vampire_solver
  | "yices" -> opt_solver := yices_solver
  | unknown -> prerr_endline ("Unrecognised SMT solver " ^ unknown)

(* SMTLIB v2.0 format is based on S-expressions so we have a
   lightweight representation of those here. *)
type sexpr = List of (sexpr list) | Atom of string

let sfun (fn : string) (xs : sexpr list) : sexpr = List (Atom fn :: xs)

let rec pp_sexpr : sexpr -> string = function
  | List xs -> "(" ^ string_of_list " " pp_sexpr xs ^ ")"
  | Atom x -> x

(** Each non-Type/Order kind in Sail mapes to a type in the SMT solver *)
let smt_type l = function
  | K_int -> Atom "Int"
  | K_bool -> Atom "Bool"
  | _ -> raise (Reporting.err_unreachable l __POS__ "Tried to pass Type or Order kinded variable to SMT solver")

let to_smt l vars constr =
  (* Numbering all SMT variables v0, ... vn, rather than generating
     names based on their Sail names (e.g. using zencode) ensures that
     alpha-equivalent constraints generate the same SMT problem, which
     is important for the SMT memoisation to work properly. *)
  let var_map = ref KBindings.empty in
  let vnum = ref (-1) in
  let smt_var v =
    match KBindings.find_opt v !var_map with
    | Some n -> Atom ("v" ^ string_of_int n)
    | None ->
       let n = !vnum + 1 in
       var_map := KBindings.add v n !var_map;
       vnum := n;
       Atom ("v" ^ string_of_int n)
  in

  (* var_decs outputs the list of variables to be used by the SMT
     solver in SMTLIB v2.0 format. It takes a kind_aux KBindings, as
     returned by Type_check.get_typ_vars *)
  let var_decs l (vars : kind_aux KBindings.t) : string =
    vars
    |> KBindings.bindings
    |> List.map (fun (v, k) -> sfun "declare-const" [smt_var v; smt_type l k])
    |> string_of_list "\n" pp_sexpr
  in
  let rec smt_nexp (Nexp_aux (aux, l) : nexp) : sexpr =
    match aux with
    | Nexp_id id -> Atom (Util.zencode_string (string_of_id id))
    | Nexp_var v -> smt_var v
    | Nexp_constant c
         when Big_int.less_equal c (Big_int.of_int (-1)) && not !opt_solver.negative_literals ->
       sfun "-" [Atom "0"; Atom (Big_int.to_string (Big_int.abs c))]
    | Nexp_constant c -> Atom (Big_int.to_string c)
    | Nexp_app (id, nexps) -> sfun (string_of_id id) (List.map smt_nexp nexps)
    | Nexp_times (nexp1, nexp2) -> sfun "*" [smt_nexp nexp1; smt_nexp nexp2]
    | Nexp_sum (nexp1, nexp2) -> sfun "+" [smt_nexp nexp1; smt_nexp nexp2]
    | Nexp_minus (nexp1, nexp2) -> sfun "-" [smt_nexp nexp1; smt_nexp nexp2]
    | Nexp_exp (Nexp_aux (Nexp_constant c, _)) when Big_int.greater c Big_int.zero ->
       Atom (Big_int.to_string (Big_int.pow_int_positive 2 (Big_int.to_int c)))
    | Nexp_exp nexp when !opt_solver.uninterpret_power -> sfun "sailexp" [smt_nexp nexp]
    | Nexp_exp nexp -> sfun "^" [Atom "2"; smt_nexp nexp]
    | Nexp_neg nexp -> sfun "-" [smt_nexp nexp]
  in
  let rec smt_constraint (NC_aux (aux, l) : n_constraint) : sexpr =
    match aux with
    | NC_equal (nexp1, nexp2) -> sfun "=" [smt_nexp nexp1; smt_nexp nexp2]
    | NC_bounded_le (nexp1, nexp2) -> sfun "<=" [smt_nexp nexp1; smt_nexp nexp2]
    | NC_bounded_ge (nexp1, nexp2) -> sfun ">=" [smt_nexp nexp1; smt_nexp nexp2]
    | NC_not_equal (nexp1, nexp2) -> sfun "not" [sfun "=" [smt_nexp nexp1; smt_nexp nexp2]]
    | NC_set (v, ints) ->
       sfun "or" (List.map (fun i -> sfun "=" [smt_var v; Atom (Big_int.to_string i)]) ints)
    | NC_or (nc1, nc2) -> sfun "or" [smt_constraint nc1; smt_constraint nc2]
    | NC_and (nc1, nc2) -> sfun "and" [smt_constraint nc1; smt_constraint nc2]
    | NC_app (id, args) ->
       sfun (string_of_id id) (List.map smt_typ_arg args)
    | NC_true -> Atom "true"
    | NC_false -> Atom "false"
    | NC_var v -> smt_var v
  and smt_typ_arg (A_aux (aux, l) : typ_arg) : sexpr =
    match aux with
    | A_nexp nexp -> smt_nexp nexp
    | A_bool nc -> smt_constraint nc
    | _ ->
       raise (Reporting.err_unreachable l __POS__ "Tried to pass Type or Order kind to SMT function")
  in
  var_decs l vars, smt_constraint constr, smt_var

let smtlib_of_constraints ?get_model:(get_model=false) l vars constr : string * (kid -> sexpr) =
  let variables, problem, var_map = to_smt l vars constr in
  !opt_solver.header
  ^ variables ^ "\n"
  ^ (if !opt_solver.uninterpret_power then "(declare-fun sailexp (Int) Int)\n" else "")
  ^ pp_sexpr (sfun "define-fun" [Atom "constraint"; List []; Atom "Bool"; problem])
  ^ "\n(assert constraint)\n(check-sat)"
  ^ (if get_model then "\n(get-model)\n" else "\n")
  ^ !opt_solver.footer,
  var_map

type smt_result = Unknown | Sat | Unsat

module DigestMap = Map.Make(Digest)

let known_problems = ref (DigestMap.empty)

let load_digests_err () =
  let in_chan = open_in_bin "z3_problems" in
  let rec load () =
    let digest = Digest.input in_chan in
    let result = input_byte in_chan in
    begin
      match result with
      | 0 -> known_problems := DigestMap.add digest Unknown !known_problems
      | 1 -> known_problems := DigestMap.add digest Sat !known_problems
      | 2 -> known_problems := DigestMap.add digest Unsat !known_problems
      | _ -> assert false
    end;
    load ()
  in
  try load () with
  | End_of_file -> close_in in_chan

let load_digests () =
  try load_digests_err () with
  | Sys_error _ -> ()

let save_digests () =
  let out_chan = open_out_bin "z3_problems" in
  let output digest result =
    Digest.output out_chan digest;
    match result with
    | Unknown -> output_byte out_chan 0
    | Sat -> output_byte out_chan 1
    | Unsat -> output_byte out_chan 2
  in
  DigestMap.iter output !known_problems;
  close_out out_chan

let call_smt' l vars constraints : smt_result =
  let problems = [constraints] in
  let smt_file, _ = smtlib_of_constraints l vars constraints in

  if !opt_smt_verbose then
    prerr_endline (Printf.sprintf "SMTLIB2 constraints are: \n%s%!" smt_file)
  else ();

  let rec input_lines chan = function
    | 0 -> []
    | n ->
       begin
         let l = input_line chan in
         let ls = input_lines chan (n - 1) in
         l :: ls
       end
  in

  let digest = Digest.string smt_file in
  try
    let result = DigestMap.find digest !known_problems in
    result
  with
  | Not_found ->
    begin
      let (input_file, tmp_chan) =
        try Filename.open_temp_file "constraint_" ".smt2" with
        | Sys_error msg -> raise (Reporting.err_general l ("Could not open temp file when calling SMT: " ^ msg))
      in
      output_string tmp_chan smt_file;
      close_out tmp_chan;
      let smt_output =
        try
          let smt_out, smt_in, smt_err = Unix.open_process_full (!opt_solver.command ^ " " ^ input_file) (Unix.environment ()) in
          let smt_output =
            try List.combine problems (input_lines smt_out (List.length problems)) with
            | End_of_file -> List.combine problems ["unknown"]
          in
          let _ = Unix.close_process_full (smt_out, smt_in, smt_err) in
          smt_output
        with
        | exn -> raise (Reporting.err_general l ("Error when calling smt: " ^ Printexc.to_string exn))
      in
      Sys.remove input_file;
      try
        let (problem, _) = List.find (fun (_, result) -> result = "unsat") smt_output in
        known_problems := DigestMap.add digest Unsat !known_problems;
        Unsat
      with
      | Not_found ->
         let unsolved = List.filter (fun (_, result) -> result = "unknown") smt_output in
         if unsolved == []
         then (known_problems := DigestMap.add digest Sat !known_problems; Sat)
         else (known_problems := DigestMap.add digest Unknown !known_problems; Unknown)
    end

let call_smt l vars constraints =
  let t = Profile.start_smt () in
  let result = call_smt' l vars constraints in
  Profile.finish_smt t;
  result

let solve_smt l vars constraints var =
  let smt_file, smt_var = smtlib_of_constraints ~get_model:true l vars constraints in
  let smt_var = pp_sexpr (smt_var var) in

  if !opt_smt_verbose then
    prerr_endline (Printf.sprintf "SMTLIB2 constraints are (solve for %s): \n%s%!" smt_var smt_file)
  else ();

  let rec input_all chan =
    try
      let l = input_line chan in
      let ls = input_all chan in
      l :: ls
    with
      End_of_file -> []
  in

  let (input_file, tmp_chan) = Filename.open_temp_file "constraint_" ".smt2" in
  output_string tmp_chan smt_file;
  close_out tmp_chan;
  let smt_output =
    try
      let smt_chan = Unix.open_process_in ("z3 -t:1000 -T:10 " ^ input_file) in
      let smt_output = String.concat " " (input_all smt_chan) in
      let _ = Unix.close_process_in smt_chan in
      smt_output
    with
    | exn ->
       raise (Reporting.err_general l ("Got error when calling smt: " ^ Printexc.to_string exn))
  in
  Sys.remove input_file;
  let regexp = {|(define-fun |} ^ smt_var ^ {| () Int[ ]+\([0-9]+\))|} in
  try
    let _ = Str.search_forward (Str.regexp regexp) smt_output 0 in
    let result = Big_int.of_string (Str.matched_group 1 smt_output) in
    Some result
  with
  | Not_found -> None

let solve_all_smt l vars constraints var =
  let rec aux results =
    let constraints = List.fold_left (fun ncs r -> (nc_and ncs (nc_neq (nconstant r) (nvar var)))) constraints results in
    match solve_smt l vars constraints var with
    | Some result -> aux (result :: results)
    | None ->
       match call_smt l vars constraints with
       | Unsat -> Some results
       | _ -> None
  in
  aux []

let solve_unique_smt l vars constraints var =
  match solve_smt l vars constraints var with
  | Some result ->
     begin match call_smt l vars (nc_and constraints (nc_neq (nconstant result) (nvar var))) with
     | Unsat -> Some result
     | _ -> None
     end
  | None -> None