Move monomorphisation after mapping rewrites

Fixes monomorphisation on files using mappings.
Also extended constant propagation to handle pattern matches on
bitvector expressions (because an earlier rewrite replaces the literals).
Also moved L_undef rewriting because monomorphisation can handle them
but not the replacement functions.

7 files changed, 136 insertions, 105 deletions

diff --git a/src/monomorphise.ml b/src/monomorphise.ml
index f4f28c41..7886fe6b 100644
--- a/src/monomorphise.ml
+++ b/src/monomorphise.ml
@@ -1591,6 +1591,31 @@ let split_defs all_errors splits defs =
             (Reporting_basic.print_err false true l' "Monomorphisation"
                "Unexpected kind of pattern for literal"; GiveUp)
        in findpat_generic checkpat "literal" assigns cases
+    | E_vector es when List.for_all (function (E_aux (E_lit _,_)) -> true | _ -> false) es ->
+       let checkpat = function
+         | P_aux (P_vector ps,_) ->
+            let matches = List.map2 (fun e p ->
+              match e, p with
+              | E_aux (E_lit (L_aux (lit,_)),_), P_aux (P_lit (L_aux (lit',_)),_) ->
+                 if lit_match (lit,lit') then DoesMatch ([],[]) else DoesNotMatch
+              | E_aux (E_lit l,_), P_aux (P_id var,_) when pat_id_is_variable env var ->
+                 DoesMatch ([var, e],[])
+              | _ -> GiveUp) es ps in
+            let final = List.fold_left (fun acc m -> match acc, m with
+              | _, GiveUp -> GiveUp
+              | GiveUp, _ -> GiveUp
+              | DoesMatch (sub,ksub), DoesMatch(sub',ksub') -> DoesMatch(sub@sub',ksub@ksub')
+              | _ -> DoesNotMatch) (DoesMatch ([],[])) matches in
+            (match final with
+            | GiveUp ->
+               (Reporting_basic.print_err false true l "Monomorphisation"
+                  "Unexpected kind of pattern for vector literal"; GiveUp)
+            | _ -> final)
+         | _ ->
+            (Reporting_basic.print_err false true l "Monomorphisation"
+               "Unexpected kind of pattern for vector literal"; GiveUp)
+       in findpat_generic checkpat "vector literal" assigns cases
+
     | E_cast (undef_typ, (E_aux (E_lit (L_aux (L_undef, lit_l)),_) as e_undef)) ->
        let checkpat = function
          | P_aux (P_lit (L_aux (lit_p, _)),_) -> DoesNotMatch
@@ -1780,21 +1805,31 @@ let split_defs all_errors splits defs =
               otherwise *)
            | Some (Some (pats,l)) ->
               let max = List.length pats - 1 in
+              let lit_like = function
+                | P_lit _ -> true
+                | P_vector ps -> List.for_all (function P_aux (P_lit _,_) -> true | _ -> false) ps
+                | _ -> false
+              in
+              let rec to_exp = function
+                | P_aux (P_lit lit,(l,ann)) -> E_aux (E_lit lit,(Generated l,ann))
+                | P_aux (P_vector ps,(l,ann)) -> E_aux (E_vector (List.map to_exp ps),(Generated l,ann))
+                | _ -> assert false
+              in
               Some (List.mapi (fun i p ->
                 match p with
-                | P_aux (P_lit lit,(pl,pannot))
-                    when (match lit with L_aux (L_undef,_) -> false | _ -> true) ->
+                | P_aux (P_lit (L_aux (L_num j,_) as lit),(pl,pannot)) ->
                    let orig_typ = Env.base_typ_of (env_of_annot (l,annot)) (typ_of_annot (l,annot)) in
-                   let kid_subst = match lit, orig_typ with
-                     | L_aux (L_num i,_),
-                       Typ_aux
+                   let kid_subst = match orig_typ with
+                     | Typ_aux
                          (Typ_app (Id_aux (Id "atom",_),
                                    [Typ_arg_aux (Typ_arg_nexp
                                                    (Nexp_aux (Nexp_var var,_)),_)]),_) ->
-                        [var,nconstant i]
+                        [var,nconstant j]
                      | _ -> []
                    in
                    p,[id,E_aux (E_lit lit,(Generated pl,pannot))],[l,(i,max,[])],kid_subst
+                | P_aux (p',(pl,pannot)) when lit_like p' ->
+                   p,[id,to_exp p],[l,(i,max,[])],[]
                 | _ ->
                    let p',subst = freshen_pat_bindings p in
                    match p' with
@@ -4193,12 +4228,8 @@ let rewrite_toplevel_nexps (Defs defs) =
 type options = {
   auto : bool;
   debug_analysis : int;
-  rewrites : bool;
-  rewrite_toplevel_nexps : bool;
-  rewrite_size_parameters : bool;
   all_split_errors : bool;
-  continue_anyway : bool;
-  dump_raw: bool
+  continue_anyway : bool
 }
 
 let recheck defs =
@@ -4208,18 +4239,10 @@ let recheck defs =
   let () = Util.opt_warnings := w in
   r
 
-let monomorphise opts splits env defs =
-  let (defs,env) =
-    if opts.rewrites then
-      let defs = MonoRewrites.mono_rewrite defs in
-      recheck defs
-    else defs,env
-  in
-  let defs,env =
-    if opts.rewrite_toplevel_nexps
-      then recheck (rewrite_toplevel_nexps defs)
-      else defs,env
-  in
+let mono_rewrites = MonoRewrites.mono_rewrite
+
+let monomorphise opts splits defs =
+  let defs, env = Type_check.check Type_check.initial_env defs in
   let ok_analysis, new_splits, extra_splits =
     if opts.auto
     then
@@ -4239,18 +4262,9 @@ let monomorphise opts splits env defs =
   let () = if (ok_analysis && ok_extras && ok_split) || opts.continue_anyway
       then ()
       else raise (Reporting_basic.err_general Unknown "Unable to monomorphise program")
-  in
-  let defs,env = recheck defs in
-  let defs = BitvectorSizeCasts.add_bitvector_casts defs in
-  (* TODO: currently doing this because constant propagation leaves numeric literals as
-     int, try to avoid this later; also use final env for DEF_spec case above, because the
-     type checker doesn't store the env at that point :( *)
-  let defs = if opts.rewrite_size_parameters then
-      let (defs,env) = recheck defs in
-      let defs = AtomToItself.rewrite_size_parameters env defs in
-      defs
-    else
-      defs
-  in
-  let () = if opts.dump_raw then Pretty_print_sail.pp_defs stdout defs else () in
-  recheck defs
+  in defs
+
+let add_bitvector_casts = BitvectorSizeCasts.add_bitvector_casts
+let rewrite_atoms_to_singletons defs =
+  let defs, env = Type_check.check Type_check.initial_env defs in
+  AtomToItself.rewrite_size_parameters env defs
diff --git a/src/monomorphise.mli b/src/monomorphise.mli
index 3baa7d12..1a82c8d0 100644
--- a/src/monomorphise.mli
+++ b/src/monomorphise.mli
@@ -51,17 +51,24 @@
 type options = {
   auto : bool;                    (* Analyse ast to find splits for monomorphisation *)
   debug_analysis : int;           (* Debug output level for the automatic analysis *)
-  rewrites : bool;                (* Experimental rewrites for variable-sized operations *)
-  rewrite_toplevel_nexps : bool;  (* Move complex nexps in function signatures into constraints *)
-  rewrite_size_parameters : bool; (* Make implicit type parameters explicit for (e.g.) lem *)
   all_split_errors : bool;
-  continue_anyway : bool;
-  dump_raw: bool
+  continue_anyway : bool
 }
 
 val monomorphise :
   options ->
   ((string * int) * string) list -> (* List of splits from the command line *)
-  Type_check.Env.t ->
   Type_check.tannot Ast.defs ->
-  Type_check.tannot Ast.defs * Type_check.Env.t
+  Type_check.tannot Ast.defs
+
+(* Rewrite (combinations of) variable-sized operations into fixed-sized operations *)
+val mono_rewrites : Type_check.tannot Ast.defs -> Type_check.tannot Ast.defs
+
+(* Move complex nexps in function signatures into constraints *)
+val rewrite_toplevel_nexps : Type_check.tannot Ast.defs -> Type_check.tannot Ast.defs
+
+(* Add casts across case splits *)
+val add_bitvector_casts : Type_check.tannot Ast.defs -> Type_check.tannot Ast.defs
+
+(* Replace atom arguments which are fixed by a type parameter for a size with a singleton type *)
+val rewrite_atoms_to_singletons : Type_check.tannot Ast.defs -> Type_check.tannot Ast.defs
diff --git a/src/process_file.ml b/src/process_file.ml
index c3e1b510..6afbae3d 100644
--- a/src/process_file.ml
+++ b/src/process_file.ml
@@ -232,27 +232,6 @@ let check_ast (env : Type_check.Env.t) (defs : unit Ast.defs) : Type_check.tanno
   let () = if !opt_just_check then exit 0 else () in
   (ast, env)
 
-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 true
-let opt_dall_split_errors = ref false
-let opt_dmono_continue = ref false
-
-let monomorphise_ast locs type_env ast =
-  let open Monomorphise in
-  let opts = {
-    auto = !opt_auto_mono;
-    debug_analysis = !opt_dmono_analysis;
-    rewrites = !opt_mono_rewrites;
-    rewrite_toplevel_nexps = !opt_mono_complex_nexps;
-    rewrite_size_parameters = !Pretty_print_lem.opt_mwords;
-    all_split_errors = !opt_dall_split_errors;
-    continue_anyway = !opt_dmono_continue;
-    dump_raw = !opt_ddump_raw_mono_ast
-  } in
-  monomorphise opts locs type_env ast
 
 let open_output_with_check file_name =
   let (temp_file_name, o) = Filename.open_temp_file "ll_temp" "" in
diff --git a/src/process_file.mli b/src/process_file.mli
index 0fec4064..a4e31890 100644
--- a/src/process_file.mli
+++ b/src/process_file.mli
@@ -52,7 +52,6 @@ val parse_file : string -> Parse_ast.defs
 val convert_ast : Ast.order -> Parse_ast.defs -> unit Ast.defs
 val preprocess_ast : Parse_ast.defs -> Parse_ast.defs
 val check_ast: Type_check.Env.t -> unit Ast.defs -> Type_check.tannot Ast.defs * Type_check.Env.t
-val monomorphise_ast : ((string * int) * string) list -> Type_check.Env.t -> Type_check.tannot Ast.defs -> Type_check.tannot Ast.defs * Type_check.Env.t
 val rewrite_ast: Type_check.tannot Ast.defs -> Type_check.tannot Ast.defs
 val rewrite_undefined: bool -> Type_check.tannot Ast.defs -> Type_check.tannot Ast.defs
 val rewrite_ast_lem : Type_check.tannot Ast.defs -> Type_check.tannot Ast.defs
@@ -69,13 +68,6 @@ val opt_just_check : bool ref
 val opt_ddump_tc_ast : bool ref
 val opt_ddump_rewrite_ast : ((string * int) option) ref
 val opt_dno_cast : bool ref
-val opt_ddump_raw_mono_ast : bool ref
-val opt_dmono_analysis : int ref
-val opt_dmono_continue : bool ref
-val opt_auto_mono : bool ref
-val opt_mono_rewrites : bool ref
-val opt_mono_complex_nexps : bool ref
-val opt_dall_split_errors : bool ref
 
 type out_type =
   | Lem_ast_out
diff --git a/src/rewrites.ml b/src/rewrites.ml
index 246a2670..2faebf9c 100644
--- a/src/rewrites.ml
+++ b/src/rewrites.ml
@@ -2336,6 +2336,11 @@ let rewrite_undefined mwords =
   let rewrite_exp_undefined = { id_exp_alg with e_aux = (fun (exp, annot) -> rewrite_e_aux (E_aux (exp, annot))) } in
   rewrite_defs_base { rewriters_base with rewrite_exp = (fun _ -> fold_exp rewrite_exp_undefined) }
 
+let rewrite_undefined_if_gen always_bitvector defs =
+  if !Initial_check.opt_undefined_gen
+  then rewrite_undefined (always_bitvector || !Pretty_print_lem.opt_mwords) defs
+  else defs
+
 let rec simple_typ (Typ_aux (typ_aux, l) as typ) = Typ_aux (simple_typ_aux typ_aux, l)
 and simple_typ_aux = function
   | Typ_id id -> Typ_id id
@@ -4269,11 +4274,54 @@ let remove_mapping_valspecs (Defs defs) =
   Defs (List.filter allowed_def defs)
 
 
+let opt_mono_rewrites = ref false
+let opt_mono_complex_nexps = ref true
+
+let mono_rewrites defs =
+  if !opt_mono_rewrites then
+    Monomorphise.mono_rewrites defs
+  else defs
+
+let rewrite_toplevel_nexps defs =
+  if !opt_mono_complex_nexps then
+    Monomorphise.rewrite_toplevel_nexps defs
+  else defs
+
+let opt_mono_split = ref ([]:((string * int) * string) list)
+let opt_dmono_analysis = ref 0
+let opt_auto_mono = ref false
+let opt_dall_split_errors = ref false
+let opt_dmono_continue = ref false
+
+let monomorphise defs =
+  let open Monomorphise in
+  monomorphise
+    { auto = !opt_auto_mono;
+      debug_analysis = !opt_dmono_analysis;
+      all_split_errors = !opt_dall_split_errors;
+      continue_anyway = !opt_dmono_continue }
+    !opt_mono_split
+    defs
+
+let if_mono f defs =
+  match !opt_mono_split, !opt_auto_mono with
+  | [], false -> defs
+  | _, _ -> f defs
+
 let rewrite_defs_lem = [
   ("realise_mappings", rewrite_defs_realise_mappings);
   ("remove_mapping_valspecs", remove_mapping_valspecs);
   ("pat_string_append", rewrite_defs_pat_string_append);
   ("mapping_builtins", rewrite_defs_mapping_patterns);
+  ("mono_rewrites", mono_rewrites);
+  ("recheck_defs", if_mono recheck_defs);
+  ("rewrite_toplevel_nexps", if_mono rewrite_toplevel_nexps);
+  ("monomorphise", if_mono monomorphise);
+  ("recheck_defs", if_mono recheck_defs);
+  ("add_bitvector_casts", if_mono Monomorphise.add_bitvector_casts);
+  ("rewrite_atoms_to_singletons", if_mono Monomorphise.rewrite_atoms_to_singletons);
+  ("recheck_defs", if_mono recheck_defs);
+  ("rewrite_undefined", rewrite_undefined_if_gen false);
   ("pat_lits", rewrite_defs_pat_lits rewrite_lit_lem);
   ("vector_concat_assignments", rewrite_vector_concat_assignments);
   ("tuple_assignments", rewrite_tuple_assignments);
@@ -4311,6 +4359,7 @@ let rewrite_defs_coq = [
   ("remove_mapping_valspecs", remove_mapping_valspecs);
   ("pat_string_append", rewrite_defs_pat_string_append);
   ("mapping_builtins", rewrite_defs_mapping_patterns);
+  ("rewrite_undefined", rewrite_undefined_if_gen true);
   ("pat_lits", rewrite_defs_pat_lits rewrite_lit_lem);
   ("vector_concat_assignments", rewrite_vector_concat_assignments);
   ("tuple_assignments", rewrite_tuple_assignments);
@@ -4351,6 +4400,7 @@ let rewrite_defs_ocaml = [
   ("realise_mappings", rewrite_defs_realise_mappings);
   ("pat_string_append", rewrite_defs_pat_string_append);
   ("mapping_builtins", rewrite_defs_mapping_patterns);
+  ("rewrite_undefined", rewrite_undefined_if_gen false);
   ("pat_lits", rewrite_defs_pat_lits rewrite_no_strings);
   ("vector_concat_assignments", rewrite_vector_concat_assignments);
   ("tuple_assignments", rewrite_tuple_assignments);
@@ -4373,6 +4423,7 @@ let rewrite_defs_c = [
   ("realise_mappings", rewrite_defs_realise_mappings);
   ("pat_string_append", rewrite_defs_pat_string_append);
   ("mapping_builtins", rewrite_defs_mapping_patterns);
+  ("rewrite_undefined", rewrite_undefined_if_gen false);
   ("pat_lits", rewrite_defs_pat_lits rewrite_no_strings);
   ("vector_concat_assignments", rewrite_vector_concat_assignments);
   ("tuple_assignments", rewrite_tuple_assignments);
@@ -4393,6 +4444,7 @@ let rewrite_defs_interpreter = [
     ("realise_mappings", rewrite_defs_realise_mappings);
     ("pat_string_append", rewrite_defs_pat_string_append);
     ("mapping_builtins", rewrite_defs_mapping_patterns);
+    ("rewrite_undefined", rewrite_undefined_if_gen false);
     ("vector_concat_assignments", rewrite_vector_concat_assignments);
     ("tuple_assignments", rewrite_tuple_assignments);
     ("simple_assignments", rewrite_simple_assignments);
diff --git a/src/rewrites.mli b/src/rewrites.mli
index 7d6bc0b2..aa793cb4 100644
--- a/src/rewrites.mli
+++ b/src/rewrites.mli
@@ -51,6 +51,15 @@
 open Ast
 open Type_check
 
+(* Monomorphisation options *)
+val opt_mono_rewrites : bool ref
+val opt_mono_complex_nexps : bool ref
+val opt_mono_split : ((string * int) * string) list ref
+val opt_dmono_analysis : int ref
+val opt_auto_mono : bool ref
+val opt_dall_split_errors : bool ref
+val opt_dmono_continue : bool ref
+
 (* Generate a fresh id with the given prefix *)
 val fresh_id : string -> l -> id
 
diff --git a/src/sail.ml b/src/sail.ml
index f0903146..6f94db4b 100644
--- a/src/sail.ml
+++ b/src/sail.ml
@@ -70,7 +70,6 @@ let opt_sanity = ref false
 let opt_libs_lem = ref ([]:string list)
 let opt_libs_coq = ref ([]:string list)
 let opt_file_arguments = ref ([]:string list)
-let opt_mono_split = ref ([]:((string * int) * string) list)
 let opt_process_elf : string option ref = ref None
 
 let options = Arg.align ([
@@ -161,19 +160,10 @@ let options = Arg.align ([
     Arg.String (fun s ->
       let l = Util.split_on_char ':' s in
       match l with
-      | [fname;line;var] -> opt_mono_split := ((fname,int_of_string line),var)::!opt_mono_split
+      | [fname;line;var] ->
+         Rewrites.opt_mono_split := ((fname,int_of_string line),var)::!Rewrites.opt_mono_split
       | _ -> raise (Arg.Bad (s ^ " not of form <filename>:<line>:<variable>"))),
       "<filename>:<line>:<variable> to case split for monomorphisation");
-  (* AA: Should use _ to be consistent with other options, but I keep
-     this case to make sure nothing breaks immediately. *)
-  ( "-mono-split",
-    Arg.String (fun s ->
-      prerr_endline (("Warning" |> Util.yellow |> Util.clear) ^ ": use -mono_split instead");
-      let l = Util.split_on_char ':' s in
-      match l with
-      | [fname;line;var] -> opt_mono_split := ((fname,int_of_string line),var)::!opt_mono_split
-      | _ -> raise (Arg.Bad (s ^ " not of form <filename>:<line>:<variable>"))),
-    "<filename>:<line>:<variable> to case split for monomorphisation");
   ( "-memo_z3",
     Arg.Set opt_memo_z3,
     " memoize calls to z3, improving performance when typechecking repeatedly");
@@ -192,26 +182,23 @@ let options = Arg.align ([
   ( "-just_check",
     Arg.Set opt_just_check,
     " (experimental) terminate immediately after typechecking");
-  ( "-ddump_raw_mono_ast",
-    Arg.Set opt_ddump_raw_mono_ast,
-    " (debug) dump the monomorphised ast before type-checking");
   ( "-dmono_analysis",
-    Arg.Set_int opt_dmono_analysis,
+    Arg.Set_int Rewrites.opt_dmono_analysis,
     " (debug) dump information about monomorphisation analysis: 0 silent, 3 max");
   ( "-auto_mono",
-    Arg.Set opt_auto_mono,
+    Arg.Set Rewrites.opt_auto_mono,
     " automatically infer how to monomorphise code");
   ( "-mono_rewrites",
-    Arg.Set Process_file.opt_mono_rewrites,
+    Arg.Set Rewrites.opt_mono_rewrites,
     " turn on rewrites for combining bitvector operations");
   ( "-dno_complex_nexps_rewrite",
-    Arg.Clear Process_file.opt_mono_complex_nexps,
+    Arg.Clear Rewrites.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,
+    Arg.Set Rewrites.opt_dall_split_errors,
     " display all case split errors from monomorphisation, rather than one");
   ( "-dmono_continue",
-    Arg.Set Process_file.opt_dmono_continue,
+    Arg.Set Rewrites.opt_dmono_continue,
     " continue despite monomorphisation errors");
   ( "-verbose",
     Arg.Set opt_print_verbose,
@@ -268,16 +255,7 @@ let load_files type_envs files =
 
   let (ast, type_envs) = check_ast type_envs ast in
 
-  let (ast, type_envs) =
-    match !opt_mono_split, !opt_auto_mono with
-    | [], false -> ast, type_envs
-    | locs, _ -> monomorphise_ast locs type_envs ast
-  in
-
-  let ast =
-    if !Initial_check.opt_undefined_gen then
-      rewrite_undefined (!Pretty_print_lem.opt_mwords || !opt_print_coq) (rewrite_ast ast)
-    else rewrite_ast ast in
+  let ast = rewrite_ast ast in
 
   let out_name = match !opt_file_out with
     | None when parsed = [] -> "out.sail"