Rewriting improvements for monomorphic aarch64_small

- allow disjoint_pat to be used on patterns that have just been introduced
  by the rewrite without rechecking
- don't rebuild the matched expression in the generated fallthrough case
  (or any wildcard fall-through) - it may be dead code and generate badly
  typed Lem
- update the type environment passed to rewrites whenever type checking
  is performed; stale information broke some rewrites

6 files changed, 52 insertions, 33 deletions

diff --git a/src/isail.ml b/src/isail.ml
index 9e9b6236..88408dcd 100644
--- a/src/isail.ml
+++ b/src/isail.ml
@@ -302,7 +302,8 @@ let rec describe_rewrite =
   | String_rewriter rw -> "<string>" :: describe_rewrite (rw "")
   | Bool_rewriter rw -> "<bool>" :: describe_rewrite (rw false)
   | Literal_rewriter rw -> "(ocaml|lem|all)" :: describe_rewrite (rw (fun _ -> true))
-  | Basic_rewriter rw -> []
+  | Basic_rewriter _
+  | Checking_rewriter _ -> []
 
 type session = {
     id : string;
@@ -592,7 +593,9 @@ let handle_input' input =
               failwith "Must provide the name of a rewrite, use :list_rewrites for a list of possible rewrites"
            end
         | ":rewrites" ->
-           Interactive.ast := Process_file.rewrite_ast_target arg !Interactive.env !Interactive.ast;
+           let new_ast, new_env = Process_file.rewrite_ast_target arg !Interactive.env !Interactive.ast in
+           Interactive.ast := new_ast;
+           Interactive.env := new_env;
            interactive_state := initial_state !Interactive.ast !Interactive.env Value.primops
         | ":prover_regstate" ->
            let env, ast = prover_regstate (Some arg) !Interactive.ast !Interactive.env in
diff --git a/src/process_file.ml b/src/process_file.ml
index 7221ec25..b988895c 100644
--- a/src/process_file.ml
+++ b/src/process_file.ml
@@ -367,9 +367,9 @@ let output libpath out_arg files =
       output1 libpath out_arg f type_env defs)
     files
 
-let rewrite_step n total env defs (name, rewriter) =
+let rewrite_step n total (defs, env) (name, rewriter) =
   let t = Profile.start () in
-  let defs = rewriter env defs in
+  let defs, env = rewriter env defs in
   Profile.finish ("rewrite " ^ name) t;
   let _ = match !(opt_ddump_rewrite_ast) with
     | Some (f, i) ->
@@ -383,15 +383,15 @@ let rewrite_step n total env defs (name, rewriter) =
       end
     | _ -> () in
   Util.progress "Rewrite " name n total;
-  defs
+  defs, env
 
 let rewrite env rewriters defs =
   let total = List.length rewriters in
-  try snd (List.fold_left (fun (n, defs) rw -> n + 1, rewrite_step n total env defs rw) (1, defs) rewriters) with
+  try snd (List.fold_left (fun (n, defsenv) rw -> n + 1, rewrite_step n total defsenv rw) (1, (defs, env)) rewriters) with
   | Type_check.Type_error (_, l, err) ->
      raise (Reporting.err_typ l (Type_error.string_of_type_error err))
 
-let rewrite_ast_initial env = rewrite env [("initial", fun _ -> Rewriter.rewrite_defs)]
+let rewrite_ast_initial env = rewrite env [("initial", fun env defs -> Rewriter.rewrite_defs defs, env)]
 
 let rewrite_ast_target tgt env = rewrite env (Rewrites.rewrite_defs_target tgt)
 
diff --git a/src/process_file.mli b/src/process_file.mli
index e144727e..91cde014 100644
--- a/src/process_file.mli
+++ b/src/process_file.mli
@@ -56,9 +56,9 @@ val clear_symbols : unit -> unit
 
 val preprocess_ast : (Arg.key * Arg.spec * Arg.doc) list -> 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 rewrite_ast_initial : Type_check.Env.t -> Type_check.tannot Ast.defs -> Type_check.tannot Ast.defs
