Update the extraction

compiler/AssociatedTypes.ml

@@ -82,10 +82,11 @@ let compute_norm_trait_types_from_preds (span : Meta.span option)
   in
   TraitTypeRefMap.of_list rbindings
 
-let ctx_add_norm_trait_types_from_preds (span : Meta.span) (ctx : eval_ctx)
-    (trait_type_constraints : trait_type_constraint list) : eval_ctx =
+let ctx_add_norm_trait_types_from_preds (span : Meta.span option)
+    (ctx : eval_ctx) (trait_type_constraints : trait_type_constraint list) :
+    eval_ctx =
   let norm_trait_types =
-    compute_norm_trait_types_from_preds (Some span) trait_type_constraints
+    compute_norm_trait_types_from_preds span trait_type_constraints
   in
   { ctx with norm_trait_types }
 
@@ -417,9 +418,9 @@ let norm_ctx_normalize_trait_type_constraint (ctx : norm_ctx)
   let ty = norm_ctx_normalize_ty ctx ty in
   { trait_ref; type_name; ty }
 
-let mk_norm_ctx (span : Meta.span) (ctx : eval_ctx) : norm_ctx =
+let mk_norm_ctx (span : Meta.span option) (ctx : eval_ctx) : norm_ctx =
   {
-    span = Some span;
+    span;
     norm_trait_types = ctx.norm_trait_types;
     type_decls = ctx.type_ctx.type_decls;
     fun_decls = ctx.fun_ctx.fun_decls;
@@ -430,17 +431,17 @@ let mk_norm_ctx (span : Meta.span) (ctx : eval_ctx) : norm_ctx =
     const_generic_vars = ctx.const_generic_vars;
   }
 
-let ctx_normalize_ty (span : Meta.span) (ctx : eval_ctx) (ty : ty) : ty =
+let ctx_normalize_ty (span : Meta.span option) (ctx : eval_ctx) (ty : ty) : ty =
   norm_ctx_normalize_ty (mk_norm_ctx span ctx) ty
 
 (** Normalize a type and erase the regions at the same time *)
 let ctx_normalize_erase_ty (span : Meta.span) (ctx : eval_ctx) (ty : ty) : ty =
-  let ty = ctx_normalize_ty span ctx ty in
+  let ty = ctx_normalize_ty (Some span) ctx ty in
   Subst.erase_regions ty
 
 let ctx_normalize_trait_type_constraint (span : Meta.span) (ctx : eval_ctx)
     (ttc : trait_type_constraint) : trait_type_constraint =
-  norm_ctx_normalize_trait_type_constraint (mk_norm_ctx span ctx) ttc
+  norm_ctx_normalize_trait_type_constraint (mk_norm_ctx (Some span) ctx) ttc
 
 (** Same as [type_decl_get_instantiated_variants_fields_types] but normalizes the types *)
 let type_decl_get_inst_norm_variants_fields_rtypes (span : Meta.span)
@@ -451,7 +452,7 @@ let type_decl_get_inst_norm_variants_fields_rtypes (span : Meta.span)
   in
   List.map
     (fun (variant_id, types) ->
-      (variant_id, List.map (ctx_normalize_ty span ctx) types))
+      (variant_id, List.map (ctx_normalize_ty (Some span) ctx) types))
     res
 
 (** Same as [type_decl_get_instantiated_field_types] but normalizes the types *)
@@ -461,15 +462,15 @@ let type_decl_get_inst_norm_field_rtypes (span : Meta.span) (ctx : eval_ctx)
   let types =
     Subst.type_decl_get_instantiated_field_types def opt_variant_id generics
   in
-  List.map (ctx_normalize_ty span ctx) types
+  List.map (ctx_normalize_ty (Some span) ctx) types
 
 (** Same as [ctx_adt_value_get_instantiated_field_rtypes] but normalizes the types *)
 let ctx_adt_value_get_inst_norm_field_rtypes (span : Meta.span) (ctx : eval_ctx)
     (adt : adt_value) (id : type_id) (generics : generic_args) : ty list =
   let types =
     Subst.ctx_adt_value_get_instantiated_field_types span ctx adt id generics
   in
