Update the way function names are handled

16 files changed, 91 insertions, 36 deletions

diff --git a/src/Assumed.ml b/src/Assumed.ml
index 527b2395..38e539b4 100644
--- a/src/Assumed.ml
+++ b/src/Assumed.ml
@@ -271,9 +271,9 @@ let get_assumed_sig (id : A.assumed_fun_id) : A.fun_sig =
   let _, sg, _, _ = get_assumed_info id in
   sg
 
-let get_assumed_name (id : A.assumed_fun_id) : Identifiers.name =
+let get_assumed_name (id : A.assumed_fun_id) : Identifiers.fun_name =
   let _, _, _, name = get_assumed_info id in
-  name
+  Identifiers.Regular name
 
 let assumed_is_monadic (id : A.assumed_fun_id) : bool =
   let _, _, b, _ = get_assumed_info id in
diff --git a/src/CfimAst.ml b/src/CfimAst.ml
index 6716c793..ab7bcb4a 100644
--- a/src/CfimAst.ml
+++ b/src/CfimAst.ml
@@ -171,7 +171,7 @@ and switch_targets =
 
 type fun_def = {
   def_id : FunDefId.id;
-  name : name;
+  name : fun_name;
   signature : fun_sig;
   arg_count : int;
   locals : var list;
diff --git a/src/CfimAstUtils.ml b/src/CfimAstUtils.ml
index 6a2f680a..04d3e6b0 100644
--- a/src/CfimAstUtils.ml
+++ b/src/CfimAstUtils.ml
@@ -26,7 +26,7 @@ let lookup_fun_sig (fun_id : fun_id) (fun_defs : fun_def FunDefId.Map.t) :
   | Assumed aid -> Assumed.get_assumed_sig aid
 
 let lookup_fun_name (fun_id : fun_id) (fun_defs : fun_def FunDefId.Map.t) :
-    Identifiers.name =
+    Identifiers.fun_name =
   match fun_id with
   | Local id -> (FunDefId.Map.find id fun_defs).name
   | Assumed aid -> Assumed.get_assumed_name aid
diff --git a/src/CfimOfJson.ml b/src/CfimOfJson.ml
index f91f7b05..8278668a 100644
--- a/src/CfimOfJson.ml
+++ b/src/CfimOfJson.ml
@@ -24,6 +24,19 @@ let log = Logging.cfim_of_json_logger
 let name_of_json (js : json) : (name, string) result =
   combine_error_msgs js "name_of_json" (list_of_json string_of_json js)
 
+let fun_name_of_json (js : json) : (fun_name, string) result =
+  combine_error_msgs js "fun_name_of_json"
+    (match js with
+    | `Assoc [ ("Regular", name) ] ->
+        let* name = name_of_json name in
+        Ok (Regular name)
+    | `Assoc [ ("Impl", `List [ type_name; impl_id; ident ]) ] ->
+        let* type_name = name_of_json type_name in
+        let* impl_id = ImplId.id_of_json impl_id in
+        let* ident = string_of_json ident in
+        Ok (Impl (type_name, impl_id, ident))
+    | _ -> Error "")
+
 let type_var_of_json (js : json) : (T.type_var, string) result =
   combine_error_msgs js "type_var_of_json"
     (match js with
@@ -606,7 +619,7 @@ let fun_def_of_json (js : json) : (A.fun_def, string) result =
           ("body", body);
         ] ->
         let* def_id = A.FunDefId.id_of_json def_id in
-        let* name = name_of_json name in
+        let* name = fun_name_of_json name in
         let* signature = fun_sig_of_json signature in
         let* arg_count = int_of_json arg_count in
         let* locals = list_of_json var_of_json locals in
diff --git a/src/ExtractToFStar.ml b/src/ExtractToFStar.ml
index dcaef438..c4323090 100644
--- a/src/ExtractToFStar.ml
+++ b/src/ExtractToFStar.ml
@@ -256,14 +256,15 @@ let mk_formatter (ctx : trans_ctx) (variant_concatenate_type_name : bool) :
    * `module::function` (if top function)
    * `module::Type::function` (for implementations)
    *)
-  let get_fun_name (name : name) : string =
+  let get_fun_name (name : fun_name) : string =
     match name with
-    | [ _module; name ] -> name
-    | [ _module; ty; name ] -> to_snake_case ty ^ "_" ^ name
+    | Regular [ _module; name ] -> name
+    | Impl ([ _module; ty ], _, name) -> to_snake_case ty ^ "_" ^ name
     | _ ->
-        raise (Failure ("Unexpected name shape: " ^ Print.name_to_string name))
+        raise
+          (Failure ("Unexpected name shape: " ^ Print.fun_name_to_string name))
   in
-  let fun_name (_fid : A.fun_id) (fname : name) (num_rgs : int)
+  let fun_name (_fid : A.fun_id) (fname : fun_name) (num_rgs : int)
       (rg : region_group_info option) (filter_info : bool * int) : string =
     let fname = get_fun_name fname in
     (* Converting to snake case should be a no-op, but it doesn't cost much *)
@@ -274,7 +275,7 @@ let mk_formatter (ctx : trans_ctx) (variant_concatenate_type_name : bool) :
     fname ^ suffix
   in
 
-  let decreases_clause_name (_fid : FunDefId.id) (fname : name) : string =
+  let decreases_clause_name (_fid : FunDefId.id) (fname : fun_name) : string =
     let fname = get_fun_name fname in
     (* Converting to snake case should be a no-op, but it doesn't cost much *)
     let fname = to_snake_case fname in
@@ -1051,7 +1052,7 @@ let extract_template_decreases_clause (ctx : extraction_ctx) (fmt : F.formatter)
   F.pp_print_break fmt 0 0;
   (* Print a comment to link the extracted type to its original rust definition *)
   F.pp_print_string fmt
-    ("(** [" ^ Print.name_to_string def.basename ^ "]: decreases clause *)");
+    ("(** [" ^ Print.fun_name_to_string def.basename ^ "]: decreases clause *)");
   F.pp_print_space fmt ();
   (* Open a box for the definition, so that whenever possible it gets printed on
    * one line *)
@@ -1107,7 +1108,8 @@ let extract_fun_def (ctx : extraction_ctx) (fmt : F.formatter)
   (* Add a break before *)
   F.pp_print_break fmt 0 0;
   (* Print a comment to link the extracted type to its original rust definition *)
-  F.pp_print_string fmt ("(** [" ^ Print.name_to_string def.basename ^ "] *)");
+  F.pp_print_string fmt
+    ("(** [" ^ Print.fun_name_to_string def.basename ^ "] *)");
   F.pp_print_space fmt ();
   (* Open a box for the definition, so that whenever possible it gets printed on
    * one line *)
@@ -1232,7 +1234,7 @@ let extract_unit_test_if_unit_fun (ctx : extraction_ctx) (fmt : F.formatter)
     F.pp_print_break fmt 0 0;
     (* Print a comment *)
     F.pp_print_string fmt
-      ("(** Unit test for [" ^ Print.name_to_string def.basename ^ "] *)");
+      ("(** Unit test for [" ^ Print.fun_name_to_string def.basename ^ "] *)");
     F.pp_print_space fmt ();
     (* Open a box for the test *)
     F.pp_open_hovbox fmt ctx.indent_incr;
diff --git a/src/Identifiers.ml b/src/Identifiers.ml
index 825b4ad9..31460dcd 100644
--- a/src/Identifiers.ml
+++ b/src/Identifiers.ml
@@ -176,6 +176,7 @@ type name = string list [@@deriving show, ord]
 (** A name such as: `std::collections::vector` (which would be represented as
     [["std"; "collections"; "vector"]]) *)
 
+(* TODO: remove? *)
 module NameOrderedType : C.OrderedType with type t = name = struct
   type t = name
 
@@ -190,3 +191,21 @@ end
 
 module NameMap = C.MakeMap (NameOrderedType)
 module NameSet = C.MakeSet (NameOrderedType)
+
+type module_name = name [@@deriving show, ord]
+
+type type_name = name [@@deriving show, ord]
+
+module ImplId = IdGen ()
+
+(** A function name *)
+type fun_name =
+  | Regular of name  (** "Regular" function name *)
+  | Impl of type_name * ImplId.id * string
+      (** The function comes from an "impl" block.
+
+          As we may have several "impl" blocks for one type, we need to use
+          a block id to disambiguate the functions (in rustc, this identifier
+          is called a "disambiguator").
+       *)
+[@@deriving show, ord]
diff --git a/src/Interpreter.ml b/src/Interpreter.ml
index 120faeda..91ad6843 100644
--- a/src/Interpreter.ml
+++ b/src/Interpreter.ml
@@ -193,7 +193,7 @@ let evaluate_function_symbolic (config : C.partial_config) (synthesize : bool)
     V.symbolic_value list * SA.expression option =
   (* Debug *)
   let name_to_string () =
-    Print.name_to_string fdef.A.name
+    Print.fun_name_to_string fdef.A.name
     ^ " ("
     ^ Print.option_to_string T.RegionGroupId.to_string back_id
     ^ ")"
@@ -258,7 +258,7 @@ module Test = struct
 
     (* Debug *)
     log#ldebug
-      (lazy ("test_unit_function: " ^ Print.name_to_string fdef.A.name));
+      (lazy ("test_unit_function: " ^ Print.fun_name_to_string fdef.A.name));
 
     (* Sanity check - *)
     assert (List.length fdef.A.signature.region_params = 0);
@@ -282,7 +282,7 @@ module Test = struct
       | _ ->
           failwith
             ("Unit test failed (concrete execution) on: "
-            ^ Print.name_to_string fdef.A.name)
+            ^ Print.fun_name_to_string fdef.A.name)
     in
 
     (* Evaluate the function *)
@@ -312,7 +312,7 @@ module Test = struct
       (fdef : A.fun_def) : unit =
     (* Debug *)
     log#ldebug
-      (lazy ("test_function_symbolic: " ^ Print.name_to_string fdef.A.name));
+      (lazy ("test_function_symbolic: " ^ Print.fun_name_to_string fdef.A.name));
 
     (* Evaluate *)
     let evaluate =
diff --git a/src/Print.ml b/src/Print.ml
index 27a932be..e64e7d73 100644
--- a/src/Print.ml
+++ b/src/Print.ml
@@ -13,6 +13,12 @@ let option_to_string (to_string : 'a -> string) (x : 'a option) : string =
 
 let name_to_string (name : name) : string = String.concat "::" name
 
+let fun_name_to_string (name : fun_name) : string =
+  match name with
+  | Regular name -> name_to_string name
+  | Impl (type_name, impl_id, ident) ->
+      name_to_string type_name ^ "{" ^ ImplId.to_string impl_id ^ "}::" ^ ident
+
 (** Pretty-printing for types *)
 module Types = struct
   let type_var_to_string (tv : T.type_var) : string = tv.name
@@ -737,7 +743,7 @@ module CfimAst = struct
     in
     let fun_def_id_to_string def_id =
       let def = C.ctx_lookup_fun_def ctx def_id in
-      name_to_string def.name
+      fun_name_to_string def.name
     in
     {
       rvar_to_string = ctx_fmt.PV.rvar_to_string;
@@ -779,7 +785,7 @@ module CfimAst = struct
     let adt_field_to_string = type_ctx_to_adt_field_to_string_fun type_defs in
     let fun_def_id_to_string def_id =
       let def = A.FunDefId.Map.find def_id fun_defs in
-      name_to_string def.name
+      fun_name_to_string def.name
     in
     {
       rvar_to_string;
@@ -1037,7 +1043,7 @@ module CfimAst = struct
     let sg = def.signature in
 
     (* Function name *)
-    let name = name_to_string def.A.name in
+    let name = fun_name_to_string def.A.name in
 
     (* Region/type parameters *)
     let regions = sg.region_params in
@@ -1124,7 +1130,7 @@ module Module = struct
     in
     let fun_def_id_to_string def_id =
       let def = A.FunDefId.Map.find def_id fun_context in
-      name_to_string def.name
+      fun_name_to_string def.name
     in
     let var_id_to_string vid =
       let var = V.VarId.nth def.locals vid in
diff --git a/src/PrintPure.ml b/src/PrintPure.ml
index 0b5d0a93..c1da7610 100644
--- a/src/PrintPure.ml
+++ b/src/PrintPure.ml
@@ -61,6 +61,8 @@ let ast_to_type_formatter (fmt : ast_formatter) : type_formatter =
 
 let name_to_string = Print.name_to_string
 
+let fun_name_to_string = Print.fun_name_to_string
+
 let option_to_string = Print.option_to_string
 
 let type_var_to_string = Print.Types.type_var_to_string
@@ -113,7 +115,7 @@ let mk_ast_formatter (type_defs : T.type_def TypeDefId.Map.t)
   in
   let fun_def_id_to_string def_id =
     let def = A.FunDefId.Map.find def_id fun_defs in
-    name_to_string def.name
+    fun_name_to_string def.name
   in
   {
     type_var_id_to_string;
@@ -474,7 +476,7 @@ and switch_to_string (fmt : ast_formatter) (indent : string)
 
 let fun_def_to_string (fmt : ast_formatter) (def : fun_def) : string =
   let type_fmt = ast_to_type_formatter fmt in
-  let name = name_to_string def.basename ^ fun_suffix def.back_id in
+  let name = fun_name_to_string def.basename ^ fun_suffix def.back_id in
   let signature = fun_sig_to_string fmt def.signature in
   let inputs = List.map (var_to_string type_fmt) def.inputs in
   let inputs =
diff --git a/src/Pure.ml b/src/Pure.ml
index bdcb35f8..0f4280c3 100644
--- a/src/Pure.ml
+++ b/src/Pure.ml
@@ -598,7 +598,7 @@ type inst_fun_sig = { inputs : ty list; outputs : ty list }
 type fun_def = {
   def_id : FunDefId.id;
   back_id : T.RegionGroupId.id option;
-  basename : name;
+  basename : fun_name;
       (** The "base" name of the function.
   
           The base name is the original name of the Rust function. We add suffixes
diff --git a/src/PureMicroPasses.ml b/src/PureMicroPasses.ml
index 6c61b9dc..f45c59bb 100644
--- a/src/PureMicroPasses.ml
+++ b/src/PureMicroPasses.ml
@@ -1072,7 +1072,7 @@ let apply_passes_to_def (config : config) (ctx : trans_ctx) (def : fun_def) :
   log#ldebug
     (lazy
       ("PureMicroPasses.apply_passes_to_def: "
-      ^ Print.name_to_string def.basename
+      ^ Print.fun_name_to_string def.basename
       ^ " ("
       ^ Print.option_to_string T.RegionGroupId.to_string def.back_id
       ^ ")"));
diff --git a/src/PureToExtract.ml b/src/PureToExtract.ml
index 21212cd0..ba638743 100644
--- a/src/PureToExtract.ml
+++ b/src/PureToExtract.ml
@@ -30,6 +30,10 @@ module StringMap = Collections.MakeMap (Collections.OrderedString)
 
 type name = Identifiers.name
 
+type type_name = Identifiers.type_name
+
+type fun_name = Identifiers.fun_name
+
 (* TODO: this should a module we give to a functor! *)
 type formatter = {
   bool_name : string;
@@ -65,9 +69,15 @@ type formatter = {
           Inputs:
           - type name
        *)
-  type_name : name -> string;  (** Provided a basename, compute a type name. *)
+  type_name : type_name -> string;
+      (** Provided a basename, compute a type name. *)
   fun_name :
-    A.fun_id -> name -> int -> region_group_info option -> bool * int -> string;
+    A.fun_id ->
+    fun_name ->
+    int ->
+    region_group_info option ->
+    bool * int ->
+    string;
       (** Inputs:
           - function id: this is especially useful to identify whether the
             function is an assumed function or a local function
@@ -82,7 +92,7 @@ type formatter = {
             (`None` if forward function)
           TODO: use the fun id for the assumed functions.
        *)
-  decreases_clause_name : FunDefId.id -> name -> string;
+  decreases_clause_name : FunDefId.id -> fun_name -> string;
       (** Generates the name of the definition used to prove/reason about
           termination. The generated code uses this clause where needed,
           but its body must be defined by the user.
@@ -346,7 +356,7 @@ let id_to_string (id : id) (ctx : extraction_ctx) : string =
       let fun_name =
         match fid with
         | A.Local fid ->
-            Print.name_to_string (FunDefId.Map.find fid fun_defs).name
+            Print.fun_name_to_string (FunDefId.Map.find fid fun_defs).name
         | A.Assumed aid -> A.show_assumed_fun_id aid
       in
       let fun_kind =
@@ -359,7 +369,7 @@ let id_to_string (id : id) (ctx : extraction_ctx) : string =
       let fun_name =
         match fid with
         | A.Local fid ->
-            Print.name_to_string (FunDefId.Map.find fid fun_defs).name
+            Print.fun_name_to_string (FunDefId.Map.find fid fun_defs).name
         | A.Assumed aid -> A.show_assumed_fun_id aid
       in
       "decreases clause for function: " ^ fun_name
diff --git a/src/SymbolicToPure.ml b/src/SymbolicToPure.ml
index 5709ce23..7f747a74 100644
--- a/src/SymbolicToPure.ml
+++ b/src/SymbolicToPure.ml
@@ -1448,7 +1448,7 @@ let translate_fun_def (config : config) (ctx : bs_ctx) (body : S.expression) :
   log#ldebug
     (lazy
       ("SymbolicToPure.translate_fun_def: "
-      ^ Print.name_to_string def.A.name
+      ^ Print.fun_name_to_string def.A.name
       ^ " ("
       ^ Print.option_to_string T.RegionGroupId.to_string bid
       ^ ")"));
diff --git a/src/Translate.ml b/src/Translate.ml
index 077cc32d..9b651288 100644
--- a/src/Translate.ml
+++ b/src/Translate.ml
@@ -56,7 +56,8 @@ let translate_function_to_symbolics (config : C.partial_config)
   (* Debug *)
   log#ldebug
     (lazy
-      ("translate_function_to_symbolics: " ^ Print.name_to_string fdef.A.name));
+      ("translate_function_to_symbolics: "
+      ^ Print.fun_name_to_string fdef.A.name));
 
   let { type_context; fun_context } = trans_ctx in
 
@@ -94,7 +95,8 @@ let translate_function_to_pure (config : C.partial_config)
     pure_fun_translation =
   (* Debug *)
   log#ldebug
-    (lazy ("translate_function_to_pure: " ^ Print.name_to_string fdef.A.name));
+    (lazy
+      ("translate_function_to_pure: " ^ Print.fun_name_to_string fdef.A.name));
 
   let { type_context; fun_context } = trans_ctx in
   let def_id = fdef.def_id in
diff --git a/src/TranslateCore.ml b/src/TranslateCore.ml
index e1a24e10..f74b552b 100644
--- a/src/TranslateCore.ml
+++ b/src/TranslateCore.ml
@@ -41,4 +41,5 @@ let fun_def_to_string (ctx : trans_ctx) (def : Pure.fun_def) : string =
   PrintPure.fun_def_to_string fmt def
 
 let fun_def_id_to_string (ctx : trans_ctx) (id : Pure.FunDefId.id) : string =
-  Print.name_to_string (Pure.FunDefId.Map.find id ctx.fun_context.fun_defs).name
+  Print.fun_name_to_string
+    (Pure.FunDefId.Map.find id ctx.fun_context.fun_defs).name
diff --git a/src/Types.ml b/src/Types.ml
index 9af97338..b099e7a5 100644
--- a/src/Types.ml
+++ b/src/Types.ml
@@ -203,7 +203,7 @@ type type_def_kind = Struct of field list | Enum of variant list
 
 type type_def = {
   def_id : TypeDefId.id;
-  name : name;
+  name : type_name;
   region_params : region_var list;
   type_params : type_var list;
   kind : type_def_kind;