Register the names for the trait decls

2 files changed, 154 insertions, 14 deletions

diff --git a/compiler/Extract.ml b/compiler/Extract.ml
index 17f850a3..5eb30daa 100644
--- a/compiler/Extract.ml
+++ b/compiler/Extract.ml
@@ -720,6 +720,41 @@ let mk_formatter (ctx : trans_ctx) (crate_name : string)
     fname ^ suffix
   in
 
+  let trait_decl_name (trait_decl : trait_decl) : string =
+    type_name_to_snake_case trait_decl.name
+  in
+
+  let trait_impl_name (trait_impl : trait_impl) : string =
+    get_fun_name trait_impl.name
+  in
+
+  let trait_parent_clause_name (trait_decl : trait_decl) (clause : trait_clause)
+      : string =
+    (* TODO: improve - it would be better to not use indices *)
+    let clause = "parent_clause_" ^ TraitClauseId.to_string clause.clause_id in
+    if !Config.record_fields_short_names then clause
+    else trait_decl_name trait_decl ^ "_" ^ clause
+  in
+  let trait_type_name (trait_decl : trait_decl) (item : string) : string =
+    if !Config.record_fields_short_names then item
+    else trait_decl_name trait_decl ^ "_" ^ item
+  in
+  let trait_const_name (trait_decl : trait_decl) (item : string) : string =
+    if !Config.record_fields_short_names then item
+    else trait_decl_name trait_decl ^ "_" ^ item
+  in
+  let trait_method_name (trait_decl : trait_decl) (item : string) : string =
+    if !Config.record_fields_short_names then item
+    else trait_decl_name trait_decl ^ "_" ^ item
+  in
+  let trait_type_clause_name (trait_decl : trait_decl) (item : string)
+      (clause : trait_clause) : string =
+    (* TODO: improve - it would be better to not use indices *)
+    trait_type_name trait_decl item
+    ^ "_clause_"
+    ^ TraitClauseId.to_string clause.clause_id
+  in
+
   let termination_measure_name (_fid : A.FunDeclId.id) (fname : fun_name)
       (num_loops : int) (loop_id : LoopId.id option) : string =
     let fname = get_fun_name fname in
@@ -933,6 +968,13 @@ let mk_formatter (ctx : trans_ctx) (crate_name : string)
     fun_name;
     termination_measure_name;
     decreases_proof_name;
+    trait_decl_name;
+    trait_impl_name;
+    trait_parent_clause_name;
+    trait_const_name;
+    trait_type_name;
+    trait_method_name;
+    trait_type_clause_name;
     opaque_pre;
     var_basename;
     type_var_basename;
@@ -1233,8 +1275,8 @@ let rec extract_ty (ctx : extraction_ctx) (fmt : F.formatter)
         extract_trait_ref ctx fmt no_params_tys false trait_ref;
         extract_generic_args ctx fmt no_params_tys generics;
         let name =
-          ctx_get_trait_assoc_type trait_ref.trait_decl_ref.trait_decl_id
-            type_name ctx
+          ctx_get_trait_type trait_ref.trait_decl_ref.trait_decl_id type_name
+            ctx
         in
         if use_brackets then F.pp_print_string fmt ")";
         F.pp_print_string fmt ("." ^ name))
@@ -1250,7 +1292,7 @@ let rec extract_ty (ctx : extraction_ctx) (fmt : F.formatter)
            types should have been normalized.
         *)
         let trait_decl_id = Option.get ctx.trait_decl_id in
-        let item_name = ctx_get_trait_assoc_type trait_decl_id type_name ctx in
+        let item_name = ctx_get_trait_type trait_decl_id type_name ctx in
         assert (generics = empty_generic_args);
         if ctx.is_provided_method then
           (* Provided method: use the trait self clause *)
@@ -3885,9 +3927,45 @@ let extract_global_decl (ctx : extraction_ctx) (fmt : F.formatter)
       F.pp_print_break fmt 0 0
 
 (** Similar to {!extract_type_decl_register_names} *)