-val rewrite_ast_target : string -> Type_check.Env.t -> Type_check.tannot Ast.defs -> Type_check.tannot Ast.defs
-val rewrite_ast_check : Type_check.Env.t -> Type_check.tannot Ast.defs -> Type_check.tannot Ast.defs
+val rewrite_ast_initial : Type_check.Env.t -> Type_check.tannot Ast.defs -> Type_check.tannot Ast.defs * Type_check.Env.t
+val rewrite_ast_target : string -> Type_check.Env.t -> Type_check.tannot Ast.defs -> Type_check.tannot Ast.defs * Type_check.Env.t
+val rewrite_ast_check : Type_check.Env.t -> Type_check.tannot Ast.defs -> Type_check.tannot Ast.defs * Type_check.Env.t
 
 val load_file_no_check : (Arg.key * Arg.spec * Arg.doc) list -> Ast.order -> string -> unit Ast.defs
 val load_file : (Arg.key * Arg.spec * Arg.doc) list -> Ast.order -> Type_check.Env.t -> string -> Type_check.tannot Ast.defs * Type_check.Env.t
diff --git a/src/rewrites.ml b/src/rewrites.ml
index 64cfe48b..becf2a88 100644
--- a/src/rewrites.ml
+++ b/src/rewrites.ml
@@ -771,23 +771,23 @@ and subsumes_fpat (FP_aux (FP_Fpat (id1,pat1),_)) (FP_aux (FP_Fpat (id2,pat2),_)
 
 (* A simple check for pattern disjointness; used for optimisation in the
    guarded pattern rewrite step *)
-let rec disjoint_pat (P_aux (p1,annot1) as pat1) (P_aux (p2,annot2) as pat2) =
+let rec disjoint_pat env (P_aux (p1,annot1) as pat1) (P_aux (p2,annot2) as pat2) =
   match p1, p2 with
-  | P_as (pat1, _), _ -> disjoint_pat pat1 pat2
-  | _, P_as (pat2, _) -> disjoint_pat pat1 pat2
-  | P_typ (_, pat1), _ -> disjoint_pat pat1 pat2
-  | _, P_typ (_, pat2) -> disjoint_pat pat1 pat2
-  | P_var (pat1, _), _ -> disjoint_pat pat1 pat2
-  | _, P_var (pat2, _) -> disjoint_pat pat1 pat2
-  | P_id id, _ when id_is_unbound id (env_of_annot annot1) -> false
-  | _, P_id id when id_is_unbound id (env_of_annot annot2) -> false
+  | P_as (pat1, _), _ -> disjoint_pat env pat1 pat2
+  | _, P_as (pat2, _) -> disjoint_pat env pat1 pat2
+  | P_typ (_, pat1), _ -> disjoint_pat env pat1 pat2
+  | _, P_typ (_, pat2) -> disjoint_pat env pat1 pat2
+  | P_var (pat1, _), _ -> disjoint_pat env pat1 pat2
+  | _, P_var (pat2, _) -> disjoint_pat env pat1 pat2
+  | P_id id, _ when id_is_unbound id env -> false
+  | _, P_id id when id_is_unbound id env -> false
   | P_id id1, P_id id2 -> Id.compare id1 id2 <> 0
   | P_app (id1, args1), P_app (id2, args2) ->
-     Id.compare id1 id2 <> 0 || List.exists2 disjoint_pat args1 args2
+     Id.compare id1 id2 <> 0 || List.exists2 (disjoint_pat env) args1 args2
   | P_vector pats1, P_vector pats2
   | P_tup pats1, P_tup pats2
   | P_list pats1, P_list pats2 ->
-     List.exists2 disjoint_pat pats1 pats2
+     List.exists2 (disjoint_pat env) pats1 pats2
   | _ -> false
 
 let equiv_pats pat1 pat2 =
@@ -846,6 +846,8 @@ let case_exp e t cs =
   let env = env_of e in
   let annot = (get_loc_exp e, Some (env_of e, t, no_effect)) in
   match cs with
+  | [(P_aux (P_wild, _), body, _)] ->
+    fix_eff_exp body
   | [(P_aux (P_id id, pannot) as pat, body, _)] ->
     fix_eff_exp (annot_exp (E_let (LB_aux (LB_val (pat, e), pannot), body)) l env t)
   | _ ->
@@ -889,10 +891,12 @@ let rewrite_guarded_clauses l env pat_typ typ cs =
                 let c' = (pat',guard',body',annot) in
                 group_aux (add_clause current c') acc cs
             | None ->
-                let pat = remove_wildcards "g__" pat in
+                let pat = match cs with _::_ -> remove_wildcards "g__" pat | _ -> pat in
                 group_aux (pat,[c],annot) (acc @ [current]) cs)
       | [] -> acc @ [current]) in
     let groups = match clauses with
