Make more updates for the Lean backend

2 files changed, 43 insertions, 26 deletions

diff --git a/compiler/Extract.ml b/compiler/Extract.ml
index e5710394..03c256ec 100644
--- a/compiler/Extract.ml
+++ b/compiler/Extract.ml
@@ -197,14 +197,25 @@ let assumed_adts () : (assumed_ty * string) list =
             (Char.uppercase_ascii s.[0])
             (String.sub s 1 (String.length s - 1)) )
       else (t, s))
-    [
-      (State, "state");
-      (Result, "result");
-      (Error, "error");
-      (Fuel, "nat");
-      (Option, "option");
-      (Vec, "vec");
-    ]
+    (match !backend with
+    | Lean ->
+        [
+          (State, "State");
+          (Result, "Result");
+          (Error, "Error");
+          (Fuel, "Nat");
+          (Option, "Option");
+          (Vec, "Vec");
+        ]
+    | Coq | FStar ->
+        [
+          (State, "state");
+          (Result, "result");
+          (Error, "error");
+          (Fuel, "nat");
+          (Option, "option");
+          (Vec, "vec");
+        ])
 
 let assumed_structs : (assumed_ty * string) list = []
 
@@ -1004,7 +1015,9 @@ let extract_type_decl_struct_body (ctx : extraction_ctx) (fmt : F.formatter)
       F.pp_print_space fmt ();
       F.pp_print_string fmt (unit_name ()))
     else if !backend = Lean && fields = [] then ()
-    else if (not is_rec) || !backend <> Coq then (
+      (* If the definition is recursive, we may need to extract it as an inductive
+         (instead of a record) *)
+    else if (not is_rec) || (!backend <> Coq && !backend <> Lean) then (
       if !backend <> Lean then F.pp_print_space fmt ();
       (* If Coq: print the constructor name *)
       (* TODO: remove superfluous test not is_rec below *)
@@ -1044,9 +1057,9 @@ let extract_type_decl_struct_body (ctx : extraction_ctx) (fmt : F.formatter)
         F.pp_print_space fmt ();
         F.pp_print_string fmt "}"))
     else (
-      (* We extract for Coq, and we have a recursive record, or a record in
+      (* We extract for Coq or Lean, and we have a recursive record, or a record in
          a group of mutually recursive types: we extract it as an inductive type *)
-      assert (is_rec && !backend = Coq);
+      assert (is_rec && (!backend = Coq || !backend = Lean));
       let with_opaque_pre = false in
       let cons_name = ctx_get_struct with_opaque_pre (AdtId def.def_id) ctx in
       let def_name = ctx_get_local_type with_opaque_pre def.def_id ctx in
@@ -2224,7 +2237,7 @@ let extract_template_fstar_decreases_clause (ctx : extraction_ctx)
   (* Add breaks to insert new lines between definitions *)
   F.pp_print_break fmt 0 0
 
-(** Extract templates for the [termination_by] and [decreases_by] clauses of a
+(** Extract templates for the [termination_by] and [decreasing_by] clauses of a
     recursive function definition.
 
     For Lean only.
@@ -2490,8 +2503,6 @@ let extract_fun_decl (ctx : extraction_ctx) (fmt : F.formatter)
     F.pp_close_box fmt ());
   (* Close the inner box for the definition *)
   F.pp_close_box fmt ();
-  (* Coq: add a "." *)
-  print_decl_end_delimiter fmt kind;
   (* Termination clause and proof for Lean *)
   if has_decreases_clause && !backend = Lean then (
     let def_body = Option.get def.body in
@@ -2558,6 +2569,8 @@ let extract_fun_decl (ctx : extraction_ctx) (fmt : F.formatter)
     F.pp_close_box fmt ();
     (* Close the box for the [decreasing by ...] *)
     F.pp_close_box fmt ());
+  (* Add the definition group end delimiter *)
+  print_decl_end_delimiter fmt kind;
   (* Close the outer box for the definition *)
   F.pp_close_box fmt ();
   (* Add breaks to insert new lines between definitions *)
diff --git a/compiler/Translate.ml b/compiler/Translate.ml
index 8250f813..409fc5d3 100644
--- a/compiler/Translate.ml
+++ b/compiler/Translate.ml
@@ -716,8 +716,7 @@ let extract_definitions (fmt : Format.formatter) (config : gen_config)
 
 let extract_file (config : gen_config) (ctx : gen_ctx) (filename : string)
     (rust_module_name : string) (module_name : string) (custom_msg : string)
-    ?(custom_variables : string list = []) (custom_imports : string list)
-    (custom_includes : string list) : unit =
+    (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
@@ -770,10 +769,7 @@ let extract_file (config : gen_config) (ctx : gen_ctx) (filename : string)
       (* Add the custom imports *)
       List.iter (fun m -> Printf.fprintf out "import %s\n" m) custom_imports;
       (* Add the custom includes *)
-      List.iter (fun m -> Printf.fprintf out "import %s\n" m) custom_includes;
-      if custom_variables <> [] then (
-        Printf.fprintf out "\n";
-        List.iter (fun m -> Printf.fprintf out "%s\n" m) custom_variables));
+      List.iter (fun m -> Printf.fprintf out "import %s\n" m) custom_includes);
   (* From now onwards, we use the formatter *)
   (* Set the margin *)
   Format.pp_set_margin fmt 80;
@@ -1066,6 +1062,18 @@ let translate_crate (filename : string) (dest_dir : string) (crate : A.crate) :
          extract_filebasename ^ file_delimiter ^ "Opaque" ^ opaque_ext
        in
        let opaque_module = module_name ^ module_delimiter ^ "Opaque" in
+       let opaque_imported_module =
+         (* In the case of Lean, we declare an interface (a record) containing
+            the opaque definitions, and we leave it to the user to provide an
+            instance of this module.
+
+            TODO: do the same for Coq.
+            TODO: do the same for the type definitions.
+         *)
+         if !Config.backend = Lean then
+           module_name ^ module_delimiter ^ "ExternalFuns"
+         else opaque_module
+       in
        let opaque_config =
          {
            base_gen_config with
@@ -1083,7 +1091,7 @@ let translate_crate (filename : string) (dest_dir : string) (crate : A.crate) :
        in
        extract_file opaque_config gen_ctx opaque_filename crate.A.name
          opaque_module ": opaque function definitions" [] [ types_module ];
-       [ opaque_module ])
+       [ opaque_imported_module ])
      else []
    in
 
@@ -1106,12 +1114,8 @@ let translate_crate (filename : string) (dest_dir : string) (crate : A.crate) :
        [ module_name ^ clauses_submodule ^ module_delimiter ^ "Clauses" ]
      else []
    in
-   let custom_variables =
-     if has_opaque_funs then [ "section variable (opaque_defs: OpaqueDefs)" ]
-     else []
-   in
    extract_file fun_config gen_ctx fun_filename crate.A.name fun_module
-     ": function definitions" [] ~custom_variables
+     ": function definitions" []
      ([ types_module ] @ opaque_funs_module @ clauses_module))
   else
     let gen_config =