Update the extraction of Lean files

2 files changed, 51 insertions, 25 deletions

diff --git a/compiler/Extract.ml b/compiler/Extract.ml
index 63d41d9a..0b8d8bdf 100644
--- a/compiler/Extract.ml
+++ b/compiler/Extract.ml
@@ -973,8 +973,9 @@ let extract_type_decl_struct_body (ctx : extraction_ctx) (fmt : F.formatter)
     if !backend = FStar && fields = [] then (
       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 (
-      F.pp_print_space fmt ();
+      if !backend <> Lean then F.pp_print_space fmt ();
       (* If Coq: print the constructor name *)
       (* TODO: remove superfluous test not is_rec below *)
       if !backend = Coq && not is_rec then (
@@ -985,10 +986,14 @@ let extract_type_decl_struct_body (ctx : extraction_ctx) (fmt : F.formatter)
       if !backend <> Lean then F.pp_print_string fmt "{";
       F.pp_print_break fmt 1 ctx.indent_incr;
       (* The body itself *)
-      F.pp_open_hvbox fmt 0;
+      (* Open a box for the body *)
+      (match !backend with
+      | Coq | FStar -> F.pp_open_hvbox fmt 0
+      | Lean -> F.pp_open_vbox fmt 0);
       (* Print the fields *)
       let print_field (field_id : FieldId.id) (f : field) : unit =
         let field_name = ctx_get_field (AdtId def.def_id) field_id ctx in
+        (* Open a box for the field *)
         F.pp_open_box fmt ctx.indent_incr;
         F.pp_print_string fmt field_name;
         F.pp_print_space fmt ();
@@ -996,20 +1001,21 @@ let extract_type_decl_struct_body (ctx : extraction_ctx) (fmt : F.formatter)
         F.pp_print_space fmt ();
         extract_ty ctx fmt false f.field_ty;
         if !backend <> Lean then F.pp_print_string fmt ";";
+        (* Close the box for the field *)
         F.pp_close_box fmt ()
       in
       let fields = FieldId.mapi (fun fid f -> (fid, f)) fields in
       Collections.List.iter_link (F.pp_print_space fmt)
         (fun (fid, f) -> print_field fid f)
         fields;
-      (* Close *)
+      (* Close the box for the body *)
       F.pp_close_box fmt ();
-      F.pp_print_space fmt ();
-      if !backend <> Lean then F.pp_print_string fmt "}")
+      if !backend <> Lean then (
+        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
-         a group of mutually recursive types: we extract it as an inductive type
-      *)
+         a group of mutually recursive types: we extract it as an inductive type *)
       assert (is_rec && !backend = Coq);
       let with_opaque_pre = false in
       let cons_name = ctx_get_struct with_opaque_pre (AdtId def.def_id) ctx in
@@ -1068,10 +1074,13 @@ let extract_type_decl (ctx : extraction_ctx) (fmt : F.formatter)
   F.pp_print_break fmt 0 0;
   (* Print a comment to link the extracted type to its original rust definition *)
   extract_comment fmt ("[" ^ Print.name_to_string def.name ^ "]");
-  F.pp_print_space fmt ();
+  F.pp_print_break fmt 0 0;
   (* Open a box for the definition, so that whenever possible it gets printed on
-   * one line *)
-  F.pp_open_hvbox fmt 0;
+   * one line. Note however that in the case of Lean line breaks are important
+   * for parsing: we thus use a hovbox. *)
+  (match !backend with
+  | Coq | FStar -> F.pp_open_hvbox fmt 0
+  | Lean -> F.pp_open_vbox fmt 0);
   (* Open a box for "type TYPE_NAME (TYPE_PARAMS) =" *)
   F.pp_open_hovbox fmt ctx.indent_incr;
   (* > "type TYPE_NAME" *)
