Add a function to find unique solution for constraints

New function Type_check.solve : Env.t -> nexp -> Big_int.num option.
Takes an environment and an n-expression (nexp), and returns either
Some u, where u is a unique solution such that nexp = u, or None which
indicates that either no unique solution could be found. It is
technically possible that a unique solution could exist, but Z3 may
not find it.

Involves two calls to Z3, one of which cannot be memoised, so should
be used carefully, as over-reliance could lead to performance issues.

4 files changed, 69 insertions, 3 deletions

diff --git a/src/constraint.ml b/src/constraint.ml
index ae72d956..1b9447f8 100644
--- a/src/constraint.ml
+++ b/src/constraint.ml
@@ -170,11 +170,13 @@ let rec sexpr_of_constraint = function
   | Boolean true -> Atom "true"
   | Boolean false -> Atom "false"
 
-let smtlib_of_constraints constr : string =
+let smtlib_of_constraints ?get_model:(get_model=false) constr : string =
   "(push)\n"
   ^ var_decs constr ^ "\n"
   ^ pp_sexpr (sfun "define-fun" [Atom "constraint"; List []; Atom "Bool"; sexpr_of_constraint constr])
-  ^ "\n(assert constraint)\n(check-sat)\n(pop)"
+  ^ "\n(assert constraint)\n(check-sat)"
+  ^ (if get_model then "\n(get-model)" else "")
+  ^ "\n(pop)"
 
 type t = nexp constraint_bool
 
@@ -259,7 +261,40 @@ let rec call_z3 constraints : smt_result =
          else (known_problems := DigestMap.add digest Unknown !known_problems; Unknown)
     end
 
-let string_of = smtlib_of_constraints
+let rec solve_z3 constraints var =
+  let problems = [constraints] in
+  let z3_file = smtlib_of_constraints ~get_model:true constraints in
+
+  (* prerr_endline (Printf.sprintf "SMTLIB2 constraints are: \n%s%!" z3_file); *)
+
+  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_" ".sat" in
+  output_string tmp_chan z3_file;
+  close_out tmp_chan;
+  let z3_chan = Unix.open_process_in ("z3 -t:1000 -T:10 " ^ input_file) in
+  let z3_output = String.concat " " (input_all z3_chan) in
+  let _ = Unix.close_process_in z3_chan in
+  Sys.remove input_file;
+  let regexp = {|(define-fun v|} ^ string_of_int var ^ {| () Int[ ]+\([0-9]+\))|} in
+  try
+    let _ = Str.search_forward (Str.regexp regexp) z3_output 0 in
+    let result = Big_int.of_string (Str.matched_group 1 z3_output) in
+    begin match call_z3 (BFun (And, constraints, CFun (NEq, NConstant result, NVar var))) with
+    | Unsat -> Some result
+    | _ -> None
+    end
+  with
+    Not_found -> None
+
+let string_of constr = smtlib_of_constraints constr
 
 (* ===== Abstract API for building constraints ===== *)
 
diff --git a/src/constraint.mli b/src/constraint.mli
index 2111a4e3..df9c8b3a 100644
--- a/src/constraint.mli
+++ b/src/constraint.mli
@@ -60,6 +60,8 @@ val save_digests : unit -> unit
 
 val call_z3 : t -> smt_result
 
+val solve_z3 : t -> int -> Big_int.num option
+
 val string_of : t -> string
 
 val implies : t -> t -> t
diff --git a/src/type_check.ml b/src/type_check.ml
index 760345a6..216e8752 100644
--- a/src/type_check.ml
+++ b/src/type_check.ml
@@ -345,6 +345,7 @@ module Env : sig
   val have_smt_op : id -> t -> bool
   (* Well formedness-checks *)
   val wf_typ : ?exs:KidSet.t -> t -> typ -> unit
+  val wf_nexp : ?exs:KidSet.t -> t -> nexp -> unit
   val wf_constraint : ?exs:KidSet.t -> t -> n_constraint -> unit
 
   (* Some of the code in the environment needs to use the Z3 prover,
@@ -1171,6 +1172,24 @@ let prove_z3 env nc =
   typ_print ("Prove " ^ string_of_list ", " string_of_n_constraint (Env.get_constraints env) ^ " |- " ^ string_of_n_constraint nc);
   prove_z3' env (fun var_of -> Constraint.negate (nc_constraint env var_of nc))
 
+let solve env nexp =
+  typ_print ("Solve " ^ string_of_list ", " string_of_n_constraint (Env.get_constraints env) ^ " |- " ^ string_of_nexp nexp ^ " = ?");
+  let bindings = ref KBindings.empty  in
+  let fresh_var kid =
+    let n = KBindings.cardinal !bindings in
+    bindings := KBindings.add kid n !bindings;
+    n
+  in
+  let var_of kid =
+    try KBindings.find kid !bindings with
+    | Not_found -> fresh_var kid
+  in
+  let env = Env.add_typ_var Parse_ast.Unknown (mk_kid "solve#") BK_nat env in
+  let constr = Constraint.conj (nc_constraints env var_of (Env.get_constraints env))
+                               (nc_constraint env var_of (nc_eq (nvar (mk_kid "solve#")) nexp))
+  in
+  Constraint.solve_z3 constr (var_of (mk_kid "solve#"))
+
 let prove env (NC_aux (nc_aux, _) as nc) =
   let compare_const f (Nexp_aux (n1, _)) (Nexp_aux (n2, _)) =
     match n1, n2 with
@@ -2212,6 +2231,14 @@ let rec check_exp env (E_aux (exp_aux, (l, ())) as exp : unit exp) (Typ_aux (typ
      if prove env nc
      then annot_exp (E_lit (L_aux (L_unit, Parse_ast.Unknown))) unit_typ
      else typ_error l ("Cannot prove " ^ string_of_n_constraint nc)
+  | E_app (f, [E_aux (E_sizeof nexp, _)]), _ when Id.compare f (mk_id "__solve") = 0 ->
+     Env.wf_nexp env nexp;
+     begin match solve env nexp with
+     | None -> typ_error l ("Coud not solve " ^ string_of_nexp nexp)
+     | Some n ->
+        print_endline ("Solved " ^ string_of_nexp nexp ^ " = " ^ Big_int.to_string n);
+        annot_exp (E_lit (L_aux (L_unit, Parse_ast.Unknown))) unit_typ
+     end
   | E_app (f, xs), _ when List.length (Env.get_overloads f env) > 0 ->
      let rec try_overload = function
        | (errs, []) -> typ_raise l (Err_no_overloading (f, errs))
diff --git a/src/type_check.mli b/src/type_check.mli
index 4685b676..2cc852c9 100644
--- a/src/type_check.mli
+++ b/src/type_check.mli
@@ -220,6 +220,8 @@ val infer_exp : Env.t -> unit exp -> tannot exp
 
 val prove : Env.t -> n_constraint -> bool
 
+val solve : Env.t -> nexp -> Big_int.num option
+
 val subtype_check : Env.t -> typ -> typ -> bool
 
 val bind_pat : Env.t -> unit pat -> typ -> tannot pat * Env.t * unit Ast.exp list