Start working on fixing the extraction of type definitions

2 files changed, 43 insertions, 9 deletions

diff --git a/src/ExtractToFStar.ml b/src/ExtractToFStar.ml
index 506b641f..63d1affd 100644
--- a/src/ExtractToFStar.ml
+++ b/src/ExtractToFStar.ml
@@ -23,6 +23,7 @@ let fstar_keywords =
     "match";
     "with";
     "assert";
+    "Type0";
   ]
 
 (**
@@ -233,6 +234,8 @@ let extract_type_def_struct_body (ctx : extraction_ctx) (fmt : F.formatter)
     F.pp_print_string fmt ":";
     F.pp_print_space fmt ();
     extract_ty ctx fmt false f.field_ty;
+    F.pp_print_string fmt ";";
+    F.pp_print_space fmt ();
     F.pp_close_box fmt ()
   in
   let fields = FieldId.mapi (fun fid f -> (fid, f)) fields in
@@ -389,10 +392,29 @@ let extract_type_def (ctx : extraction_ctx) (fmt : F.formatter) (def : type_def)
   let ctx_body, type_params = ctx_add_type_params def.type_params ctx in
   (* Open a box for the definition *)
   F.pp_open_hovbox fmt 0;
-  (* > "type TYPE_NAME =" *)
-  F.pp_print_string fmt "type";
+  (* Print a comment to link the extracted type to its original rust definition *)
+  F.pp_print_string fmt ("(** [" ^ Print.name_to_string def.name ^ "] *)");
   F.pp_print_space fmt ();
-  F.pp_print_string fmt def_name;
+  (* Open a box for "type TYPE_NAME (TYPE_PARAMS) =" *)
+  F.pp_open_hovbox fmt ctx.indent_incr;
+  (* > "type TYPE_NAME" *)
+  F.pp_print_string fmt ("type " ^ def_name);
+  (* Print the type parameters *)
+  if def.type_params <> [] then (
+    F.pp_print_space fmt ();
+    F.pp_print_string fmt "(";
+    List.iter
+      (fun (p : type_var) ->
+        let pname = ctx_get_type_var p.index ctx_body in
+        F.pp_print_string fmt pname;
+        F.pp_print_space fmt ())
+      def.type_params;
+    F.pp_print_string fmt " : Type0)");
+  (* Print the "=" *)
+  F.pp_print_space fmt ();
+  F.pp_print_string fmt "=";
+  (* Close the box for "type TYPE_NAME (TYPE_PARAMS) =" *)
+  F.pp_close_box fmt ();
   (match def.kind with
   | Struct fields -> extract_type_def_struct_body ctx_body fmt def fields
   | Enum variants ->
diff --git a/src/Translate.ml b/src/Translate.ml
index ed9cd22c..e43d5741 100644
--- a/src/Translate.ml
+++ b/src/Translate.ml
@@ -291,17 +291,23 @@ let translate_module (filename : string) (config : C.partial_config)
   (* TODO: register the functions *)
 
   (* Open the output file *)
-  (* First compute the filename by replacing the extension *)
-  let filename =
+  (* First compute the filename by replacing the extension and converting the
+   * case (rust module names are snake case) *)
+  let module_name, extract_filename =
     match Filename.chop_suffix_opt ~suffix:".cfim" filename with
     | None ->
         (* Note that we already checked the suffix upon opening the file *)
         failwith "Unreachable"
     | Some filename ->
+        (* Split between basename and dirname *)
+        let dirname = Filename.dirname filename in
+        let basename = Filename.basename filename in
+        (* Convert the case *)
+        let module_name = StringUtils.to_camel_case basename in
         (* We add the extension for F* *)
-        filename ^ ".fst"
+        (module_name, Filename.concat dirname (module_name ^ ".fst"))
   in
-  let out = open_out filename in
+  let out = open_out extract_filename in
   let fmt = Format.formatter_of_out_channel out in
 
   (* Put the translated definitions in maps *)
@@ -319,6 +325,13 @@ let translate_module (filename : string) (config : C.partial_config)
   (* Create a vertical box *)
   Format.pp_open_vbox fmt 0;
 
+  (* Add the module name *)
+  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_break fmt 0 0;
+
   (* Export the definition groups to the file, in the proper order *)
   let export_type (id : Pure.TypeDefId.id) : unit =
     let def = Pure.TypeDefId.Map.find id trans_types in
@@ -338,5 +351,4 @@ let translate_module (filename : string) (config : C.partial_config)
   List.iter export_decl m.declarations;
 
   (* Close the box and end the formatting *)
-  Format.pp_close_box fmt ();
-  Format.pp_print_newline fmt ()
+  Format.pp_close_box fmt ()