Make the traits work for Coq

6 files changed, 353 insertions, 193 deletions

diff --git a/compiler/Extract.ml b/compiler/Extract.ml
index e22f1385..d04f5c1d 100644
--- a/compiler/Extract.ml
+++ b/compiler/Extract.ml
@@ -320,8 +320,11 @@ and extract_App (ctx : extraction_ctx) (fmt : F.formatter) (inside : bool)
             ctx_get_trait_const trait_ref.trait_decl_ref.trait_decl_id
               const_name ctx
           in
+          let add_brackets (s : string) =
+            if !backend = Coq then "(" ^ s ^ ")" else s
+          in
           if use_brackets then F.pp_print_string fmt ")";
-          F.pp_print_string fmt ("." ^ name))
+          F.pp_print_string fmt ("." ^ add_brackets name))
   | _ ->
       (* "Regular" expression *)
       (* Open parentheses *)
@@ -430,7 +433,10 @@ and extract_function_call (ctx : extraction_ctx) (fmt : F.formatter)
               ctx_get_trait_method trait_ref.trait_decl_ref.trait_decl_id
                 method_name rg_id ctx
             in
-            F.pp_print_string fmt ("." ^ fun_name))
+            let add_brackets (s : string) =
+              if !backend = Coq then "(" ^ s ^ ")" else s
+            in
+            F.pp_print_string fmt ("." ^ add_brackets fun_name))
           else
             (* Provided method: we see it as a regular function call, and use
                the function name *)
@@ -2021,12 +2027,15 @@ let extract_trait_decl_register_names (ctx : extraction_ctx)
     SimpleNameMap.find_opt sname (builtin_trait_decls_map ())
   in
   let ctx =
-    let trait_name =
+    let trait_name, trait_constructor =
       match builtin_info with
-      | None -> ctx.fmt.trait_decl_name trait_decl
-      | Some info -> info.extract_name
+      | None ->
+          ( ctx.fmt.trait_decl_name trait_decl,
+            ctx.fmt.trait_decl_constructor trait_decl )
+      | Some info -> (info.extract_name, info.constructor)
     in
-    ctx_add (TraitDeclId trait_decl.def_id) trait_name ctx
+    let ctx = ctx_add (TraitDeclId trait_decl.def_id) trait_name ctx in
+    ctx_add (TraitDeclConstructorId trait_decl.def_id) trait_constructor ctx
   in
   (* Parent clauses *)
   let ctx =
@@ -2108,7 +2117,7 @@ let extract_trait_decl_item (ctx : extraction_ctx) (fmt : F.formatter)
 
 let extract_trait_impl_item (ctx : extraction_ctx) (fmt : F.formatter)
     (item_name : string) (ty : unit -> unit) : unit =
-  let assign = match !Config.backend with Lean -> ":=" | _ -> "=" in
+  let assign = match !Config.backend with Lean | Coq -> ":=" | _ -> "=" in
   extract_trait_item ctx fmt item_name assign ty
 
 (** Small helper - TODO: move *)
@@ -2215,87 +2224,173 @@ let extract_trait_decl (ctx : extraction_ctx) (fmt : F.formatter)
     cg_params trait_clauses;
 
   F.pp_print_space fmt ();