@@ -1111,7 +1120,7 @@ let extract_type_decl (ctx : extraction_ctx) (fmt : F.formatter)
     in
     F.pp_print_string fmt eq)
   else (
-    (* Otherwise print ": Type0" *)
+    (* Otherwise print ": Type" *)
     if use_forall then F.pp_print_string fmt ","
     else (
       F.pp_print_space fmt ();
@@ -1384,7 +1393,7 @@ let extract_state_type (fmt : F.formatter) (ctx : extraction_ctx)
   F.pp_print_break fmt 0 0;
   (* Print a comment  *)
   extract_comment fmt "The state type used in the state-error monad";
-  F.pp_print_space fmt ();
+  F.pp_print_break fmt 0 0;
   (* Open a box for the definition, so that whenever possible it gets printed on
    * one line *)
   F.pp_open_hvbox fmt 0;
@@ -1392,8 +1401,13 @@ let extract_state_type (fmt : F.formatter) (ctx : extraction_ctx)
   let state_name = ctx_get_assumed_type State ctx in
   (* The syntax for Lean and Coq is almost identical. *)
   let print_axiom () =
-    if !backend = Coq then F.pp_print_string fmt "Axiom"
-    else F.pp_print_string fmt "axiom";
+    let axiom =
+      match !backend with
+      | Coq -> "Axiom"
+      | Lean -> "axiom"
+      | FStar -> raise (Failure "Unexpected")
+    in
+    F.pp_print_string fmt axiom;
     F.pp_print_space fmt ();
     F.pp_print_string fmt state_name;
     F.pp_print_space fmt ();
diff --git a/compiler/Translate.ml b/compiler/Translate.ml
index ab1addab..3278aa6a 100644
--- a/compiler/Translate.ml
+++ b/compiler/Translate.ml
@@ -679,18 +679,20 @@ let extract_definitions (fmt : Format.formatter) (config : gen_config)
      of the functions marked as Opaque. We emit the `structure ...` bit here,
      then rely on `extract_fun_decl` to be aware of this, and skip the keyword
      (e.g. "axiom" or "val") so as to generate valid syntax for records. *)
-  if config.extract_opaque && config.extract_fun_decls && !Config.backend = Lean
-  then (
+  let wrap_in_sig =
+    config.extract_opaque && config.extract_fun_decls && !Config.backend = Lean
+  in
+  if wrap_in_sig then (
     Format.pp_print_break fmt 0 0;
     Format.pp_open_vbox fmt ctx.extract_ctx.indent_incr;
     Format.pp_print_string fmt "structure OpaqueDefs where";
     Format.pp_print_break fmt 0 0);
   List.iter export_decl_group ctx.crate.declarations;
-  if config.extract_opaque && !Config.backend = Lean then
-    Format.pp_close_box fmt ();
 
   if config.extract_state_type && not config.extract_fun_decls then
-    export_state_type ()
+    export_state_type ();
+
+  if wrap_in_sig then Format.pp_close_box fmt ()
 
 let extract_file (config : gen_config) (ctx : gen_ctx) (filename : string)
     (rust_module_name : string) (module_name : string) (custom_msg : string)
@@ -1117,6 +1119,19 @@ let translate_crate (filename : string) (dest_dir : string) (crate : A.crate) :
       *)
   | Lean ->
       (*
+       * Generate the library entry point, if the crate is split between
+       * different files.
+       *)
+      if !Config.split_files then (
+        let filename = Filename.concat dest_dir (module_name ^ ".lean") in
+        let out = open_out filename in
+        (* Write *)
+        Printf.fprintf out "import %s.Funs\n" module_name;
+        (* Flush and close the file, log *)
+        close_out out;
+        log#linfo (lazy ("Generated: " ^ filename)));
+
+      (*
        * Generate the lakefile.lean file
        *)
 
@@ -1131,14 +1146,11 @@ let translate_crate (filename : string) (dest_dir : string) (crate : A.crate) :
         "  \"https://github.com/leanprover-community/mathlib4.git\"\n\n";
 
       let package_name = StringUtils.to_snake_case module_name in
-      Printf.fprintf out "package «%s» {\n" package_name;
-      Printf.fprintf out "  -- add package configuration options here\n}\n\n";
+      Printf.fprintf out "package «%s» {}\n\n" package_name;
 
-      Printf.fprintf out "lean_lib «Base» {\n";
-      Printf.fprintf out "  -- add library configuration options here\n}\n\n";
+      Printf.fprintf out "lean_lib «Base» {}\n\n";
 
-      Printf.fprintf out "lean_lib «%s» {\n" module_name;
-      Printf.fprintf out "  -- add library configuration options here\n}\n\n";
+      Printf.fprintf out "@[default_target]\nlean_lib «%s» {}\n" module_name;
 
       (* No default target for now.
          Format would be: