Add an option to crudely slice a function out of a Sail model

Not ideal because it keeps everything that's not a function, but good
enough for quick tests extracted from arm.

3 files changed, 26 insertions, 2 deletions

diff --git a/src/sail.ml b/src/sail.ml
index be2a6198..41aee119 100644
--- a/src/sail.ml
+++ b/src/sail.ml
@@ -72,6 +72,7 @@ let opt_libs_coq = ref ([]:string list)
 let opt_file_arguments = ref ([]:string list)
 let opt_process_elf : string option ref = ref None
 let opt_ocaml_generators = ref ([]:string list)
+let opt_slice = ref ([]:string list)
 
 let options = Arg.align ([
   ( "-o",
@@ -268,6 +269,9 @@ let options = Arg.align ([
   ( "-dprofile",
     Arg.Set Profile.opt_profile,
     " (debug) provides basic profiling information for rewriting passes within Sail");
+  ( "-slice",
+    Arg.String (fun s -> opt_slice := s::!opt_slice),
+    "<id> produce version of input restricted to the given function");
   ( "-v",
     Arg.Set opt_print_version,
     " print version");
@@ -380,6 +384,12 @@ let main() =
       (if !(opt_print_verbose)
        then ((Pretty_print_sail.pp_defs stdout) ast)
        else ());
+      (match !opt_slice with
+       | [] -> ()
+       | ids ->
+          let ids = List.map Ast_util.mk_id ids in
+          let ids = Ast_util.IdSet.of_list ids in
+          Pretty_print_sail.pp_defs stdout (Specialize.slice_defs type_envs ast ids));
       (if !(opt_print_ocaml)
        then
          let ast_ocaml = rewrite_ast_ocaml ast in
diff --git a/src/specialize.ml b/src/specialize.ml
index e7f686d8..00357557 100644
--- a/src/specialize.ml
+++ b/src/specialize.ml
@@ -379,8 +379,11 @@ let specialize_id_overloads instantiations id (Defs defs) =
    therefore remove all unused valspecs. Remaining polymorphic
    valspecs are then re-specialized. This process is iterated until
    the whole spec is specialized. *)
-let remove_unused_valspecs env ast =
-  let calls = ref (IdSet.of_list [mk_id "main"; mk_id "__SetConfig"; mk_id "__ListConfig"; mk_id "execute"; mk_id "decode"; mk_id "initialize_registers"; mk_id "append_64"]) in
+
+let initial_calls = (IdSet.of_list [mk_id "main"; mk_id "__SetConfig"; mk_id "__ListConfig"; mk_id "execute"; mk_id "decode"; mk_id "initialize_registers"; mk_id "append_64"])
+
+let remove_unused_valspecs ?(initial_calls=initial_calls) env ast =
+  let calls = ref initial_calls in
   let vs_ids = Initial_check.val_spec_ids ast in
 
   let inspect_exp = function
@@ -413,6 +416,14 @@ let remove_unused_valspecs env ast =
 
   List.fold_left (fun ast id -> Defs (remove_unused ast id)) ast (IdSet.elements unused)
 
+let slice_defs env (Defs defs) keep_ids =
+  let keep = function
+    | DEF_fundef fd -> IdSet.mem (id_of_fundef fd) keep_ids
+    | _ -> true
+  in
+  let defs = List.filter keep defs in
+  remove_unused_valspecs env (Defs defs) ~initial_calls:keep_ids
+
 let specialize_id id ast =
   let instantiations = instantiations_of id ast in
   let ast = specialize_id_valspec instantiations id ast in
diff --git a/src/specialize.mli b/src/specialize.mli
index f2c94a48..28029747 100644
--- a/src/specialize.mli
+++ b/src/specialize.mli
@@ -71,3 +71,6 @@ val specialize : tannot defs -> Env.t -> tannot defs * Env.t
 val instantiations_of : id -> tannot defs -> typ_arg KBindings.t list
 
 val string_of_instantiation : typ_arg KBindings.t -> string
+
+(* Remove all function definitions except for the given set *)
+val slice_defs : Env.t -> tannot defs -> IdSet.t -> tannot defs