-  (match !backend with
-  | Lean -> F.pp_print_string fmt "where"
-  | FStar -> if not is_empty then F.pp_print_string fmt "= {"
-  | _ -> F.pp_print_string fmt "{");
-  if !backend = FStar && is_empty then F.pp_print_string fmt "= unit";
+  if is_empty && !backend = FStar then (
+    F.pp_print_string fmt "= unit";
+    (* Outer box *)
+    F.pp_close_box fmt ())
+  else if is_empty && !backend = Coq then (
+    (* Coq is not very good at infering constructors *)
+    let cons = ctx_get_trait_constructor decl.def_id ctx in
+    F.pp_print_string fmt (":= " ^ cons ^ "{}.");
+    (* Outer box *)
+    F.pp_close_box fmt ())
+  else (
+    (match !backend with
+    | Lean -> F.pp_print_string fmt "where"
+    | FStar -> F.pp_print_string fmt "= {"
+    | Coq ->
+        let cons = ctx_get_trait_constructor decl.def_id ctx in
+        F.pp_print_string fmt (":= " ^ cons ^ " {")
+    | _ -> F.pp_print_string fmt "{");
 
-  (* Close the box for the name + generics *)
-  F.pp_close_box fmt ();
+    (* Close the box for the name + generics *)
+    F.pp_close_box fmt ();
 
-  (*
-   * Extract the items
-   *)
+    (*
+     * Extract the items
+     *)
 
-  (* The constants *)
-  List.iter
-    (fun (name, (ty, _)) ->
-      let item_name = ctx_get_trait_const decl.def_id name ctx in
-      let ty () =
-        let inside = false in
-        F.pp_print_space fmt ();
-        extract_ty ctx fmt TypeDeclId.Set.empty inside ty
-      in
-      extract_trait_decl_item ctx fmt item_name ty)
-    decl.consts;
+    (* The constants *)
+    List.iter
+      (fun (name, (ty, _)) ->
+        let item_name = ctx_get_trait_const decl.def_id name ctx in
+        let ty () =
+          let inside = false in
+          F.pp_print_space fmt ();
+          extract_ty ctx fmt TypeDeclId.Set.empty inside ty
+        in
+        extract_trait_decl_item ctx fmt item_name ty)
+      decl.consts;
 
-  (* The types *)
-  List.iter
-    (fun (name, (clauses, _)) ->
-      (* Extract the type *)
-      let item_name = ctx_get_trait_type decl.def_id name ctx in
-      let ty () =
-        F.pp_print_space fmt ();
-        F.pp_print_string fmt (type_keyword ())
-      in
-      extract_trait_decl_item ctx fmt item_name ty;
-      (* Extract the clauses *)
-      List.iter
-        (fun clause ->
-          let item_name =
-            ctx_get_trait_item_clause decl.def_id name clause.clause_id ctx
-          in
-          let ty () =
-            F.pp_print_space fmt ();
-            extract_trait_clause_type ctx fmt TypeDeclId.Set.empty clause
-          in
-          extract_trait_decl_item ctx fmt item_name ty)
-        clauses)
-    decl.types;
+    (* The types *)
+    List.iter
+      (fun (name, (clauses, _)) ->
+        (* Extract the type *)
+        let item_name = ctx_get_trait_type decl.def_id name ctx in
+        let ty () =
+          F.pp_print_space fmt ();
+          F.pp_print_string fmt (type_keyword ())
+        in
+        extract_trait_decl_item ctx fmt item_name ty;
+        (* Extract the clauses *)
+        List.iter
+          (fun clause ->
+            let item_name =
+              ctx_get_trait_item_clause decl.def_id name clause.clause_id ctx
+            in
+            let ty () =
+              F.pp_print_space fmt ();
+              extract_trait_clause_type ctx fmt TypeDeclId.Set.empty clause
+            in
+            extract_trait_decl_item ctx fmt item_name ty)
+          clauses)
+      decl.types;
 
-  (* The parent clauses - note that the parent clauses may refer to the types
-     and const generics: for this reason we extract them *after* *)
-  List.iter
-    (fun clause ->
-      let item_name =
-        ctx_get_trait_parent_clause decl.def_id clause.clause_id ctx
-      in
-      let ty () =
-        F.pp_print_space fmt ();
-        extract_trait_clause_type ctx fmt TypeDeclId.Set.empty clause
-      in
-      extract_trait_decl_item ctx fmt item_name ty)
-    decl.parent_clauses;
+    (* The parent clauses - note that the parent clauses may refer to the types
+       and const generics: for this reason we extract them *after* *)
+    List.iter
+      (fun clause ->
+        let item_name =
+          ctx_get_trait_parent_clause decl.def_id clause.clause_id ctx
+        in
+        let ty () =
+          F.pp_print_space fmt ();
+          extract_trait_clause_type ctx fmt TypeDeclId.Set.empty clause
+        in
+        extract_trait_decl_item ctx fmt item_name ty)
+      decl.parent_clauses;
 
-  (* The required methods *)
-  List.iter
-    (fun (name, id) -> extract_trait_decl_method_items ctx fmt decl name id)
-    decl.required_methods;
-
-  (* Close the outer boxes for the definition *)
-  if !Config.backend <> Lean then F.pp_close_box fmt ();
-  (* Close the brackets *)
-  (match !Config.backend with
-  | Lean -> ()
-  | _ ->
-      if (not (!backend = FStar)) || not is_empty then (
+    (* The required methods *)
+    List.iter
+      (fun (name, id) -> extract_trait_decl_method_items ctx fmt decl name id)
+      decl.required_methods;
+
+    (* Close the outer boxes for the definition *)
+    if !Config.backend <> Lean then F.pp_close_box fmt ();
+    (* Close the brackets *)
+    match !Config.backend with
+    | Lean -> ()
+    | Coq ->
+        F.pp_print_space fmt ();
+        F.pp_print_string fmt "}."
+    | _ ->
         F.pp_print_space fmt ();
-        F.pp_print_string fmt "}"));
+        F.pp_print_string fmt "}");
   F.pp_close_box fmt ();
   (* Add breaks to insert new lines between definitions *)
   F.pp_print_break fmt 0 0
 
