Implement extration to different files

3 files changed, 128 insertions, 42 deletions

diff --git a/Makefile b/Makefile
index b12f7dae..33f98089 100644
--- a/Makefile
+++ b/Makefile
@@ -23,7 +23,7 @@ build:
 test: build translate-no_nested_borrows translate-hashmap
 
 # Add specific options to some tests
-translate-no_nested_borrows: TRANS_OPTIONS:=$(TRANS_OPTIONS) -test-units -no-decrease-clauses
+translate-no_nested_borrows: TRANS_OPTIONS:=$(TRANS_OPTIONS) -test-units -no-split -no-decrease-clauses
 translate-hashmap: TRANS_OPTIONS:=$(TRANS_OPTIONS) -template-clauses
 
 # Generic rule to extract the CFIM from a rust file
diff --git a/src/Translate.ml b/src/Translate.ml
index ec817b50..7247279a 100644
--- a/src/Translate.ml
+++ b/src/Translate.ml
@@ -15,6 +15,11 @@ let log = TranslateCore.log
 type config = {
   eval_config : Contexts.partial_config;
   mp_config : Micro.config;
+  split_files : bool;
+      (** Controls whether we split the generated definitions between different
+          files for the types, clauses and functions, or if we group them in
+          one file.
+       *)
   test_unit_functions : bool;
       (** If true, insert tests in the generated files to check that the
           unit functions normalize to `Success _`.
@@ -421,6 +426,62 @@ let extract_definitions (fmt : Format.formatter) (config : gen_config)
 
   List.iter export_decl ctx.m.declarations
 
+let extract_file (config : gen_config) (ctx : gen_ctx) (filename : string)
+    (rust_module_name : string) (module_name : string) (custom_msg : string)
+    (custom_imports : string list) (custom_includes : string list) : unit =
+  (* Open the file and create the formatter *)
+  let out = open_out filename in
+  let fmt = Format.formatter_of_out_channel out in
+
+  (* TODO: we actually don't want to use the formatter for the headers,
+   * because we want to make sure we introduce proper line breaks (it is
+   * as long as we end with a line break, so that the formatter is
+   * not lost in its count *)
+
+  (* Set the margin *)
+  Format.pp_set_margin fmt 80;
+
+  (* Create a vertical box *)
+  Format.pp_open_vbox fmt 0;
+
+  (* Create the header *)
+  Format.pp_print_string fmt
+    "(** THIS FILE WAS AUTOMATICALLY GENERATED BY AENEAS *)";
+  Format.pp_print_break fmt 0 0;
+  Format.pp_print_string fmt
+    ("(** [" ^ rust_module_name ^ "]" ^ custom_msg ^ " *)");
+  Format.pp_print_break fmt 0 0;
+  Format.pp_print_string fmt ("module " ^ module_name);
+  Format.pp_print_break fmt 0 0;
+  Format.pp_print_string fmt "open Primitives";
+  (* Add the custom imports *)
+  List.iter
+    (fun m ->
+      Format.pp_print_break fmt 0 0;
+      Format.pp_print_string fmt ("open " ^ m))
+    custom_imports;
+  (* Add the custom includes *)
+  List.iter
+    (fun m ->
+      Format.pp_print_break fmt 0 0;
+      Format.pp_print_string fmt ("include " ^ m))
+    custom_includes;
+  (* Z3 options *)
+  Format.pp_print_break fmt 0 0;
+  Format.pp_print_break fmt 0 0;
+  Format.pp_print_string fmt "#set-options \"--z3rlimit 50 --fuel 0 --ifuel 1\"";
+  Format.pp_print_break fmt 0 0;
+
+  (* Extract the definitions *)
+  extract_definitions fmt config ctx;
+
+  (* Close the box and end the formatting *)
+  Format.pp_close_box fmt ();
+  Format.pp_print_newline fmt ();
+
+  (* Some logging *)
+  log#linfo (lazy ("Generated: " ^ filename))
+
 (** Translate a module and write the synthesized code to an output file. *)
 let translate_module (filename : string) (dest_dir : string) (config : config)
     (m : M.cfim_module) : unit =
@@ -489,10 +550,6 @@ let translate_module (filename : string) (dest_dir : string) (config : config)
         (* Concatenate *)
         (module_name, Filename.concat dest_dir module_name)
   in
-  (* Add the extension for F* *)
-  let extract_filename = extract_filebasename ^ ".fst" in
-  let out = open_out extract_filename in
-  let fmt = Format.formatter_of_out_channel out in
 
   (* Put the translated definitions in maps *)
   let trans_types =
@@ -507,27 +564,7 @@ let translate_module (filename : string) (dest_dir : string) (config : config)
          trans_funs)
   in
 
-  (* Set the margin *)
-  Format.pp_set_margin fmt 80;
-
-  (* Create a vertical box *)
-  Format.pp_open_vbox fmt 0;
-
-  (* Create the header *)
-  Format.pp_print_string fmt
-    "(** THIS FILE WAS AUTOMATICALLY GENERATED BY AENEAS *)";
-  Format.pp_print_break fmt 0 0;
-  Format.pp_print_string fmt ("(** [" ^ m.M.name ^ "] *)");
-  Format.pp_print_break fmt 0 0;
-  Format.pp_print_string fmt ("module " ^ module_name);
-  Format.pp_print_break fmt 0 0;
-  Format.pp_print_string fmt "open Primitives";
-  Format.pp_print_break fmt 0 0;
-  Format.pp_print_break fmt 0 0;
-  Format.pp_print_string fmt "#set-options \"--z3rlimit 50 --fuel 0 --ifuel 1\"";
-  Format.pp_print_break fmt 0 0;
-
-  (* Extract the definitions *)
+  (* Extract the file(s) *)
   let gen_ctx =
     {
       m;
@@ -537,21 +574,64 @@ let translate_module (filename : string) (dest_dir : string) (config : config)
       functions_with_decrease_clause = rec_functions;
     }
   in
-  let gen_config =
-    {
-      extract_types = true;
-      extract_decrease_clauses = config.extract_decrease_clauses;
-      extract_template_decrease_clauses =
-        config.extract_template_decrease_clauses;
-      extract_fun_defs = true;
-      test_unit_functions = config.test_unit_functions;
-    }
-  in
-  extract_definitions fmt gen_config gen_ctx;
 
-  (* Close the box and end the formatting *)
-  Format.pp_close_box fmt ();
-  Format.pp_print_newline fmt ();
+  (* Extract one or several files, depending on the configuration *)
+  if config.split_files then (
+    let base_gen_config =
+      {
+        extract_types = false;
+        extract_decrease_clauses = config.extract_decrease_clauses;
+        extract_template_decrease_clauses = false;
+        extract_fun_defs = false;
+        test_unit_functions = false;
+      }
+    in
 
-  (* Some logging *)
-  log#linfo (lazy ("Generated: " ^ extract_filename))
+    (* Extract the types *)
+    let types_filename = extract_filebasename ^ ".Types.fst" in
+    let types_module = module_name ^ ".Types" in
+    let types_config = { base_gen_config with extract_types = true } in
+    extract_file types_config gen_ctx types_filename m.M.name types_module
+      ": types definitions" [] [];
+
+    (* Extract the template clauses *)
+    (if
+     config.extract_decrease_clauses && config.extract_template_decrease_clauses
+    then
+     let clauses_filename = extract_filebasename ^ ".Clauses.Template.fst" in
+     let clauses_module = module_name ^ ".Clauses.Template" in
+     let clauses_config =
+       { base_gen_config with extract_template_decrease_clauses = true }
+     in
+     extract_file clauses_config gen_ctx clauses_filename m.M.name
+       clauses_module ": templates for the decrease clauses" [ types_module ] []);
+
+    (* Extract the functions *)
+    let fun_filename = extract_filebasename ^ ".Funs.fst" in
+    let fun_module = module_name ^ ".Funs" in
+    let fun_config =
+      {
+        base_gen_config with
+        extract_fun_defs = true;
+        test_unit_functions = config.test_unit_functions;
+      }
+    in
+    let clauses_module = module_name ^ ".Clauses" in
+    extract_file fun_config gen_ctx fun_filename m.M.name fun_module
+      ": function definitions" []
+      [ types_module; clauses_module ])
+  else
+    let gen_config =
+      {
+        extract_types = true;
+        extract_decrease_clauses = config.extract_decrease_clauses;
+        extract_template_decrease_clauses =
+          config.extract_template_decrease_clauses;
+        extract_fun_defs = true;
+        test_unit_functions = config.test_unit_functions;
+      }
+    in
+    (* Add the extension for F* *)
+    let extract_filename = extract_filebasename ^ ".fst" in
+    extract_file gen_config gen_ctx extract_filename m.M.name module_name "" []
+      []
diff --git a/src/main.ml b/src/main.ml
index e0f16d8c..fdd0576e 100644
--- a/src/main.ml
+++ b/src/main.ml
@@ -32,6 +32,7 @@ let () =
   let test_trans_units = ref false in
   let no_decrease_clauses = ref false in
   let template_decrease_clauses = ref false in
+  let no_split = ref false in
 
   let spec =
     [
@@ -72,6 +73,10 @@ let () =
          `\n\
         \                                dedicated file. Incompatible with \
          -no-decrease-clauses" );
+      ( "-no-split",
+        Arg.Set no_split,
+        " Don't split the definitions between different files for types, \
+         functions, etc." );
     ]
   in
   (* Sanity check: -template-clauses ==> not -no-decrease-clauses *)
@@ -169,6 +174,7 @@ let () =
         {
           Translate.eval_config;
           mp_config = micro_passes_config;
+          split_files = not !no_split;
           test_unit_functions;
           extract_decrease_clauses = not !no_decrease_clauses;
           extract_template_decrease_clauses = !template_decrease_clauses;