Improve the generation of variant names for Lean

3 files changed, 33 insertions, 15 deletions

diff --git a/compiler/Extract.ml b/compiler/Extract.ml
index f2b5f00f..476359f0 100644
--- a/compiler/Extract.ml
+++ b/compiler/Extract.ml
@@ -915,8 +915,10 @@ let extract_type_decl_enum_body (ctx : extraction_ctx) (fmt : F.formatter)
 
      Note that we already printed: [type s =]
   *)
-  let print_variant variant_id v =
-    let cons_name = ctx_get_variant (AdtId def.def_id) variant_id ctx in
+  let print_variant _variant_id (v : variant) =
+    (* We don't lookup the name, because it may have a prefix for the type
+       id (in the case of Lean) *)
+    let cons_name = ctx.fmt.variant_name def.name v.variant_name in
     let fields = v.fields in
     extract_type_decl_variant ctx fmt def_name type_params cons_name fields
   in
@@ -1535,12 +1537,14 @@ let extract_adt_g_value
         (* The ADT shouldn't be opaque *)
         let with_opaque_pre = false in
         match variant_id with
-        | Some vid ->
-            if !backend = Lean then
-              ctx_get_type with_opaque_pre adt_id ctx
-              ^ "."
-              ^ ctx_get_variant adt_id vid ctx
-            else ctx_get_variant adt_id vid ctx
+        | Some vid -> (
+            (* In the case of Lean, we might have to add the type name as a prefix *)
+            match (!backend, adt_id) with
+            | Lean, Assumed _ ->
+                ctx_get_type with_opaque_pre adt_id ctx
+                ^ "."
+                ^ ctx_get_variant adt_id vid ctx
+            | _ -> ctx_get_variant adt_id vid ctx)
         | None -> ctx_get_struct with_opaque_pre adt_id ctx
       in
       if inside && field_values <> [] then F.pp_print_string fmt "(";
@@ -1731,6 +1735,7 @@ and extract_adt_cons (ctx : extraction_ctx) (fmt : F.formatter) (inside : bool)
           args;
         F.pp_print_string fmt ")")
   | _ ->
+      (* TODO: this duplicates [extrac_adt_g_value] *)
       (* "Regular" ADT *)
       (* We print something of the form: [Cons field0 ... fieldn].
        * We could update the code to print something of the form:
@@ -1741,12 +1746,14 @@ and extract_adt_cons (ctx : extraction_ctx) (fmt : F.formatter) (inside : bool)
         (* The ADT shouldn't be opaque *)
         let with_opaque_pre = false in
         match adt_cons.variant_id with
-        | Some vid ->
-            if !backend = Lean then
-              ctx_get_type with_opaque_pre adt_cons.adt_id ctx
-              ^ "."
-              ^ ctx_get_variant adt_cons.adt_id vid ctx
-            else ctx_get_variant adt_cons.adt_id vid ctx
+        | Some vid -> (
+            (* In the case of Lean, we might have to add the type name as a prefix *)
+            match (!backend, adt_cons.adt_id) with
+            | Lean, Assumed _ ->
+                ctx_get_type with_opaque_pre adt_cons.adt_id ctx
+                ^ "."
+                ^ ctx_get_variant adt_cons.adt_id vid ctx
+            | _ -> ctx_get_variant adt_cons.adt_id vid ctx)
         | None -> ctx_get_struct with_opaque_pre adt_cons.adt_id ctx
       in
       let use_parentheses = inside && args <> [] in
diff --git a/compiler/ExtractBase.ml b/compiler/ExtractBase.ml
index 86bb0cff..81bdd202 100644
--- a/compiler/ExtractBase.ml
+++ b/compiler/ExtractBase.ml
@@ -805,6 +805,13 @@ let ctx_add_variant (def : type_decl) (variant_id : VariantId.id)
     (variant : variant) (ctx : extraction_ctx) : extraction_ctx * string =
   let is_opaque = false in
   let name = ctx.fmt.variant_name def.name variant.variant_name in
+  (* Add the type name prefix for Lean *)
+  let name =
+    if !Config.backend = Lean then
+      let type_name = ctx.fmt.type_name def.name in
+      type_name ^ "." ^ name
+    else name
+  in
   let ctx =
     ctx_add is_opaque (VariantId (AdtId def.def_id, variant_id)) name ctx
   in
diff --git a/compiler/Translate.ml b/compiler/Translate.ml
index 347052a8..781766ee 100644
--- a/compiler/Translate.ml
+++ b/compiler/Translate.ml
@@ -796,7 +796,11 @@ let translate_crate (filename : string) (dest_dir : string) (crate : A.crate) :
   (* Initialize the extraction context - for now we extract only to F*.
    * We initialize the names map by registering the keywords used in the
    * language, as well as some primitive names ("u32", etc.) *)
-  let variant_concatenate_type_name = true in
+  let variant_concatenate_type_name =
+    (* For Lean, we exploit the fact that the variant name should always be
+       prefixed with the type name to prevent collisions *)
+    match !Config.backend with Coq | FStar -> true | Lean -> false
+  in
   let mk_formatter_and_names_map = Extract.mk_formatter_and_names_map in
   let fmt, names_map =
     mk_formatter_and_names_map trans_ctx crate.name