Use short names for the structure fields in Lean

5 files changed, 78 insertions, 19 deletions

diff --git a/compiler/Config.ml b/compiler/Config.ml
index f58ba89b..0475899c 100644
--- a/compiler/Config.ml
+++ b/compiler/Config.ml
@@ -304,3 +304,14 @@ let filter_useless_functions = ref true
     called opaque_defs, of type OpaqueDefs.
   *)
 let wrap_opaque_in_sig = ref false
+
+(** Use short names for the record fields.
+
+    Some backends can't disambiguate records when their field names have collisions.
+    When this happens, we use long names, by which we concatenate the record
+    names with the field names, and check whether there are name collisions.
+
+    For backends which can disambiguate records (typically by using the typing
+    information), we use short names (i.e., the original field names).
+ *)
+let record_fields_short_names = ref false
diff --git a/compiler/Driver.ml b/compiler/Driver.ml
index f935a717..166ef11b 100644
--- a/compiler/Driver.ml
+++ b/compiler/Driver.ml
@@ -163,7 +163,9 @@ let () =
         (* We don't support fuel for the Lean backend *)
         if !use_fuel then (
           log#error "The Lean backend doesn't support the -use-fuel option";
-          fail ())
+          fail ());
+        (* Lean can disambiguate the field names *)
+        record_fields_short_names := true
     | HOL4 ->
         (* We don't support fuel for the HOL4 backend *)
         if !use_fuel then (
diff --git a/compiler/Extract.ml b/compiler/Extract.ml
index cacb9b96..558a981d 100644
--- a/compiler/Extract.ml
+++ b/compiler/Extract.ml
@@ -625,10 +625,15 @@ let mk_formatter (ctx : trans_ctx) (crate_name : string)
   in
   let field_name (def_name : name) (field_id : FieldId.id)
       (field_name : string option) : string =
-    let def_name = type_name_to_snake_case def_name ^ "_" in
-    match field_name with
-    | Some field_name -> def_name ^ field_name
-    | None -> def_name ^ FieldId.to_string field_id
+    let field_name =
+      match field_name with
+      | Some field_name -> field_name
+      | None -> FieldId.to_string field_id
+    in
+    if !Config.record_fields_short_names then field_name
+    else
+      let def_name = type_name_to_snake_case def_name ^ "_" in
+      def_name ^ field_name
   in
   let variant_name (def_name : name) (variant : string) : string =
     match !backend with
diff --git a/compiler/ExtractBase.ml b/compiler/ExtractBase.ml
index ede7af29..655bb033 100644
--- a/compiler/ExtractBase.ml
+++ b/compiler/ExtractBase.ml
@@ -416,7 +416,7 @@ module IdSet = Collections.MakeSet (IdOrderedType)
 
     We use it for lookups (during the translation) and to check for name clashes.
 
-    [id_to_string] is for debugging.
+    [id_to_name] is for debugging.
   *)
 type names_map = {
   id_to_name : string IdMap.t;
@@ -427,7 +427,9 @@ type names_map = {
        *)
   names_set : StringSet.t;
   opaque_ids : IdSet.t;
-      (** The set of opaque definitions.
+      (** TODO: this is obsolete. Remove.
+
+          The set of opaque definitions.
 
           See {!formatter.opaque_pre} for detailed explanations about why
           we need to know which definitions are opaque to compute names.
@@ -488,6 +490,20 @@ let names_map_add_function (id_to_string : id -> string) (is_opaque : bool)
     (fid : fun_id) (name : string) (nm : names_map) : names_map =
   names_map_add id_to_string is_opaque (FunId fid) name nm
 
+(** The unsafe names map stores mappings from identifiers to names which might
+    collide. For some backends and some names, it might be acceptable to have
+    collisions. For instance, in Lean, different records can have fields with
+    the same name because Lean uses the typing information to resolve the
+    ambiguities.
+
+    This map complements the {!names_map}, which checks for collisions.
+  *)
+type unsafe_names_map = { id_to_name : string IdMap.t }
+
+let unsafe_names_map_add (id : id) (name : string) (nm : unsafe_names_map) :
+    unsafe_names_map =
+  { id_to_name = IdMap.add id name nm.id_to_name }
+
 (** Make a (variable) basename unique (by adding an index).
 
     We do this in an inefficient manner (by testing all indices starting from
@@ -532,6 +548,11 @@ type fun_name_info = { keep_fwd : bool; num_backs : int }
 type extraction_ctx = {
   trans_ctx : trans_ctx;
   names_map : names_map;
+      (** The map for id to names, where we forbid name collisions
+          (ex.: we always forbid function name collisions). *)
+  unsafe_names_map : unsafe_names_map;
+      (** The map for id to names, where we allow name collisions
+          (ex.: we might allow record field name collisions). *)
   fmt : formatter;
   indent_incr : int;
       (** The indent increment we insert whenever we need to indent more *)
@@ -690,23 +711,42 @@ let id_to_string (id : id) (ctx : extraction_ctx) : string =
   | TypeVarId id -> "type_var_id: " ^ TypeVarId.to_string id
   | VarId id -> "var_id: " ^ VarId.to_string id
 
+(** We might not check for collisions for some specific ids (ex.: field names) *)
+let allow_collisions (id : id) : bool =
+  match id with
+  | FieldId (_, _) -> !Config.record_fields_short_names
+  | _ -> false
+
 let ctx_add (is_opaque : bool) (id : id) (name : string) (ctx : extraction_ctx)
     : extraction_ctx =
-  (* The id_to_string function to print nice debugging messages if there are
-   * collisions *)
-  let id_to_string (id : id) : string = id_to_string id ctx in
-  let names_map = names_map_add id_to_string is_opaque id name ctx.names_map in
-  { ctx with names_map }
+  (* We do not use the same name map if we allow/disallow collisions *)
+  if allow_collisions id then (
+    assert (not is_opaque);
+    {
+      ctx with
+      unsafe_names_map = unsafe_names_map_add id name ctx.unsafe_names_map;
+    })
+  else
+    (* The id_to_string function to print nice debugging messages if there are
+     * collisions *)
+    let id_to_string (id : id) : string = id_to_string id ctx in
+    let names_map =
+      names_map_add id_to_string is_opaque id name ctx.names_map
+    in
+    { ctx with names_map }
 
 (** [with_opaque_pre]: if [true] and the definition is opaque, add the opaque prefix *)
 let ctx_get (with_opaque_pre : bool) (id : id) (ctx : extraction_ctx) : string =
-  match IdMap.find_opt id ctx.names_map.id_to_name with
-  | Some s ->
-      let is_opaque = IdSet.mem id ctx.names_map.opaque_ids in
-      if with_opaque_pre && is_opaque then ctx.fmt.opaque_pre () ^ s else s
-  | None ->
-      log#serror ("Could not find: " ^ id_to_string id ctx);
-      raise Not_found
+  (* We do not use the same name map if we allow/disallow collisions *)
+  if allow_collisions id then IdMap.find id ctx.unsafe_names_map.id_to_name
+  else
+    match IdMap.find_opt id ctx.names_map.id_to_name with
+    | Some s ->
+        let is_opaque = IdSet.mem id ctx.names_map.opaque_ids in
+        if with_opaque_pre && is_opaque then ctx.fmt.opaque_pre () ^ s else s
+    | None ->
+        log#serror ("Could not find: " ^ id_to_string id ctx);
+        raise Not_found
 
 let ctx_get_global (with_opaque_pre : bool) (id : A.GlobalDeclId.id)
     (ctx : extraction_ctx) : string =
diff --git a/compiler/Translate.ml b/compiler/Translate.ml
index 444642c0..c5f7df92 100644
--- a/compiler/Translate.ml
+++ b/compiler/Translate.ml
@@ -922,6 +922,7 @@ let translate_crate (filename : string) (dest_dir : string) (crate : A.crate) :
     {
       ExtractBase.trans_ctx;
       names_map;
+      unsafe_names_map = { id_to_name = ExtractBase.IdMap.empty };
       fmt;
       indent_incr = 2;
       use_opaque_pre = !Config.split_files;