Use solver properly to simplify nexps in mono analysis, Lem printing

Turn on complex nexp rewriting for mono by default
(NB: solving is currently quite slow, will optimise)

4 files changed, 35 insertions, 20 deletions

diff --git a/src/monomorphise.ml b/src/monomorphise.ml
index 121b55d3..b261e485 100644
--- a/src/monomorphise.ml
+++ b/src/monomorphise.ml
@@ -2779,16 +2779,23 @@ let refine_dependency env (E_aux (e,(l,annot)) as exp) pexps =
   | _ -> None
 
 let simplify_size_nexp env typ_env (Nexp_aux (ne,l) as nexp) =
-  let is_equal kid =
-    prove typ_env (NC_aux (NC_equal (Nexp_aux (Nexp_var kid,Unknown), nexp),Unknown))
-  in
-  match ne with
-  | Nexp_var _
-  | Nexp_constant _ -> nexp
-  | _ ->
-     match List.find is_equal env.top_kids with
-     | kid -> Nexp_aux (Nexp_var kid,Generated l)
-     | exception Not_found -> nexp
+  match solve typ_env nexp with
+  | Some n -> nconstant n
+  | None ->
+     let is_equal kid =
+       prove typ_env (NC_aux (NC_equal (Nexp_aux (Nexp_var kid,Unknown), nexp),Unknown))
+     in
+     match ne with
+     | Nexp_var _
+     | Nexp_constant _ -> nexp
+     | _ ->
+        match List.find is_equal env.top_kids with
+        | kid -> Nexp_aux (Nexp_var kid,Generated l)
+        | exception Not_found -> nexp
+
+let simplify_size_uvar env typ_env = function
+  | U_nexp nexp -> U_nexp (simplify_size_nexp env typ_env nexp)
+  | x -> x
 
 (* Takes an environment of dependencies on vars, type vars, and flow control,
    and dependencies on mutable variables.  The latter are quite conservative,
@@ -2841,7 +2848,8 @@ let rec analyse_exp fn_id env assigns (E_aux (e,(l,annot)) as exp) =
     | E_cast (_,e) -> analyse_exp fn_id env assigns e
     | E_app (id,args) ->
        let deps, assigns, r = non_det args in
-       let (_,fn_typ) = Env.get_val_spec id (env_of_annot (l,annot)) in
+       let typ_env = env_of_annot (l,annot) in
+       let (_,fn_typ) = Env.get_val_spec id typ_env in
        let fn_effect = match fn_typ with
          | Typ_aux (Typ_fn (_,_,eff),_) -> eff
          | _ -> Effect_aux (Effect_set [],Unknown)
@@ -2851,6 +2859,7 @@ let rec analyse_exp fn_id env assigns (E_aux (e,(l,annot)) as exp) =
          | _ -> Unknown (l, "Effects from function application")
        in
        let kid_inst = instantiation_of exp in
+       let kid_inst = KBindings.map (simplify_size_uvar env typ_env) kid_inst in
        (* Change kids in instantiation to the canonical ones from the type signature *)
        let kid_inst = KBindings.fold (fun kid -> KBindings.add (orig_kid kid)) kid_inst KBindings.empty in
        let kid_deps = KBindings.map (deps_of_uvar l fn_id env deps) kid_inst in
diff --git a/src/pretty_print_lem.ml b/src/pretty_print_lem.ml
index 97b618e8..18fcbf69 100644
--- a/src/pretty_print_lem.ml
+++ b/src/pretty_print_lem.ml
@@ -336,11 +336,17 @@ let replace_typ_size ctxt env (Typ_aux (t,a)) =
   match t with
   | Typ_app (Id_aux (Id "vector",_) as id, [Typ_arg_aux (Typ_arg_nexp size,_);ord;typ']) ->
      begin
-       let is_equal nexp =
-         prove env (NC_aux (NC_equal (size,nexp),Parse_ast.Unknown))
-       in match List.find is_equal (NexpSet.elements ctxt.bound_nexps) with
-       | nexp -> Some (Typ_aux (Typ_app (id, [Typ_arg_aux (Typ_arg_nexp nexp,Parse_ast.Unknown);ord;typ']),a))
-       | exception Not_found -> None
+       let mk_typ nexp = 
+         Some (Typ_aux (Typ_app (id, [Typ_arg_aux (Typ_arg_nexp nexp,Parse_ast.Unknown);ord;typ']),a))
+       in
+       match Type_check.solve env size with
+       | Some n -> mk_typ (nconstant n)
+       | None ->
+          let is_equal nexp =
+            prove env (NC_aux (NC_equal (size,nexp),Parse_ast.Unknown))
+          in match List.find is_equal (NexpSet.elements ctxt.bound_nexps) with
+          | nexp -> mk_typ nexp
+          | exception Not_found -> None
      end
   | _ -> None
 
diff --git a/src/process_file.ml b/src/process_file.ml
index e9d04c5b..3f24f4c3 100644
--- a/src/process_file.ml
+++ b/src/process_file.ml
@@ -190,7 +190,7 @@ let opt_ddump_raw_mono_ast = ref false
 let opt_dmono_analysis = ref 0
 let opt_auto_mono = ref false
 let opt_mono_rewrites = ref false
-let opt_mono_complex_nexps = ref false
+let opt_mono_complex_nexps = ref true
 let opt_dall_split_errors = ref false
 let opt_dmono_continue = ref false
 
diff --git a/src/sail.ml b/src/sail.ml
index 2632f9a4..65473d86 100644
--- a/src/sail.ml
+++ b/src/sail.ml
@@ -164,9 +164,9 @@ let options = Arg.align ([
   ( "-mono_rewrites",
     Arg.Set Process_file.opt_mono_rewrites,
     " turn on rewrites for combining bitvector operations");
-  ( "-mono_complex_nexps",
-    Arg.Set Process_file.opt_mono_complex_nexps,
-    " move complex size expressions in function signatures into constraints");
+  ( "-dno_complex_nexps_rewrite",
+    Arg.Clear Process_file.opt_mono_complex_nexps,
+    " do not move complex size expressions in function signatures into constraints (monomorphisation)");
   ( "-dall_split_errors",
     Arg.Set Process_file.opt_dall_split_errors,
     " display all case split errors from monomorphisation, rather than one");