Add rewriting step for moving execute clauses into auxiliary functions

For example, generates an auxiliary function execute_ADD (rs, rt, rd) for the
clause execute (ADD (rs,rt,rd)) = ...

Without this rewriting, the execute function easily becomes too large to be
handled by Isabelle (e.g., for CHERI-MIPS; for MIPS alone, it seems to be just
about small enough).

This used to be implemented in the pretty-printer, but that code was commented
out recently in order to support a recursive execute function for RISC-V
compressed instructions.

4 files changed, 119 insertions, 84 deletions

diff --git a/src/ast_util.ml b/src/ast_util.ml
index 591e8df2..82ef660d 100644
--- a/src/ast_util.ml
+++ b/src/ast_util.ml
@@ -940,6 +940,11 @@ and tyvars_of_typ_arg (Typ_arg_aux (ta,_)) =
   | Typ_arg_typ typ -> tyvars_of_typ typ
   | Typ_arg_order _ -> KidSet.empty
 
+let tyvars_of_quant_item (QI_aux (qi, _)) = match qi with
+  | QI_id (KOpt_aux ((KOpt_none kid | KOpt_kind (_, kid)), _)) ->
+     KidSet.singleton kid
+  | QI_const nc -> tyvars_of_nc nc
+
 let is_kid_generated kid = String.contains (string_of_kid kid) '#'
 
 let rec undefined_of_typ mwords l annot (Typ_aux (typ_aux, _) as typ) =
diff --git a/src/ast_util.mli b/src/ast_util.mli
index de925fc3..89004b0d 100644
--- a/src/ast_util.mli
+++ b/src/ast_util.mli
@@ -296,6 +296,7 @@ val union_effects : effect -> effect -> effect
 
 val tyvars_of_nexp : nexp -> KidSet.t
 val tyvars_of_typ : typ -> KidSet.t
