Add option to linearize constraints containing exponentials

When adding a constraint such as `x <= 2^n-1` to an environment
containing e.g. `n in {32, 64}` or similar, we can enumerate all
values of n and add `n == 32 & x <= 2^32-1 | n == 64 & x <= 2^64-1`
instead. The exponentials then get simplified away, which means that
we stay on the SMT solver's happy path.

This is enabled by the (experimental, name subject to change) flag
-smt_linearize, as this both allows some things to typecheck that
didn't before (see pow_32_64.sail in test/typecheck/pass), but also
may potentially cause some things that typecheck to no longer
typecheck for SMT solvers like z3 that have some support for reasoning
with power functions, or where we could simply treat the exponential
as a uninterpreted function.

Also make the -dsmt_verbose option print the smtlib2 files for the
solve functions in constraint.ml. We currently ignore the -smt_solver
option for these, because smtlib does not guarantee a standard format
for the output of get-model, so we always use z3 in this case.

Following the last commit where solve got renamed solve_unique, there
are now 3 functions for solving constraints:

* Constraint.solve_smt, which finds a solution if one exists

* Constraint.solve_all_smt, which returns all possible
  solutions. Currently there's no bound, so this could loop forever if
  there are infinite solutions.

* Constraint.solve_unique, which returns a unique solution if one exists

Fix a bug where $option pragmas were dropped unnecessarily.

7 files changed, 128 insertions, 63 deletions

diff --git a/src/ast_util.ml b/src/ast_util.ml
index fcfa619e..e4287249 100644
--- a/src/ast_util.ml
+++ b/src/ast_util.ml
@@ -417,11 +417,16 @@ let nc_lteq n1 n2 = NC_aux (NC_bounded_le (n1, n2), Parse_ast.Unknown)
 let nc_gteq n1 n2 = NC_aux (NC_bounded_ge (n1, n2), Parse_ast.Unknown)
 let nc_lt n1 n2 = nc_lteq (nsum n1 (nint 1)) n2
 let nc_gt n1 n2 = nc_gteq n1 (nsum n2 (nint 1))
-let nc_or nc1 nc2 = mk_nc (NC_or (nc1, nc2))
 let nc_var kid = mk_nc (NC_var kid)
 let nc_true = mk_nc NC_true
 let nc_false = mk_nc NC_false
 
+let nc_or nc1 nc2 =
+  match nc1, nc2 with
+  | _, NC_aux (NC_false, _) -> nc1
+  | NC_aux (NC_false, _), _ -> nc2
+  | _, _ -> mk_nc (NC_or (nc1, nc2))
+
 let nc_and nc1 nc2 =
   match nc1, nc2 with
   | _, NC_aux (NC_true, _) -> nc1
@@ -439,7 +444,7 @@ let arg_kopt (KOpt_aux (KOpt_kind (K_aux (k, _), v), l)) =
   | K_order -> arg_order (Ord_aux (Ord_var v, l))
   | K_bool -> arg_bool (nc_var v)
   | K_type -> arg_typ (mk_typ (Typ_var v))
-           
+
 let nc_not nc = mk_nc (NC_app (mk_id "not", [arg_bool nc]))
 
 let mk_typschm typq typ = TypSchm_aux (TypSchm_ts (typq, typ), Parse_ast.Unknown)
diff --git a/src/constraint.ml b/src/constraint.ml
index 2c5150a5..5402f6f7 100644
--- a/src/constraint.ml
+++ b/src/constraint.ml
@@ -321,12 +321,13 @@ let call_smt l vars constraints =
   Profile.finish_smt t;
   result
 
-let rec solve_unique_smt l vars constraints var =
+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
 
-  (* prerr_endline (Printf.sprintf "SMTLIB2 constraints are: \n%s%!" smt_file);
-     prerr_endline ("Solving for " ^ smt_var); *)
+  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
@@ -355,9 +356,27 @@ let rec solve_unique_smt l vars constraints var =
   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
-    begin match call_smt l vars (nc_and constraints (nc_neq (nconstant result) (nvar var))) with
-    | Unsat -> Some result
-    | _ -> None
-    end
+    Some result
   with
-    Not_found -> None
+  | 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
diff --git a/src/constraint.mli b/src/constraint.mli
index fb6e0fcd..b5d6ff6b 100644
--- a/src/constraint.mli
+++ b/src/constraint.mli
@@ -63,4 +63,8 @@ val save_digests : unit -> unit
 
 val call_smt : l -> kind_aux KBindings.t -> n_constraint -> smt_result
 