-let extract_trait_decl_register_names (ctx : extraction_ctx) (d : trait_decl) :
-    extraction_ctx =
-  raise (Failure "TODO")
+let extract_trait_decl_register_names (ctx : extraction_ctx)
+    (trait_decl : trait_decl) : extraction_ctx =
+  let {
+    def_id = _;
+    name = _;
+    generics;
+    preds = _;
+    all_trait_clauses = _;
+    consts;
+    types;
+    required_methods;
+    provided_methods = _;
+  } =
+    trait_decl
+  in
+  let ctx = ctx_add_trait_decl trait_decl ctx in
+  let ctx =
+    List.fold_left
+      (fun ctx clause -> ctx_add_trait_parent_clause trait_decl clause ctx)
+      ctx generics.trait_clauses
+  in
+  let ctx =
+    List.fold_left
+      (fun ctx (name, (_, _)) -> ctx_add_trait_const trait_decl name ctx)
+      ctx consts
+  in
+  let ctx =
+    List.fold_left
+      (fun ctx (name, (clauses, _)) ->
+        let ctx = ctx_add_trait_type trait_decl name ctx in
+        List.fold_left
+          (fun ctx clause ->
+            ctx_add_trait_type_clause trait_decl name clause ctx)
+          ctx clauses)
+      ctx types
+  in
+  List.fold_left
+    (fun ctx (name, _) -> ctx_add_trait_method trait_decl name ctx)
+    ctx required_methods
 
 (** Similar to {!extract_type_decl_register_names} *)
 let extract_trait_impl_register_names (ctx : extraction_ctx) (d : trait_impl) :
diff --git a/compiler/ExtractBase.ml b/compiler/ExtractBase.ml
index 2855b3b9..7e6a2d40 100644
--- a/compiler/ExtractBase.ml
+++ b/compiler/ExtractBase.ml
@@ -242,6 +242,13 @@ type formatter = {
             the same purpose as in {!field:fun_name}.
           - loop identifier, if this is for a loop
        *)
+  trait_decl_name : trait_decl -> string;
+  trait_impl_name : trait_impl -> string;
+  trait_parent_clause_name : trait_decl -> trait_clause -> string;
+  trait_const_name : trait_decl -> string -> string;
+  trait_type_name : trait_decl -> string -> string;
+  trait_method_name : trait_decl -> string -> string;
+  trait_type_clause_name : trait_decl -> string -> trait_clause -> string;
   opaque_pre : unit -> string;
       (** TODO: obsolete, remove.
 
@@ -410,7 +417,7 @@ type id =
   | TraitDeclId of TraitDeclId.id
   | TraitImplId of TraitImplId.id
   | LocalTraitClauseId of TraitClauseId.id
-  | TraitAssocTypeId of TraitDeclId.id * string  (** A trait associated type *)
+  | TraitItemId of TraitDeclId.id * string  (** A trait associated item *)
   | TraitParentClauseId of TraitDeclId.id * TraitClauseId.id
   | TraitItemClauseId of TraitDeclId.id * string * TraitClauseId.id
   | TraitSelfClauseId
@@ -798,16 +805,15 @@ let id_to_string (id : id) (ctx : extraction_ctx) : string =
       "trait_item_clause_id: decl_id:" ^ TraitDeclId.to_string id
       ^ ", item name: " ^ item_name ^ ", clause_id: "
       ^ TraitClauseId.to_string clause_id
-  | TraitAssocTypeId (id, type_name) ->
-      "trait_assoc_type_id: decl_id:" ^ TraitDeclId.to_string id
-      ^ ", type name: " ^ type_name
+  | TraitItemId (id, name) ->
+      "trait_item_id: decl_id:" ^ TraitDeclId.to_string id ^ ", type name: "
+      ^ name
   | TraitSelfClauseId -> "trait_self_clause"
 
 (** We might not check for collisions for some specific ids (ex.: field names) *)
 let allow_collisions (id : id) : bool =
   match id with
-  | FieldId _ | TraitItemClauseId _ | TraitParentClauseId _ | TraitAssocTypeId _
-    ->
+  | FieldId _ | TraitItemClauseId _ | TraitParentClauseId _ | TraitItemId _ ->
       !Config.record_fields_short_names
   | _ -> false
 
@@ -910,10 +916,22 @@ let ctx_get_trait_impl (with_opaque_pre : bool) (id : trait_impl_id)
     (ctx : extraction_ctx) : string =
   ctx_get with_opaque_pre (TraitImplId id) ctx
 
-let ctx_get_trait_assoc_type (id : trait_decl_id) (type_name : string)
+let ctx_get_trait_item (id : trait_decl_id) (item_name : string)
     (ctx : extraction_ctx) : string =
   let is_opaque = false in
-  ctx_get is_opaque (TraitAssocTypeId (id, type_name)) ctx
+  ctx_get is_opaque (TraitItemId (id, item_name)) ctx
+
+let ctx_get_trait_const (id : trait_decl_id) (item_name : string)
+    (ctx : extraction_ctx) : string =
+  ctx_get_trait_item id item_name ctx
+
+let ctx_get_trait_type (id : trait_decl_id) (item_name : string)
+    (ctx : extraction_ctx) : string =
+  ctx_get_trait_item id item_name ctx
+
+let ctx_get_trait_method (id : trait_decl_id) (item_name : string)
+    (ctx : extraction_ctx) : string =
+  ctx_get_trait_item id item_name ctx
 
 let ctx_get_trait_parent_clause (id : trait_decl_id) (clause : trait_clause_id)
     (ctx : extraction_ctx) : string =
@@ -1205,6 +1223,50 @@ let ctx_add_fun_decl (trans_group : bool * pure_fun_translation)
         ctx.fun_name_info;
   }
 