+val tyvars_of_quant_item : quant_item -> KidSet.t
 val is_kid_generated : kid -> bool
 
 val undefined_of_typ : bool -> Ast.l -> (typ -> 'annot) -> typ -> 'annot exp
diff --git a/src/pretty_print_lem.ml b/src/pretty_print_lem.ml
index 4bb3b553..d347c1bd 100644
--- a/src/pretty_print_lem.ml
+++ b/src/pretty_print_lem.ml
@@ -1213,71 +1213,6 @@ let doc_mutrec_lem = function
 let rec doc_fundef_lem (FD_aux(FD_function(r, typa, efa, fcls),fannot) as fd) =
   match fcls with
   | [] -> failwith "FD_function with empty function list"
-  (* TODO: Move splitting of execute function to the rewriter *)
-  (*| FCL_aux (FCL_Funcl(id,_),_) :: _
-    when string_of_id id = "execute" (*|| string_of_id id = "initial_analysis"*) ->
-    let (_,auxiliary_functions,clauses) =
-      List.fold_left
-        (fun (already_used_fnames,auxiliary_functions,clauses) funcl ->
-          match funcl with
-          | FCL_aux (FCL_Funcl (Id_aux (Id _,l),pexp),annot) ->
-             let pat,guard,exp,(pexp_l,_) = destruct_pexp pexp in
-             let _ = match guard with
-               | None -> ()
-               | _ ->
-                  raise (Reporting_basic.err_unreachable pexp_l
-                           "guarded pattern expression should have been rewritten before pretty-printing")
-             in
-             let ctor, l, argspat, pannot = (match pat with
-               | P_aux (P_app (Id_aux (Id ctor,l),argspat),pannot) ->
-                 (ctor, l, argspat, pannot)
-               | P_aux (P_id (Id_aux (Id ctor,l)), pannot) ->
-                 (ctor, l, [], pannot)
-               | _ ->
-                 raise (Reporting_basic.err_unreachable l
-                   "unsupported parameter pattern in function clause")) in
-             let rec pick_name_not_clashing_with already_used candidate =
-               if StringSet.mem candidate already_used then
-                 pick_name_not_clashing_with already_used (candidate ^ "'")
-               else candidate in
-             let unit_pat = P_aux (P_lit (mk_lit L_unit), (l, Some (Env.empty, unit_typ, no_effect))) in
-             let aux_fname = pick_name_not_clashing_with already_used_fnames (string_of_id id ^ "_" ^ ctor) in
-             let aux_args = if argspat = [] then unit_pat else P_aux (P_tup argspat,pannot) in
-             let already_used_fnames = StringSet.add aux_fname already_used_fnames in
-             let fcl = FCL_aux (FCL_Funcl (Id_aux (Id aux_fname,l),
-                                           Pat_aux (Pat_exp (aux_args,exp),(pexp_l,None)))
-                                  ,annot) in
-             let auxiliary_functions =
-                auxiliary_functions ^^ hardline ^^ hardline ^^
-                  doc_fundef_lem (FD_aux (FD_function(r,typa,efa,[fcl]),fannot)) in
-             (* Bind complex patterns to names so that we can pass them to the
-                auxiliary function *)
-             let name_pat idx (P_aux (p,a)) = match p with
-               | P_as (pat,_) -> P_aux (p,a) (* already named *)
-               | P_lit _ -> P_aux (p,a) (* no need to name a literal *)
-               | P_id _ -> P_aux (p,a) (* no need to name an identifier *)
-               | _ -> P_aux (P_as (P_aux (p,a), Id_aux (Id ("arg" ^ string_of_int idx),l)),a) in
-             let named_argspat = List.mapi name_pat argspat in
-             let named_pat = P_aux (P_app (Id_aux (Id ctor,l),named_argspat),pannot) in
-             let named_args = if argspat = [] then [unit_pat] else named_argspat in
-             let doc_arg idx (P_aux (p,(l,a))) = match p with
-               | P_as (pat,id) -> doc_id_lem id
-               | P_lit lit -> doc_lit_lem lit
-               | P_id id -> doc_id_lem id
-               | _ -> string ("arg" ^ string_of_int idx) in
-             let clauses =
-               clauses ^^ (break 1) ^^
-                 (separate space
-                    [pipe;doc_pat_lem empty_ctxt false named_pat;arrow;
-                     string aux_fname;
-                     separate space (List.mapi doc_arg named_args)]) in
-             (already_used_fnames,auxiliary_functions,clauses)
-        ) (StringSet.empty,empty,empty) fcls in
-
-    auxiliary_functions ^^ hardline ^^ hardline ^^
-    (prefix 2 1)
-      ((separate space) [string "let" ^^ doc_rec_lem false r ^^ doc_id_lem id;equals;string "function"])
-      (clauses ^/^ string "end")*)
   | FCL_aux (FCL_Funcl(id,_),annot) :: _
     when not (Env.is_extern id (env_of_annot annot) "lem") ->
     string "let" ^^ (doc_rec_lem (List.length fcls > 1) r) ^^ (doc_fundef_rhs_lem fd)
diff --git a/src/rewrites.ml b/src/rewrites.ml
index af621b47..30230710 100644
--- a/src/rewrites.ml
+++ b/src/rewrites.ml
@@ -1783,6 +1783,118 @@ let rewrite_defs_early_return (Defs defs) =
     { rewriters_base with rewrite_fun = rewrite_fun_early_return }
     (Defs (early_ret_spec @ defs))
 
+let swaptyp typ (l,tannot) = match tannot with
+  | Some (env, typ', eff) -> (l, Some (env, typ, eff))
+  | _ -> raise (Reporting_basic.err_unreachable l "swaptyp called with empty type annotation")
+
+let is_funcl_rec (FCL_aux (FCL_Funcl (id, pexp), _)) =
+  let pat,guard,exp,pannot = destruct_pexp pexp in
+  let exp = match guard with None -> exp
+    | Some exp' -> E_aux (E_block [exp';exp],(Parse_ast.Unknown,None)) in
+  fst (fold_exp
+    { (compute_exp_alg false (||) ) with
+      e_app = (fun (f, es) ->
+        let (rs, es) = List.split es in
+        (List.fold_left (||) (string_of_id f = string_of_id id) rs,
+         E_app (f, es)));
+      e_app_infix = (fun ((r1,e1), f, (r2,e2)) ->
+        (r1 || r2 || (string_of_id f = string_of_id id),
+         E_app_infix (e1, f, e2))) }
+    exp)
+
+(* Split out function clauses for individual union constructor patterns
+   (e.g. AST nodes) into auxiliary functions.  Used for the execute function. *)
+let rewrite_split_fun_constr_pats fun_name (Defs defs) =
+  let rewrite_fundef typquant (FD_aux (FD_function (r_o, t_o, e_o, clauses), ((l, _) as fdannot))) =
+    let rec_clauses, clauses = List.partition is_funcl_rec clauses in
+    let clauses, aux_funs =
+      List.fold_left
+        (fun (clauses, aux_funs) (FCL_aux (FCL_Funcl (id, pexp), fannot) as clause) ->
+           let pat, guard, exp, annot = destruct_pexp pexp in
+           match pat with
+             | P_aux (P_app (constr_id, args), pannot) ->
+                let argstup_typ = tuple_typ (List.map pat_typ_of args) in
+                let pannot' = swaptyp argstup_typ pannot in
+                let pat' = P_aux (P_tup args, pannot') in
+                let pexp' = construct_pexp (pat', guard, exp, annot) in
+                let aux_fun_id = prepend_id (fun_name ^ "_") constr_id in
+                let aux_funcl = FCL_aux (FCL_Funcl (aux_fun_id, pexp'), pannot') in
+                begin
+                  match Bindings.find_opt aux_fun_id aux_funs with
+                    | Some aux_clauses ->
+                       clauses,
+                       Bindings.add aux_fun_id (aux_clauses @ [aux_funcl]) aux_funs
+                    | None ->
+                       let argpats, argexps = List.split (List.mapi
+                         (fun idx (P_aux (paux, a)) ->
+                            let id = match paux with
+                              | P_as (_, id) | P_id id -> id
+                              | _ -> mk_id ("arg" ^ string_of_int idx)
+                            in
+                            P_aux (P_id id, a), E_aux (E_id id, a))
+                         args)
+                       in
+                       let pexp = construct_pexp
+                         (P_aux (P_app (constr_id, argpats), pannot),
+                          None,
+                          E_aux (E_app (aux_fun_id, argexps), annot),
+                          annot)
+                       in
+                       clauses @ [FCL_aux (FCL_Funcl (id, pexp), fannot)],
+                       Bindings.add aux_fun_id [aux_funcl] aux_funs
+                end
+             | _ -> clauses @ [clause], aux_funs)
+        ([], Bindings.empty) clauses
+    in
+    let add_aux_def id funcls defs =
+      let env, args_typ, ret_typ = match funcls with
+        | FCL_aux (FCL_Funcl (_, pexp), _) :: _ ->
+           let pat, _, exp, _ = destruct_pexp pexp in
+           env_of exp, pat_typ_of pat, typ_of exp
+        | _ ->
+           raise (Reporting_basic.err_unreachable l
+             "rewrite_split_fun_constr_pats: empty auxiliary function")
+      in
+      let eff = List.fold_left
+        (fun eff (FCL_aux (FCL_Funcl (_, pexp), _)) ->
+           let _, _, exp, _ = destruct_pexp pexp in
+           union_effects eff (effect_of exp))
+        no_effect funcls
+      in
+      let fun_typ = function_typ args_typ ret_typ eff in
+      let typquant = match typquant with
+        | TypQ_aux (TypQ_tq qis, l) ->
+           let qis =
+             List.filter
+               (fun qi -> KidSet.subset (tyvars_of_quant_item qi) (tyvars_of_typ fun_typ))
+               qis
+           in
+           TypQ_aux (TypQ_tq qis, l)
+        | _ -> typquant
+      in
+      let val_spec =
+        VS_aux (VS_val_spec
+          (mk_typschm typquant fun_typ, id, (fun _ -> None), false),
+          (Parse_ast.Unknown, None))
+      in
+      let fundef = FD_aux (FD_function (r_o, t_o, e_o, funcls), fdannot) in
+      [DEF_spec val_spec; DEF_fundef fundef] @ defs
+    in
+    Bindings.fold add_aux_def aux_funs
+      [DEF_fundef (FD_aux (FD_function (r_o, t_o, e_o, rec_clauses @ clauses), fdannot))]
+  in
+  let typquant = List.fold_left (fun tq def -> match def with
+    | DEF_spec (VS_aux (VS_val_spec (TypSchm_aux (TypSchm_ts (tq, _), _), id, _, _), _))
+      when string_of_id id = fun_name -> tq
+    | _ -> tq) (mk_typquant []) defs
+  in
+  let defs = List.fold_right (fun def defs -> match def with
+    | DEF_fundef fundef when string_of_id (id_of_fundef fundef) = fun_name ->
+       rewrite_fundef typquant fundef @ defs
+    | _ -> def :: defs) defs []
+  in
+  Defs defs
+
 (* Propagate effects of functions, if effect checking and propagation
    have not been performed already by the type checker. *)
 let rewrite_fix_val_specs (Defs defs) =