+val solve_smt : l -> kind_aux KBindings.t -> n_constraint -> kid -> Big_int.num option
+
+val solve_all_smt : l -> kind_aux KBindings.t -> n_constraint -> kid -> Big_int.num list option
+
 val solve_unique_smt : l -> kind_aux KBindings.t -> n_constraint -> kid -> Big_int.num option
diff --git a/src/process_file.ml b/src/process_file.ml
index 52e0cd08..e7bf8d30 100644
--- a/src/process_file.ml
+++ b/src/process_file.ml
@@ -159,7 +159,7 @@ let rec preprocess opts = function
      symbols := StringSet.add symbol !symbols;
      preprocess opts defs
 
-  | Parse_ast.DEF_pragma ("option", command, l) :: defs ->
+  | (Parse_ast.DEF_pragma ("option", command, l) as opt_pragma) :: defs ->
      begin
        try
          let args = Str.split (Str.regexp " +") command in
@@ -167,7 +167,7 @@ let rec preprocess opts = function
        with
        | Arg.Bad message | Arg.Help message -> raise (Reporting.err_general l message)
      end;
-     preprocess opts defs
+     opt_pragma :: preprocess opts defs
 
   | Parse_ast.DEF_pragma ("ifndef", symbol, l) :: defs ->
      let then_defs, else_defs, defs = cond_pragma l defs in