-  List.map (ctx_normalize_ty span ctx) types
+  List.map (ctx_normalize_ty (Some span) ctx) types
 
 (** Same as [ctx_adt_value_get_instantiated_field_types] but normalizes the types
     and erases the regions. *)
@@ -479,7 +480,7 @@ let type_decl_get_inst_norm_field_etypes (span : Meta.span) (ctx : eval_ctx)
   let types =
     Subst.type_decl_get_instantiated_field_types def opt_variant_id generics
   in
-  let types = List.map (ctx_normalize_ty span ctx) types in
+  let types = List.map (ctx_normalize_ty (Some span) ctx) types in
   List.map Subst.erase_regions types
 
 (** Same as [ctx_adt_get_instantiated_field_types] but normalizes the types and
@@ -491,7 +492,7 @@ let ctx_adt_get_inst_norm_field_etypes (span : Meta.span) (ctx : eval_ctx)
     Subst.ctx_adt_get_instantiated_field_types ctx def_id opt_variant_id
       generics
   in
-  let types = List.map (ctx_normalize_ty span ctx) types in
+  let types = List.map (ctx_normalize_ty (Some span) ctx) types in
   List.map Subst.erase_regions types
 
 (** Same as [substitute_signature] but normalizes the types *)
@@ -507,8 +508,8 @@ let ctx_subst_norm_signature (span : Meta.span) (ctx : eval_ctx)
       sg regions_hierarchy
   in
   let { regions_hierarchy; inputs; output; trait_type_constraints } = sg in
-  let inputs = List.map (ctx_normalize_ty span ctx) inputs in
-  let output = ctx_normalize_ty span ctx output in
+  let inputs = List.map (ctx_normalize_ty (Some span) ctx) inputs in
+  let output = ctx_normalize_ty (Some span) ctx output in
   let trait_type_constraints =
     List.map
       (ctx_normalize_trait_type_constraint span ctx)

compiler/Assumed.ml

@@ -215,10 +215,11 @@ type assumed_fun_info = {
        *)
 }
 
-let mk_assumed_fun_info (raw : raw_assumed_fun_info) : assumed_fun_info =
+let mk_assumed_fun_info (raw : raw_assumed_fun_info) :
+    assumed_fun_id * assumed_fun_info =
   let fun_id, fun_sig, can_fail, keep_types = raw in
   let name = Charon.PrintExpressions.assumed_fun_id_to_string fun_id in
-  { fun_id; fun_sig; can_fail; name; keep_types }
+  (fun_id, { fun_id; fun_sig; can_fail; name; keep_types })
 
 (** The list of assumed functions and all their information:
     - their signature
@@ -260,11 +261,23 @@ let raw_assumed_fun_infos : raw_assumed_fun_info list =
       None );
   ]
 
-let assumed_fun_infos : assumed_fun_info list =
-  List.map mk_assumed_fun_info raw_assumed_fun_infos
+module OrderedAssumedFunId :
+  Collections.OrderedType with type t = assumed_fun_id = struct
+  type t = assumed_fun_id
+
+  let compare x y = compare_assumed_fun_id x y
+  let to_string x = show_assumed_fun_id x
+  let pp_t fmt x = Format.pp_print_string fmt (show_assumed_fun_id x)
+  let show_t x = show_assumed_fun_id x
+end
+
+module AssumedFunIdMap = Collections.MakeMap (OrderedAssumedFunId)
+
+let assumed_fun_infos : assumed_fun_info AssumedFunIdMap.t =
+  AssumedFunIdMap.of_list (List.map mk_assumed_fun_info raw_assumed_fun_infos)
 
 let get_assumed_fun_info (id : assumed_fun_id) : assumed_fun_info =
-  match List.find_opt (fun x -> id = x.fun_id) assumed_fun_infos with
+  match AssumedFunIdMap.find_opt id assumed_fun_infos with
   | Some info -> info
   | None ->
       raise