@@ -1841,21 +1953,6 @@ let rewrite_fix_val_specs (Defs defs) =
     let (val_specs, funcls) = List.fold_left rewrite_funcl (val_specs, []) funcls in
     (* Repeat once to cross-propagate effects between clauses *)
     let (val_specs, funcls) = List.fold_left rewrite_funcl (val_specs, []) funcls in
-    let is_funcl_rec (FCL_aux (FCL_Funcl (id, pexp), _)) =
-      let pat,guard,exp,pannot = destruct_pexp pexp in
-      let exp = match guard with None -> exp
-        | Some exp' -> E_aux (E_block [exp';exp],(Parse_ast.Unknown,None)) in
-      fst (fold_exp
-        { (compute_exp_alg false (||) ) with
-          e_app = (fun (f, es) ->
-            let (rs, es) = List.split es in
-            (List.fold_left (||) (string_of_id f = string_of_id id) rs,
-             E_app (f, es)));
-          e_app_infix = (fun ((r1,e1), f, (r2,e2)) ->
-            (r1 || r2 || (string_of_id f = string_of_id id),
-             E_app_infix (e1, f, e2))) }
-        exp)
-    in
     let recopt =
       if List.exists is_funcl_rec funcls then
         Rec_aux (Rec_rec, Parse_ast.Unknown)
@@ -2592,10 +2689,6 @@ let rewrite_defs_pat_lits =
  * internal let-expressions, or internal plet-expressions ended by a term that does not
  * access memory or registers and does not update variables *)
 
-let swaptyp typ (l,tannot) = match tannot with
-  | Some (env, typ', eff) -> (l, Some (env, typ, eff))
-  | _ -> raise (Reporting_basic.err_unreachable l "swaptyp called with empty type annotation")
-
 type 'a updated_term =
   | Added_vars of 'a exp * 'a pat
   | Same_vars of 'a exp
@@ -2984,6 +3077,7 @@ let rewrite_defs_lem = [
   ("guarded_pats", rewrite_defs_guarded_pats);
   (* ("register_ref_writes", rewrite_register_ref_writes); *)
   ("fix_val_specs", rewrite_fix_val_specs);
+  ("split_execute", rewrite_split_fun_constr_pats "execute");
   ("recheck_defs", recheck_defs);
   ("exp_lift_assign", rewrite_defs_exp_lift_assign);
   (* ("constraint", rewrite_constraint); *)