+let ctx_add_trait_decl (d : trait_decl) (ctx : extraction_ctx) : extraction_ctx
+    =
+  let is_opaque = false in
+  let name = ctx.fmt.trait_decl_name d in
+  ctx_add is_opaque (TraitDeclId d.def_id) name ctx
+
+let ctx_add_trait_impl (d : trait_impl) (ctx : extraction_ctx) : extraction_ctx
+    =
+  let is_opaque = false in
+  let name = ctx.fmt.trait_impl_name d in
+  ctx_add is_opaque (TraitImplId d.def_id) name ctx
+
+let ctx_add_trait_const (d : trait_decl) (item : string) (ctx : extraction_ctx)
+    : extraction_ctx =
+  let is_opaque = false in
+  let name = ctx.fmt.trait_const_name d item in
+  ctx_add is_opaque (TraitItemId (d.def_id, item)) name ctx
+
+let ctx_add_trait_type (d : trait_decl) (item : string) (ctx : extraction_ctx) :
+    extraction_ctx =
+  let is_opaque = false in
+  let name = ctx.fmt.trait_type_name d item in
+  ctx_add is_opaque (TraitItemId (d.def_id, item)) name ctx
+
+let ctx_add_trait_method (d : trait_decl) (item : string) (ctx : extraction_ctx)
+    : extraction_ctx =
+  let is_opaque = false in
+  let name = ctx.fmt.trait_method_name d item in
+  ctx_add is_opaque (TraitItemId (d.def_id, item)) name ctx
+
+let ctx_add_trait_parent_clause (d : trait_decl) (clause : trait_clause)
+    (ctx : extraction_ctx) : extraction_ctx =
+  let is_opaque = false in
+  let name = ctx.fmt.trait_parent_clause_name d clause in
+  ctx_add is_opaque (TraitParentClauseId (d.def_id, clause.clause_id)) name ctx
+
+let ctx_add_trait_type_clause (d : trait_decl) (item : string)
+    (clause : trait_clause) (ctx : extraction_ctx) : extraction_ctx =
+  let is_opaque = false in
+  let name = ctx.fmt.trait_type_clause_name d item clause in
+  ctx_add is_opaque
+    (TraitItemClauseId (d.def_id, item, clause.clause_id))
+    name ctx
+
 type names_map_init = {
   keywords : string list;
   assumed_adts : (assumed_ty * string) list;