+(** Generate the [Arguments] instructions for the trait declarationsin Coq, so
+    that we don't have to provide the implicit arguments when projecting the fields. *)
+let extract_trait_decl_coq_arguments (ctx : extraction_ctx) (fmt : F.formatter)
+    (decl : trait_decl) : unit =
+  (* Generating the [Arguments] instructions is useful only if there are parameters *)
+  let num_params =
+    List.length decl.generics.types
+    + List.length decl.generics.const_generics
+    + List.length decl.generics.trait_clauses
+  in
+  if num_params > 0 then (
+    (* The constructor *)
+    let cons_name = ctx_get_trait_constructor decl.def_id ctx in
+    extract_coq_arguments_instruction ctx fmt cons_name num_params;
+    (* The constants *)
+    List.iter
+      (fun (name, _) ->
+        let item_name = ctx_get_trait_const decl.def_id name ctx in
+        extract_coq_arguments_instruction ctx fmt item_name num_params)
+      decl.consts;
+    (* The types *)
+    List.iter
+      (fun (name, (clauses, _)) ->
+        (* The type *)
+        let item_name = ctx_get_trait_type decl.def_id name ctx in
+        extract_coq_arguments_instruction ctx fmt item_name num_params;
+        (* The type clauses *)
+        List.iter
+          (fun clause ->
+            let item_name =
+              ctx_get_trait_item_clause decl.def_id name clause.clause_id ctx
+            in
+            extract_coq_arguments_instruction ctx fmt item_name num_params)
+          clauses)
+      decl.types;
+    (* The parent clauses *)
+    List.iter
+      (fun clause ->
+        let item_name =
+          ctx_get_trait_parent_clause decl.def_id clause.clause_id ctx
+        in
+        extract_coq_arguments_instruction ctx fmt item_name num_params)
+      decl.parent_clauses;
+    (* The required methods *)
+    List.iter
+      (fun (item_name, id) ->
+        (* Lookup the definition *)
+        let trans = A.FunDeclId.Map.find id ctx.trans_funs in
+        (* Extract the items *)
+        let funs =
+          if trans.keep_fwd then trans.fwd :: trans.backs else trans.backs
+        in
+        let extract_for_method (f : fun_and_loops) =
+          let f = f.f in
+          let item_name =
+            ctx_get_trait_method decl.def_id item_name f.back_id ctx
+          in
+          extract_coq_arguments_instruction ctx fmt item_name num_params
+        in
+        List.iter extract_for_method funs)
+      decl.required_methods;
+    (* Add a space *)
+    F.pp_print_space fmt ())
+
+(** See {!extract_trait_decl_coq_arguments} *)
+let extract_trait_decl_extra_info (ctx : extraction_ctx) (fmt : F.formatter)
+    (trait_decl : trait_decl) : unit =
+  match !backend with
+  | Coq -> extract_trait_decl_coq_arguments ctx fmt trait_decl
+  | _ -> ()
+
 (** Small helper.
 
     Extract the items for a method in a trait impl.
@@ -2425,76 +2520,92 @@ let extract_trait_impl (ctx : extraction_ctx) (fmt : F.formatter)
   let is_empty = trait_impl_is_empty { impl with provided_methods = [] } in
 
   F.pp_print_space fmt ();
-  if !Config.backend = Lean then F.pp_print_string fmt ":= {"
-  else if !Config.backend = FStar && is_empty then F.pp_print_string fmt "= ()"
-  else F.pp_print_string fmt "= {";
+  if is_empty && !Config.backend = FStar then (
+    F.pp_print_string fmt "= ()";
+    (* Outer box *)
+    F.pp_close_box fmt ())
+  else if is_empty && !Config.backend = Coq then (
+    (* Coq is not very good at infering constructors *)
+    let cons = ctx_get_trait_constructor impl.impl_trait.trait_decl_id ctx in
+    F.pp_print_string fmt (":= " ^ cons ^ ".");
+    (* Outer box *)
+    F.pp_close_box fmt ())
+  else (
+    if !Config.backend = Lean then F.pp_print_string fmt ":= {"
+    else if !Config.backend = Coq then F.pp_print_string fmt ":= {|"
+    else F.pp_print_string fmt "= {";
 
-  (* Close the box for the name + generics *)
-  F.pp_close_box fmt ();
+    (* Close the box for the name + generics *)
+    F.pp_close_box fmt ();
 
-  (*
-   * Extract the items
-   *)
-  let trait_decl_id = impl.impl_trait.trait_decl_id in
+    (*
+     * Extract the items
+     *)
+    let trait_decl_id = impl.impl_trait.trait_decl_id in
 
-  (* The constants *)
-  List.iter
-    (fun (name, (_, id)) ->
-      let item_name = ctx_get_trait_const trait_decl_id name ctx in
-      let ty () =
-        F.pp_print_space fmt ();
-        F.pp_print_string fmt (ctx_get_global id ctx)
-      in
+    (* The constants *)
+    List.iter
+      (fun (name, (_, id)) ->
+        let item_name = ctx_get_trait_const trait_decl_id name ctx in
+        let ty () =
+          F.pp_print_space fmt ();
+          F.pp_print_string fmt (ctx_get_global id ctx)
+        in
 
-      extract_trait_impl_item ctx fmt item_name ty)
-    impl.consts;
+        extract_trait_impl_item ctx fmt item_name ty)
+      impl.consts;
 
-  (* The types *)
-  List.iter
-    (fun (name, (trait_refs, ty)) ->
-      (* Extract the type *)
-      let item_name = ctx_get_trait_type trait_decl_id name ctx in
-      let ty () =
-        F.pp_print_space fmt ();
-        extract_ty ctx fmt TypeDeclId.Set.empty false ty
-      in
-      extract_trait_impl_item ctx fmt item_name ty;
-      (* Extract the clauses *)
-      TraitClauseId.iteri
-        (fun clause_id trait_ref ->
-          let item_name =
-            ctx_get_trait_item_clause trait_decl_id name clause_id ctx
-          in
-          let ty () =
-            F.pp_print_space fmt ();
-            extract_trait_ref ctx fmt TypeDeclId.Set.empty false trait_ref
-          in
-          extract_trait_impl_item ctx fmt item_name ty)
-        trait_refs)
-    impl.types;
-
-  (* The parent clauses *)
-  TraitClauseId.iteri
-    (fun clause_id trait_ref ->
-      let item_name = ctx_get_trait_parent_clause trait_decl_id clause_id ctx in
-      let ty () =
-        F.pp_print_space fmt ();
-        extract_trait_ref ctx fmt TypeDeclId.Set.empty false trait_ref
-      in
-      extract_trait_impl_item ctx fmt item_name ty)
-    impl.parent_trait_refs;
+    (* The types *)
+    List.iter
+      (fun (name, (trait_refs, ty)) ->
+        (* Extract the type *)
+        let item_name = ctx_get_trait_type trait_decl_id name ctx in
+        let ty () =
+          F.pp_print_space fmt ();
+          extract_ty ctx fmt TypeDeclId.Set.empty false ty
+        in
+        extract_trait_impl_item ctx fmt item_name ty;
+        (* Extract the clauses *)
+        TraitClauseId.iteri
+          (fun clause_id trait_ref ->
+            let item_name =
+              ctx_get_trait_item_clause trait_decl_id name clause_id ctx
+            in
+            let ty () =
+              F.pp_print_space fmt ();
+              extract_trait_ref ctx fmt TypeDeclId.Set.empty false trait_ref
+            in
+            extract_trait_impl_item ctx fmt item_name ty)
+          trait_refs)
+      impl.types;
+
+    (* The parent clauses *)
+    TraitClauseId.iteri
+      (fun clause_id trait_ref ->
+        let item_name =
+          ctx_get_trait_parent_clause trait_decl_id clause_id ctx
+        in
+        let ty () =
+          F.pp_print_space fmt ();
+          extract_trait_ref ctx fmt TypeDeclId.Set.empty false trait_ref
+        in
+        extract_trait_impl_item ctx fmt item_name ty)
+      impl.parent_trait_refs;
 
-  (* The required methods *)
-  List.iter
-    (fun (name, id) ->
-      extract_trait_impl_method_items ctx fmt impl name id all_generics)
-    impl.required_methods;
+    (* The required methods *)
+    List.iter
+      (fun (name, id) ->
+        extract_trait_impl_method_items ctx fmt impl name id all_generics)
+      impl.required_methods;
 
-  (* Close the outer boxes for the definition, as well as the brackets *)
-  F.pp_close_box fmt ();
-  if (not (!backend = FStar)) || not is_empty then (
-    F.pp_print_space fmt ();
-    F.pp_print_string fmt "}");
+    (* Close the outer boxes for the definition, as well as the brackets *)
+    F.pp_close_box fmt ();
+    if !backend = Coq then (
+      F.pp_print_space fmt ();
+      F.pp_print_string fmt "|}.")
+    else if (not (!backend = FStar)) || not is_empty then (
+      F.pp_print_space fmt ();
+      F.pp_print_string fmt "}"));
   F.pp_close_box fmt ();
   (* Add breaks to insert new lines between definitions *)
   F.pp_print_break fmt 0 0
diff --git a/compiler/ExtractBase.ml b/compiler/ExtractBase.ml
index 55b1bca3..31b1a447 100644
--- a/compiler/ExtractBase.ml
+++ b/compiler/ExtractBase.ml
@@ -246,6 +246,7 @@ type formatter = {
        *)
   trait_decl_name : trait_decl -> string;
   trait_impl_name : trait_decl -> trait_impl -> string;
+  trait_decl_constructor : trait_decl -> string;
   trait_parent_clause_name : trait_decl -> trait_clause -> string;
   trait_const_name : trait_decl -> string -> string;
   trait_type_name : trait_decl -> string -> string;
@@ -388,6 +389,7 @@ type id =
   | TraitDeclId of TraitDeclId.id
   | TraitImplId of TraitImplId.id
   | LocalTraitClauseId of TraitClauseId.id
+  | TraitDeclConstructorId of TraitDeclId.id
   | TraitMethodId of TraitDeclId.id * string * T.RegionGroupId.id option
       (** Something peculiar with trait methods: because we have to take into
           account forward/backward functions, we may need to generate fields
@@ -801,6 +803,8 @@ let id_to_string (id : id) (ctx : extraction_ctx) : string =
   | TraitImplId id -> "trait_impl_id: " ^ TraitImplId.to_string id
   | LocalTraitClauseId id ->
       "local_trait_clause_id: " ^ TraitClauseId.to_string id
+  | TraitDeclConstructorId id ->
+      "trait_decl_constructor: " ^ trait_decl_id_to_string id
   | TraitParentClauseId (id, clause_id) ->
       "trait_parent_clause_id: " ^ trait_decl_id_to_string id ^ ", clause_id: "
       ^ TraitClauseId.to_string clause_id
@@ -959,6 +963,10 @@ let ctx_get_local_type (id : TypeDeclId.id) (ctx : extraction_ctx) : string =
 let ctx_get_assumed_type (id : assumed_ty) (ctx : extraction_ctx) : string =
   ctx_get_type (Assumed id) ctx
 
+let ctx_get_trait_constructor (id : trait_decl_id) (ctx : extraction_ctx) :
+    string =
+  ctx_get (TraitDeclConstructorId id) ctx
+
 let ctx_get_trait_self_clause (ctx : extraction_ctx) : string =
   ctx_get TraitSelfClauseId ctx
 
diff --git a/compiler/ExtractBuiltin.ml b/compiler/ExtractBuiltin.ml
index c6bde9c2..a54ab604 100644
--- a/compiler/ExtractBuiltin.ml
+++ b/compiler/ExtractBuiltin.ml
@@ -420,6 +420,7 @@ let builtin_fun_effects_map =
 type builtin_trait_decl_info = {
   rust_name : string;
   extract_name : string;
+  constructor : string;
   parent_clauses : string list;
   consts : (string * string) list;
   types : (string * (string * string list)) list;
@@ -444,6 +445,7 @@ let builtin_trait_decls_info () =
           | Coq | FStar | HOL4 -> String.concat "_" rust_name
           | Lean -> String.concat "." rust_name)
     in
+    let constructor = mk_struct_constructor extract_name in
     let consts = [] in
     let types =
       let mk_type item_name =
@@ -479,7 +481,15 @@ let builtin_trait_decls_info () =
       List.map mk_method methods
     in
     let rust_name = String.concat "::" rust_name in
-    { rust_name; extract_name; parent_clauses; consts; types; methods }
+    {
+      rust_name;
+      extract_name;
+      constructor;
+      parent_clauses;
+      consts;
+      types;
+      methods;
+    }
   in
   [
     (* Deref *)
diff --git a/compiler/ExtractTypes.ml b/compiler/ExtractTypes.ml
index a294d4ca..77f76bb4 100644
--- a/compiler/ExtractTypes.ml
+++ b/compiler/ExtractTypes.ml
@@ -697,6 +697,11 @@ let mk_formatter (ctx : trans_ctx) (crate_name : string)
     | Coq | HOL4 | Lean -> name
   in
 
+  let trait_decl_constructor (trait_decl : trait_decl) : string =
+    let name = trait_decl_name trait_decl in
+    ExtractBuiltin.mk_struct_constructor 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 *)
@@ -937,6 +942,7 @@ let mk_formatter (ctx : trans_ctx) (crate_name : string)
     decreases_proof_name;
     trait_decl_name;
     trait_impl_name;
+    trait_decl_constructor;
     trait_parent_clause_name;
     trait_const_name;
     trait_type_name;
@@ -1254,6 +1260,9 @@ let rec extract_ty (ctx : extraction_ctx) (fmt : F.formatter)
           ctx_get_trait_type trait_ref.trait_decl_ref.trait_decl_id type_name
             ctx
         in
+        let add_brackets (s : string) =
+          if !backend = Coq then "(" ^ s ^ ")" else s
+        in
         (* There may be a special treatment depending on the instance id.
            See the comments for {!extract_trait_instance_id_with_dot}.
            TODO: there should be a cleaner way to do. The annoying thing
@@ -1276,7 +1285,7 @@ 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;
             if use_brackets then F.pp_print_string fmt ")";
-            F.pp_print_string fmt ("." ^ type_name))
+            F.pp_print_string fmt ("." ^ add_brackets type_name))
 
 and extract_trait_ref (ctx : extraction_ctx) (fmt : F.formatter)
     (no_params_tys : TypeDeclId.Set.t) (inside : bool) (tr : trait_ref) : unit =
@@ -1376,6 +1385,7 @@ and extract_trait_instance_id_with_dot (ctx : extraction_ctx)
 and extract_trait_instance_id (ctx : extraction_ctx) (fmt : F.formatter)
     (no_params_tys : TypeDeclId.Set.t) (inside : bool) (id : trait_instance_id)
     : unit =
+  let add_brackets (s : string) = if !backend = Coq then "(" ^ s ^ ")" else s in
   match id with
   | Self ->
       (* This has a specific treatment depending on the item we're extracting
@@ -1393,12 +1403,12 @@ and extract_trait_instance_id (ctx : extraction_ctx) (fmt : F.formatter)
       (* Use the trait decl id to lookup the name *)
       let name = ctx_get_trait_parent_clause decl_id clause_id ctx in
       extract_trait_instance_id_with_dot ctx fmt no_params_tys true inst_id;
-      F.pp_print_string fmt name
+      F.pp_print_string fmt (add_brackets name)
   | ItemClause (inst_id, decl_id, item_name, clause_id) ->
       (* Use the trait decl id to lookup the name *)
       let name = ctx_get_trait_item_clause decl_id item_name clause_id ctx in
       extract_trait_instance_id_with_dot ctx fmt no_params_tys true inst_id;
-      F.pp_print_string fmt name
+      F.pp_print_string fmt (add_brackets name)
   | TraitRef trait_ref ->
       extract_trait_ref ctx fmt no_params_tys inside trait_ref
   | UnknownTrait _ ->
@@ -2156,49 +2166,59 @@ let extract_type_decl (ctx : extraction_ctx) (fmt : F.formatter)
         extract_type_decl_gen ctx fmt type_decl_group kind def extract_body
     | HOL4 -> extract_type_decl_hol4_opaque ctx fmt def
 
+(** Generate a [Argument] instruction in Coq to allow omitting implicit
+    arguments for variants, fields, etc..
+
+    For instance, provided we have this definition:
+    {[
+      Inductive result A :=
+      | Return : A -> result A
+      | Fail_ : error -> result A.
+    ]}
+
+    We may want to generate those instructions:
+    {[
+      Arguments Return {_} a.
+      Arguments Fail_ {_}.
+    ]}
+ *)
+let extract_coq_arguments_instruction (ctx : extraction_ctx) (fmt : F.formatter)
+    (cons_name : string) (num_implicit_params : int) : unit =
+  (* Add a break before *)
+  F.pp_print_break fmt 0 0;
+  (* Open a box *)
+  F.pp_open_hovbox fmt ctx.indent_incr;
+  F.pp_print_break fmt 0 0;
+  F.pp_print_string fmt "Arguments";
+  F.pp_print_space fmt ();
+  F.pp_print_string fmt cons_name;
+  (* Print the type/const params and the trait clauses (`{T}`) *)
+  F.pp_print_space fmt ();
+  F.pp_print_string fmt "{";
+  Collections.List.iter_times num_implicit_params (fun () ->
+      F.pp_print_space fmt ();
+      F.pp_print_string fmt "_");
+  F.pp_print_space fmt ();
+  F.pp_print_string fmt "}.";
+
+  (* Close the box *)
+  F.pp_close_box fmt ()
+
 (** Auxiliary function.
 
-    Generate [Arguments] instructions in Coq.
+    Generate [Arguments] instructions in Coq for type definitions.
  *)
 let extract_type_decl_coq_arguments (ctx : extraction_ctx) (fmt : F.formatter)
     (kind : decl_kind) (decl : type_decl) : unit =
   assert (!backend = Coq);
-  (* Generating the [Arguments] instructions is useful only if there are type parameters *)
-  if decl.generics.types = [] && decl.generics.const_generics = [] then ()
+  (* Generating the [Arguments] instructions is useful only if there are parameters *)
+  let num_params =
+    List.length decl.generics.types
+    + List.length decl.generics.const_generics
+    + List.length decl.generics.trait_clauses
+  in
+  if num_params = 0 then ()
   else
-    (* Add the type params - note that we need those bindings only for the
-     * body translation (they are not top-level) *)
-    let _ctx_body, type_params, cg_params, trait_clauses =
-      ctx_add_generic_params decl.generics ctx
-    in
-    (* Auxiliary function to extract an [Arguments Cons {T} _ _.] instruction *)
-    let extract_arguments_info (cons_name : string) (fields : 'a list) : unit =
-      (* Add a break before *)
-      F.pp_print_break fmt 0 0;
-      (* Open a box *)
-      F.pp_open_hovbox fmt ctx.indent_incr;
-      F.pp_print_break fmt 0 0;
-      F.pp_print_string fmt "Arguments";
-      F.pp_print_space fmt ();
-      F.pp_print_string fmt cons_name;
-      (* Print the type/const params and the trait clauses (`{T}`) *)
-      List.iter
-        (fun (var : string) ->
-          F.pp_print_space fmt ();
-          F.pp_print_string fmt ("{" ^ var ^ "}"))
-        (List.concat [ type_params; cg_params; trait_clauses ]);
-      (* Print the fields (`_`) *)
-      List.iter
-        (fun _ ->
-          F.pp_print_space fmt ();
-          F.pp_print_string fmt "_")
-        fields;
-      F.pp_print_string fmt ".";
-
-      (* Close the box *)
-      F.pp_close_box fmt ()
-    in
-
     (* Generate the [Arguments] instruction *)
     match decl.kind with
     | Opaque -> ()
@@ -2206,23 +2226,23 @@ let extract_type_decl_coq_arguments (ctx : extraction_ctx) (fmt : F.formatter)
         let adt_id = AdtId decl.def_id in
         (* Generate the instruction for the record constructor *)
         let cons_name = ctx_get_struct adt_id ctx in
-        extract_arguments_info cons_name fields;
+        extract_coq_arguments_instruction ctx fmt cons_name num_params;
         (* Generate the instruction for the record projectors, if there are *)
         let is_rec = decl_is_from_rec_group kind in
         if not is_rec then
           FieldId.iteri
             (fun fid _ ->
               let cons_name = ctx_get_field adt_id fid ctx in
-              extract_arguments_info cons_name [])
+              extract_coq_arguments_instruction ctx fmt cons_name num_params)
             fields;
         (* Add breaks to insert new lines between definitions *)
         F.pp_print_break fmt 0 0
     | Enum variants ->
         (* Generate the instructions *)
         VariantId.iteri
-          (fun vid (v : variant) ->
+          (fun vid (_ : variant) ->
             let cons_name = ctx_get_variant (AdtId decl.def_id) vid ctx in
-            extract_arguments_info cons_name v.fields)
+            extract_coq_arguments_instruction ctx fmt cons_name num_params)
           variants;
         (* Add breaks to insert new lines between definitions *)
         F.pp_print_break fmt 0 0
diff --git a/compiler/FunsAnalysis.ml b/compiler/FunsAnalysis.ml
index 69c0df71..e17ea16f 100644
--- a/compiler/FunsAnalysis.ml
+++ b/compiler/FunsAnalysis.ml
@@ -104,9 +104,10 @@ let analyze_module (m : crate) (funs_map : fun_decl FunDeclId.Map.t)
                 (* None of the assumed functions can diverge nor are considered stateful *)
                 can_fail := !can_fail || Assumed.assumed_fun_can_fail id
             | TraitMethod _ ->
-                (* We consider trait functions can fail, diverge, and are not stateful *)
-                can_fail := true;
-                can_diverge := true);
+                (* We consider trait functions can fail, but can not diverge and are not stateful.
+                   TODO: this may cause issues if we use use a fuel parameter.
+                *)
+                can_fail := true);
             super#visit_Call env call
 
           method! visit_Panic env =
diff --git a/compiler/Translate.ml b/compiler/Translate.ml
index cb23198a..a3d96023 100644
--- a/compiler/Translate.ml
+++ b/compiler/Translate.ml
@@ -751,14 +751,17 @@ let export_functions_group (fmt : Format.formatter) (config : gen_config)
 
 (** Export a trait declaration. *)
 let export_trait_decl (fmt : Format.formatter) (_config : gen_config)
-    (ctx : gen_ctx) (trait_decl_id : Pure.trait_decl_id) : unit =
+    (ctx : gen_ctx) (trait_decl_id : Pure.trait_decl_id) (extract_decl : bool)
+    (extract_extra_info : bool) : unit =
   let trait_decl = T.TraitDeclId.Map.find trait_decl_id ctx.trans_trait_decls in
   (* Check if the trait declaration is builtin, in which case we ignore it *)
   let open ExtractBuiltin in
   let sname = name_to_simple_name trait_decl.name in
-  if SimpleNameMap.find_opt sname (builtin_trait_decls_map ()) = None then
+  if SimpleNameMap.find_opt sname (builtin_trait_decls_map ()) = None then (
     let ctx = { ctx with trait_decl_id = Some trait_decl.def_id } in
-    Extract.extract_trait_decl ctx fmt trait_decl
+    if extract_decl then Extract.extract_trait_decl ctx fmt trait_decl;
+    if extract_extra_info then
+      Extract.extract_trait_decl_extra_info ctx fmt trait_decl)
   else ()
 
 (** Export a trait implementation. *)
@@ -796,7 +799,12 @@ let extract_definitions (fmt : Format.formatter) (config : gen_config)
   let export_functions_group = export_functions_group fmt config ctx in
   let export_global = export_global fmt config ctx in
   let export_types_group = export_types_group fmt config ctx in
-  let export_trait_decl = export_trait_decl fmt config ctx in
+  let export_trait_decl_group id =
+    export_trait_decl fmt config ctx id true false
+  in
+  let export_trait_decl_group_extra_info id =
+    export_trait_decl fmt config ctx id false true
+  in
   let export_trait_impl = export_trait_impl fmt config ctx in
 
   let export_state_type () : unit =
@@ -833,8 +841,10 @@ let extract_definitions (fmt : Format.formatter) (config : gen_config)
         export_functions_group pure_funs
     | Global id -> export_global id
     | TraitDecl id ->
-        if config.extract_trait_decls && config.extract_transparent then
-          export_trait_decl id
+        (* TODO: update to extract groups *)
+        if config.extract_trait_decls && config.extract_transparent then (
+          export_trait_decl_group id;
+          export_trait_decl_group_extra_info id)
     | TraitImpl id ->
         if config.extract_trait_impls && config.extract_transparent then
           export_trait_impl id