+      | [(pat,guard,body,annot) as c] ->
+          [(pat, [c], annot)]
       | ((pat,guard,body,annot) as c) :: cs ->
           group_aux (remove_wildcards "g__" pat, [c], annot) [] cs
       | _ ->
@@ -924,7 +928,7 @@ let rewrite_guarded_clauses l env pat_typ typ cs =
         (* For singleton clauses with a guard, use fallthrough clauses if the
            guard is not satisfied, but only those fallthrough clauses that are
            not disjoint with the current pattern *)
-        let overlapping_clause (pat, _, _) = not (disjoint_pat current_pat pat) in
+        let overlapping_clause (pat, _, _) = not (disjoint_pat env current_pat pat) in
         let fallthrough = List.filter overlapping_clause fallthrough in
         (match guard, fallthrough with
           | Some exp, _ :: _ ->
@@ -4679,11 +4683,11 @@ let rewrite_loops_with_escape_effect env defs =
   in
   rewrite_defs_base { rewriters_base with rewrite_exp } defs
 
-let recheck_defs env defs = fst (Type_error.check initial_env defs)
+let recheck_defs env defs = Type_error.check initial_env defs
 let recheck_defs_without_effects env defs =
   let old = !opt_no_effects in
   let () = opt_no_effects := true in
-  let result,_ = Type_error.check initial_env defs in
+  let result = Type_error.check initial_env defs in
   let () = opt_no_effects := old in
   result
 
@@ -4783,12 +4787,20 @@ let if_mono f env defs =
   | [], false -> defs
   | _, _ -> f env defs
 
+let if_mono_env f env defs =
+  match !opt_mono_split, !opt_auto_mono with
+  | [], false -> defs, env
+  | _, _ -> f env defs
+
 (* Also turn mwords stages on when we're just trying out mono *)
 let if_mwords f env defs =
   if !Pretty_print_lem.opt_mwords then f env defs else if_mono f env defs
+let if_mwords_env f env defs =
+  if !Pretty_print_lem.opt_mwords then f env defs else if_mono_env f env defs
 
 type rewriter =
   | Basic_rewriter of (Env.t -> tannot defs -> tannot defs)
+  | Checking_rewriter of (Env.t -> tannot defs -> tannot defs * Env.t)
   | Bool_rewriter of (bool -> rewriter)
   | String_rewriter of (string -> rewriter)
   | Literal_rewriter of ((lit -> bool) -> rewriter)
@@ -4812,6 +4824,8 @@ let instantiate_rewrite rewriter args =
     match rewriter, arg with
     | Basic_rewriter rw, If_mono_arg -> Basic_rewriter (if_mono rw)
     | Basic_rewriter rw, If_mwords_arg -> Basic_rewriter (if_mwords rw)
+    | Checking_rewriter rw, If_mono_arg -> Checking_rewriter (if_mono_env rw)
+    | Checking_rewriter rw, If_mwords_arg -> Checking_rewriter (if_mwords_env rw)
     | Bool_rewriter rw, Bool_arg b -> rw b
     | String_rewriter rw, String_arg str -> rw str
     | Literal_rewriter rw, Literal_arg selector -> rw (selector_function selector)
@@ -4819,14 +4833,15 @@ let instantiate_rewrite rewriter args =
        raise (Reporting.err_unreachable Parse_ast.Unknown __POS__ "Invalid rewrite argument")
   in
   match List.fold_left instantiate rewriter args with
-  | Basic_rewriter rw -> rw
+  | Basic_rewriter rw -> fun env defs -> rw env defs, env
+  | Checking_rewriter rw -> rw
   | _ ->
      raise (Reporting.err_general Parse_ast.Unknown "Rewrite not fully instantiated")
 
 let all_rewrites = [
     ("no_effect_check", Basic_rewriter (fun _ defs -> opt_no_effects := true; defs));
-    ("recheck_defs", Basic_rewriter recheck_defs);
-    ("recheck_defs_without_effects", Basic_rewriter recheck_defs_without_effects);
+    ("recheck_defs", Checking_rewriter recheck_defs);
+    ("recheck_defs_without_effects", Checking_rewriter recheck_defs_without_effects);
     ("optimize_recheck_defs", Basic_rewriter (fun _ -> Optimize.recheck));
     ("realise_mappings", Basic_rewriter rewrite_defs_realise_mappings);
     ("remove_mapping_valspecs", Basic_rewriter remove_mapping_valspecs);
@@ -5083,5 +5098,5 @@ let rewrite_check_annot =
                                           rewrite_pat = (fun _ -> check_pat) }
 
 let rewrite_defs_check = [
-  ("check_annotations", fun _ -> rewrite_check_annot);
+  ("check_annotations", fun env defs -> rewrite_check_annot defs, env);
   ]
diff --git a/src/rewrites.mli b/src/rewrites.mli
index e30a4206..3b572d51 100644
--- a/src/rewrites.mli
+++ b/src/rewrites.mli
@@ -70,10 +70,11 @@ val move_loop_measures : 'a defs -> 'a defs
 val rewrite_undefined : bool -> Env.t -> tannot defs -> tannot defs
 
 (* Perform rewrites to create an AST supported for a specific target *)
-val rewrite_defs_target : string -> (string * (Env.t -> tannot defs -> tannot defs)) list
+val rewrite_defs_target : string -> (string * (Env.t -> tannot defs -> tannot defs * Env.t)) list
 
 type rewriter =
   | Basic_rewriter of (Env.t -> tannot defs -> tannot defs)
+  | Checking_rewriter of (Env.t -> tannot defs -> tannot defs * Env.t)
   | Bool_rewriter of (bool -> rewriter)
   | String_rewriter of (string -> rewriter)
   | Literal_rewriter of ((lit -> bool) -> rewriter)
@@ -96,6 +97,6 @@ val opt_coq_warn_nonexhaustive : bool ref
 
 (* This is a special rewriter pass that checks AST invariants without
    actually doing any re-writing *)
-val rewrite_defs_check : (string * (Env.t -> tannot defs -> tannot defs)) list
+val rewrite_defs_check : (string * (Env.t -> tannot defs -> tannot defs * Env.t)) list
 
 val simple_typ : typ -> typ
diff --git a/src/sail.ml b/src/sail.ml
index 1ee8214f..e9b1914d 100644
--- a/src/sail.ml
+++ b/src/sail.ml
@@ -409,7 +409,7 @@ let load_files ?check:(check=false) type_envs files =
     ("out.sail", ast, type_envs)
   else
     let ast = Scattered.descatter ast in
-    let ast = rewrite_ast_initial type_envs ast in
+    let ast, type_envs = rewrite_ast_initial type_envs ast in
 
     let out_name = match !opt_file_out with
       | None when parsed = [] -> "out.sail"
@@ -579,7 +579,7 @@ let main () =
       else ();
 
       let type_envs, ast = prover_regstate !opt_target ast type_envs in
-      let ast = match !opt_target with Some tgt -> rewrite_ast_target tgt type_envs ast | None -> ast in
+      let ast, type_envs = match !opt_target with Some tgt -> rewrite_ast_target tgt type_envs ast | None -> ast, type_envs in
       target !opt_target out_name ast type_envs;
 
       if !Interactive.opt_interactive then