Merge pull request #31 from AeneasVerif/son_lean_backend

Improve the Lean backend

1 files changed, 100 insertions, 20 deletions

diff --git a/compiler/ExtractBase.ml b/compiler/ExtractBase.ml
index 0a5d7df2..655bb033 100644
--- a/compiler/ExtractBase.ml
+++ b/compiler/ExtractBase.ml
@@ -240,7 +240,9 @@ type formatter = {
           - loop identifier, if this is for a loop
        *)
   opaque_pre : unit -> string;
-      (** The prefix to use for opaque definitions.
+      (** TODO: obsolete, remove.
+
+          The prefix to use for opaque definitions.
 
           We need this because for some backends like Lean and Coq, we group
           opaque definitions in module signatures, meaning that using those
@@ -414,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;
@@ -425,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.
@@ -486,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
@@ -518,6 +536,8 @@ let basename_to_unique (names_set : StringSet.t)
   in
   if StringSet.mem basename names_set then gen 0 else basename
 
+type fun_name_info = { keep_fwd : bool; num_backs : int }
+
 (** Extraction context.
 
     Note that the extraction context contains information coming from the
@@ -528,6 +548,11 @@ let basename_to_unique (names_set : StringSet.t)
 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 *)
@@ -539,6 +564,25 @@ type extraction_ctx = {
           use it.
           Also see {!names_map.opaque_ids}.
        *)
+  use_dep_ite : bool;
+      (** For Lean: do we use dependent-if then else expressions?
+
+          Example:
+          {[
+            if h: b then ... else ...
+            -- ^^
+            -- makes the if then else dependent
+          ]}
+        *)
+  fun_name_info : fun_name_info PureUtils.RegularFunIdMap.t;
+      (** Information used to filter and name functions - we use it
+          to print comments in the generated code, to help link
+          the generated code to the original code (information such
+          as: "this function is the backward function of ...", or
+          "this function is the merged forward/backward function of ..."
+          in case a Rust function only has one backward translation
+          and we filter the forward function because it returns unit.
+        *)
 }
 
 (** Debugging function, used when communicating name collisions to the user,
@@ -667,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 =
@@ -918,9 +981,15 @@ let ctx_add_fun_decl (trans_group : bool * pure_fun_translation)
     ctx.fmt.fun_name def.basename def.num_loops def.loop_id num_rgs rg_info
       (keep_fwd, num_backs)
   in
-  ctx_add is_opaque
-    (FunId (FromLlbc (A.Regular def_id, def.loop_id, def.back_id)))
-    def_name ctx
+  let fun_id = (A.Regular def_id, def.loop_id, def.back_id) in
+  let ctx = ctx_add is_opaque (FunId (FromLlbc fun_id)) def_name ctx in
+  (* Add the name info *)
+  {
+    ctx with
+    fun_name_info =
+      PureUtils.RegularFunIdMap.add fun_id { keep_fwd; num_backs }
+        ctx.fun_name_info;
+  }
 
 type names_map_init = {
   keywords : string list;
@@ -1041,13 +1110,24 @@ let default_fun_suffix (num_loops : int) (loop_id : LoopId.id option)
      definitions (in particular between type and function definitions).
   *)
   let rg_suff =
+    (* TODO: make all the backends match what is done for Lean *)
     match rg with
-    | None -> "_fwd"
+    | None -> (
+        match !Config.backend with
+        | FStar | Coq | HOL4 -> "_fwd"
+        | Lean ->
+            (* In order to avoid name conflicts:
+             * - if the forward is eliminated, we add the suffix "_fwd" (it won't be used)
+             * - otherwise, no suffix (because the backward functions will have a suffix)
+             *)
+            if num_backs = 1 && not keep_fwd then "_fwd" else "")
     | Some rg ->
         assert (num_region_groups > 0 && num_backs > 0);
         if num_backs = 1 then
           (* Exactly one backward function *)
-          if not keep_fwd then "_fwd_back" else "_back"
+          match !Config.backend with
+          | FStar | Coq | HOL4 -> if not keep_fwd then "_fwd_back" else "_back"
+          | Lean -> if not keep_fwd then "" else "_back"
         else if
           (* Several region groups/backward functions:
              - if all the regions in the group have names, we use those names