diff --git a/src/sail.ml b/src/sail.ml
index 9cf87af8..64ccd341 100644
--- a/src/sail.ml
+++ b/src/sail.ml
@@ -115,8 +115,11 @@ let options = Arg.align ([
   ( "-smt_solver",
     Arg.String (fun s -> Constraint.set_solver (String.trim s)),
     "<solver> choose SMT solver. Supported solvers are z3 (default), alt-ergo, cvc4, mathsat, vampire and yices.");
+  ( "-smt_linearize",
+    Arg.Set Type_check.opt_smt_linearize,
+    "(experimental) force linearization for constraints involving exponentials");
   ( "-latex",
-    Arg.Tuple [Arg.Set opt_print_latex; Arg.Clear Type_check.opt_expand_valspec ],
+    Arg.Tuple [Arg.Set opt_print_latex; Arg.Clear Type_check.opt_expand_valspec],
     " pretty print the input to LaTeX");
   ( "-latex_prefix",
     Arg.String (fun prefix -> Latex.opt_prefix := prefix),
diff --git a/src/type_check.ml b/src/type_check.ml
index ada04c24..0da7f753 100644
--- a/src/type_check.ml
+++ b/src/type_check.ml
@@ -72,6 +72,10 @@ let opt_no_lexp_bounds_check = ref false
    We prefer not to do it for latex output but it is otherwise a good idea. *)
 let opt_expand_valspec = ref true
 
+(* Linearize cases involving power where we would otherwise require
+   the SMT solver to use non-linear arithmetic. *)
+let opt_smt_linearize = ref false
+
 let depth = ref 0
 
 let rec indent n = match n with
@@ -247,6 +251,50 @@ and strip_kinded_id_aux = function
 and strip_kind = function
   | K_aux (k_aux, _) -> K_aux (k_aux, Parse_ast.Unknown)
 
+let rec typ_nexps (Typ_aux (typ_aux, l)) =
+  match typ_aux with
+  | Typ_internal_unknown -> []
+  | Typ_id v -> []
+  | Typ_var kid -> []
+  | Typ_tup typs -> List.concat (List.map typ_nexps typs)
+  | Typ_app (f, args) -> List.concat (List.map typ_arg_nexps args)
+  | Typ_exist (kids, nc, typ) -> typ_nexps typ
+  | Typ_fn (arg_typs, ret_typ, _) ->
+     List.concat (List.map typ_nexps arg_typs) @ typ_nexps ret_typ
+  | Typ_bidir (typ1, typ2) ->
+     typ_nexps typ1 @ typ_nexps typ2
+and typ_arg_nexps (A_aux (typ_arg_aux, l)) =
+  match typ_arg_aux with
+  | A_nexp n -> [n]
+  | A_typ typ -> typ_nexps typ
+  | A_bool nc -> constraint_nexps nc
+  | A_order ord -> []
+and constraint_nexps (NC_aux (nc_aux, l)) =
+  match nc_aux with
+  | NC_equal (n1, n2) | NC_bounded_ge (n1, n2) | NC_bounded_le (n1, n2) | NC_not_equal (n1, n2) ->
+     [n1; n2]
+  | NC_set _ | NC_true | NC_false | NC_var _ -> []
+  | NC_or (nc1, nc2) | NC_and (nc1, nc2) -> constraint_nexps nc1 @ constraint_nexps nc2
+  | NC_app (_, args) -> List.concat (List.map typ_arg_nexps args)
+
+(* Return a KidSet containing all the type variables appearing in
+   nexp, where nexp occurs underneath a Nexp_exp, i.e. 2^nexp *)
+let rec nexp_power_variables (Nexp_aux (aux, _)) =
+  match aux with
+  | Nexp_times (n1, n2) | Nexp_sum (n1, n2) | Nexp_minus (n1, n2) ->
+     KidSet.union (nexp_power_variables n1) (nexp_power_variables n2)
+  | Nexp_neg n ->
+     nexp_power_variables n
+  | Nexp_id _ | Nexp_var _ | Nexp_constant _ ->
+     KidSet.empty
+  | Nexp_app (_, ns) ->
+     List.fold_left KidSet.union KidSet.empty (List.map nexp_power_variables ns)
+  | Nexp_exp n ->
+     tyvars_of_nexp n
+
+let constraint_power_variables nc =
+  List.fold_left KidSet.union KidSet.empty (List.map nexp_power_variables (constraint_nexps nc))
+
 let rec name_pat (P_aux (aux, _)) =
   match aux with
   | P_id id | P_as (_, id) -> Some ("_" ^ string_of_id id)
@@ -262,7 +310,7 @@ let fresh_existential k =
 let named_existential k = function
   | Some n -> mk_kopt k (mk_kid n)
   | None -> fresh_existential k
- 
+
 let destruct_exist_plain ?name:(name=None) typ =
   match typ with
   | Typ_aux (Typ_exist ([kopt], nc, typ), _) ->
@@ -1089,7 +1137,7 @@ end = struct
   let add_typ_var l (KOpt_aux (KOpt_kind (K_aux (k, _), v), _)) env =
     if KBindings.mem v env.typ_vars then begin
         let n = match KBindings.find_opt v env.shadow_vars with Some n -> n | None -> 0 in
-        let s_l, s_k = KBindings.find v env.typ_vars in 
+        let s_l, s_k = KBindings.find v env.typ_vars in
         let s_v = Kid_aux (Var (string_of_kid v ^ "#" ^ string_of_int n), l) in
         typ_print (lazy (Printf.sprintf "%stype variable (shadowing %s) %s : %s" adding (string_of_kid s_v) (string_of_kid v) (string_of_kind_aux k)));
         { env with
@@ -1108,11 +1156,34 @@ end = struct
   let add_constraint constr env =
     wf_constraint env constr;
     let (NC_aux (nc_aux, l) as constr) = constraint_simp (expand_constraint_synonyms env constr) in
-    match nc_aux with
-    | NC_true -> env
-    | _ ->
-       typ_print (lazy (adding ^ "constraint " ^ string_of_n_constraint constr));
-       { env with constraints = constr :: env.constraints }
+    let power_vars = constraint_power_variables constr in
+    if KidSet.cardinal power_vars > 1 && !opt_smt_linearize then
+      typ_error env l ("Cannot add constraint " ^ string_of_n_constraint constr
+                       ^ " where more than two variables appear within an exponential")
+    else if KidSet.cardinal power_vars = 1 && !opt_smt_linearize then
+      let v = KidSet.choose power_vars in
+      let constrs = List.fold_left nc_and nc_true (get_constraints env) in
+      begin match Constraint.solve_all_smt l (get_typ_vars env) constrs v with
+      | Some solutions ->
+         typ_print (lazy (Util.("Linearizing " |> red |> clear) ^ string_of_n_constraint constr
+                          ^ " for " ^ string_of_kid v ^ " in " ^ Util.string_of_list ", " Big_int.to_string solutions));
+         let linearized =
+           List.fold_left
+             (fun c s -> nc_or c (nc_and (nc_eq (nvar v) (nconstant s)) (constraint_subst v (arg_nexp (nconstant s)) constr)))
+             nc_false solutions
+         in
+         typ_print (lazy (adding ^ "constraint " ^ string_of_n_constraint linearized));
+         { env with constraints = linearized :: env.constraints }
+      | None ->
+         typ_error env l ("Type variable " ^ string_of_kid v
+                          ^ " must have a finite number of solutions to add " ^ string_of_n_constraint constr)
+      end
+    else
+      match nc_aux with
+      | NC_true -> env
+      | _ ->
+         typ_print (lazy (adding ^ "constraint " ^ string_of_n_constraint constr));
+         { env with constraints = constr :: env.constraints }
 
   let get_ret_typ env = env.ret_typ
 
@@ -1326,32 +1397,6 @@ and simp_typ_aux = function
 
   | typ_aux -> typ_aux
 
-let rec typ_nexps (Typ_aux (typ_aux, l)) =
-  match typ_aux with
-  | Typ_internal_unknown -> []
-  | Typ_id v -> []
-  | Typ_var kid -> []
-  | Typ_tup typs -> List.concat (List.map typ_nexps typs)
-  | Typ_app (f, args) -> List.concat (List.map typ_arg_nexps args)
-  | Typ_exist (kids, nc, typ) -> typ_nexps typ
-  | Typ_fn (arg_typs, ret_typ, _) ->
-     List.concat (List.map typ_nexps arg_typs) @ typ_nexps ret_typ
-  | Typ_bidir (typ1, typ2) ->
-     typ_nexps typ1 @ typ_nexps typ2
-and typ_arg_nexps (A_aux (typ_arg_aux, l)) =
-  match typ_arg_aux with
-  | A_nexp n -> [n]
-  | A_typ typ -> typ_nexps typ
-  | A_bool nc -> constraint_nexps nc
-  | A_order ord -> []
-and constraint_nexps (NC_aux (nc_aux, l)) =
-  match nc_aux with
-  | NC_equal (n1, n2) | NC_bounded_ge (n1, n2) | NC_bounded_le (n1, n2) | NC_not_equal (n1, n2) ->
-     [n1; n2]
-  | NC_set _ | NC_true | NC_false | NC_var _ -> []
-  | NC_or (nc1, nc2) | NC_and (nc1, nc2) -> constraint_nexps nc1 @ constraint_nexps nc2
-  | NC_app (_, args) -> List.concat (List.map typ_arg_nexps args)
-
 (* Here's how the constraint generation works for subtyping
 
 X(b,c...) --> {a. Y(a,b,c...)} \subseteq {a. Z(a,b,c...)}
@@ -1373,22 +1418,6 @@ this is equivalent to
 which is then a problem we can feed to the constraint solver expecting unsat.
  *)
 
-let rec nexp_variable_power (Nexp_aux (aux, _)) =
-  match aux with
-  | Nexp_times (n1, n2) | Nexp_sum (n1, n2) | Nexp_minus (n1, n2) ->
-     nexp_variable_power n1 || nexp_variable_power n2
-  | Nexp_neg n ->
-     nexp_variable_power n
-  | Nexp_id _ | Nexp_var _ | Nexp_constant _ ->
-     false
-  | Nexp_app (_, ns) ->
-     List.exists nexp_variable_power ns
-  | Nexp_exp n ->
-     not (KidSet.is_empty (tyvars_of_nexp n))
-
-let constraint_variable_power nc =
-  List.exists nexp_variable_power (constraint_nexps nc)
-
 let prove_smt env (NC_aux (_, l) as nc) =
   let vars = Env.get_typ_vars env in
   let vars = KBindings.filter (fun _ k -> match k with K_int | K_bool -> true | _ -> false) vars in
@@ -1400,7 +1429,7 @@ let prove_smt env (NC_aux (_, l) as nc) =
      (* Work around versions of z3 that are confused by 2^n in
         constraints, even when such constraints are irrelevant *)
      let ncs' = List.concat (List.map constraint_conj ncs) in
-     let ncs' = List.filter (fun nc -> not (constraint_variable_power nc)) ncs' in
+     let ncs' = List.filter (fun nc -> KidSet.is_empty (constraint_power_variables nc)) ncs' in
      match Constraint.call_smt l vars (List.fold_left nc_and (nc_not nc) ncs') with
      | Constraint.Unsat -> typ_debug (lazy "unsat"); true
      | Constraint.Sat | Constraint.Unknown -> typ_debug (lazy "sat/unknown"); false
diff --git a/src/type_check.mli b/src/type_check.mli
index 00439412..cfdfa2c6 100644
--- a/src/type_check.mli
+++ b/src/type_check.mli
@@ -71,6 +71,10 @@ val opt_no_lexp_bounds_check : bool ref
    We prefer not to do it for latex output but it is otherwise a good idea. *)
 val opt_expand_valspec : bool ref
 
+(** Linearize cases involving power where we would otherwise require
+   the SMT solver to use non-linear arithmetic. *)
+val opt_smt_linearize : bool ref
+
 (** {2 Type errors} *)
 
 type type_error =
@@ -316,6 +320,7 @@ val check_val_spec : Env.t -> 'a val_spec -> tannot def list * Env.t
    untrustworthy. *)
 val prove : (string * int * int * int) -> Env.t -> n_constraint -> bool
 
+(** Returns Some c if there is a unique c such that nexp = c *)
 val solve_unique : Env.t -> nexp -> Big_int.num option
 
 val canonicalize : Env.t -> typ -> typ