Normalize the function signatures before translation to pure

9 files changed, 294 insertions, 150 deletions

diff --git a/compiler/AssociatedTypes.ml b/compiler/AssociatedTypes.ml
index 92cd464e..992dade9 100644
--- a/compiler/AssociatedTypes.ml
+++ b/compiler/AssociatedTypes.ml
@@ -26,6 +26,7 @@ let trait_type_ref_substitute (subst : ('r, 'r1) Subst.subst)
   let trait_ref = Subst.trait_ref_substitute subst trait_ref in
   { C.trait_ref; type_name }
 
+(* TODO: how not to duplicate below? *)
 module RTyOrd = struct
   type t = T.rty
 
@@ -35,58 +36,114 @@ module RTyOrd = struct
   let show_t = T.show_rty
 end
 
+module STyOrd = struct
+  type t = T.sty
+
+  let compare = T.compare_sty
+  let to_string = T.show_sty
+  let pp_t = T.pp_sty
+  let show_t = T.show_sty
+end
+
 module RTyMap = Collections.MakeMap (RTyOrd)
+module STyMap = Collections.MakeMap (STyOrd)
+
+(* TODO: is it possible not to have this? *)
+module type TypeWrapper = sig
+  type t
+end
+
+(* TODO: don't manage to get the syntax right so using a functor *)
+module MakeNormalizer
+    (R : TypeWrapper)
+    (RTyMap : Collections.Map with type key = R.t T.region T.ty)
+    (M : Collections.Map with type key = R.t T.region C.trait_type_ref) =
+struct
+  let compute_norm_trait_types_from_preds
+      (trait_type_constraints : R.t T.region T.trait_type_constraint list) :
+      R.t T.region T.ty M.t =
+    (* Compute a union-find structure by recursively exploring the predicates and clauses *)
+    let norm : R.t T.region T.ty UF.elem RTyMap.t ref = ref RTyMap.empty in
+    let get_ref (ty : R.t T.region T.ty) : R.t T.region T.ty UF.elem =
+      match RTyMap.find_opt ty !norm with
+      | Some r -> r
+      | None ->
+          let r = UF.make ty in
+          norm := RTyMap.add ty r !norm;
+          r
+    in
+    let add_trait_type_constraint (c : R.t T.region T.trait_type_constraint) =
+      let trait_ty = T.TraitType (c.trait_ref, c.generics, c.type_name) in
+      let trait_ty_ref = get_ref trait_ty in
+      let ty_ref = get_ref c.ty in
+      let new_repr = UF.get ty_ref in
+      let merged = UF.union trait_ty_ref ty_ref in
+      (* Not sure the set operation is necessary, but I want to control which
+         representative is chosen *)
+      UF.set merged new_repr
+    in
+    (* Explore the local predicates *)
+    List.iter add_trait_type_constraint trait_type_constraints;
+    (* TODO: explore the local clauses *)
+    (* Compute the norm maps *)
+    let rbindings =
+      List.map (fun (k, v) -> (k, UF.get v)) (RTyMap.bindings !norm)
+    in
+    (* Filter the keys to keep only the trait type aliases *)
+    let rbindings =
+      List.filter_map
+        (fun (k, v) ->
+          match k with
+          | T.TraitType (trait_ref, generics, type_name) ->
+              assert (generics = TypesUtils.mk_empty_generic_args);
+              Some ({ C.trait_ref; type_name }, v)
+          | _ -> None)
+        rbindings
+    in
+    M.of_list rbindings
+end
+
+(** Compute the representative classes of trait associated types, for normalization *)
+let compute_norm_trait_stypes_from_preds
+    (trait_type_constraints : T.strait_type_constraint list) :
+    T.sty C.STraitTypeRefMap.t =
+  (* Compute the normalization map for the types with regions *)
+  let module R = struct
+    type t = T.region_var_id
+  end in
+  let module M = C.STraitTypeRefMap in
+  let module Norm = MakeNormalizer (R) (STyMap) (M) in
+  Norm.compute_norm_trait_types_from_preds trait_type_constraints
 
 (** Compute the representative classes of trait associated types, for normalization *)
 let compute_norm_trait_types_from_preds
     (trait_type_constraints : T.rtrait_type_constraint list) :
     T.ety C.ETraitTypeRefMap.t * T.rty C.RTraitTypeRefMap.t =
-  (* Compute a union-find structure by recursively exploring the predicates and clauses *)
-  let norm : T.rty UF.elem RTyMap.t ref = ref RTyMap.empty in
-  let get_ref (ty : T.rty) : T.rty UF.elem =
-    match RTyMap.find_opt ty !norm with
-    | Some r -> r
-    | None ->
-        let r = UF.make ty in
-        norm := RTyMap.add ty r !norm;
-        r
-  in
-  let add_trait_type_constraint (c : T.rtrait_type_constraint) =
-    let trait_ty = T.TraitType (c.trait_ref, c.generics, c.type_name) in
-    let trait_ty_ref = get_ref trait_ty in
-    let ty_ref = get_ref c.ty in
-    let new_repr = UF.get ty_ref in
-    let merged = UF.union trait_ty_ref ty_ref in
-    (* Not sure the set operation is necessary, but I want to control which
-       representative is chosen *)
-    UF.set merged new_repr
-  in
-  (* Explore the local predicates *)
-  List.iter add_trait_type_constraint trait_type_constraints;
-  (* TODO: explore the local clauses *)
-  (* Compute the norm maps *)
+  (* Compute the normalization map for the types with regions *)
+  let module R = struct
+    type t = T.region_id
+  end in
+  let module M = C.RTraitTypeRefMap in
+  let module Norm = MakeNormalizer (R) (RTyMap) (M) in
   let rbindings =
-    List.map (fun (k, v) -> (k, UF.get v)) (RTyMap.bindings !norm)
-  in
-  (* Filter the keys to keep only the trait type aliases *)
-  let rbindings =
-    List.filter_map
-      (fun (k, v) ->
-        match k with
-        | T.TraitType (trait_ref, generics, type_name) ->
-            assert (generics = TypesUtils.mk_empty_generic_args);
-            Some ({ C.trait_ref; type_name }, v)
-        | _ -> None)
-      rbindings
+    Norm.compute_norm_trait_types_from_preds trait_type_constraints
   in
+  (* Compute the normalization map for the types with erased regions *)
   let ebindings =
     List.map
       (fun (k, v) ->
         ( trait_type_ref_substitute Subst.erase_regions_subst k,
           Subst.erase_regions v ))
-      rbindings
+      (M.bindings rbindings)
   in
-  (C.ETraitTypeRefMap.of_list ebindings, C.RTraitTypeRefMap.of_list rbindings)
+  (C.ETraitTypeRefMap.of_list ebindings, rbindings)
+
+let ctx_add_norm_trait_stypes_from_preds (ctx : C.eval_ctx)
+    (trait_type_constraints : T.strait_type_constraint list) : C.eval_ctx =
+  let norm_trait_stypes =
+    compute_norm_trait_stypes_from_preds trait_type_constraints
+  in
+  { ctx with C.norm_trait_stypes }
 
 let ctx_add_norm_trait_types_from_preds (ctx : C.eval_ctx)
     (trait_type_constraints : T.rtrait_type_constraint list) : C.eval_ctx =
@@ -398,6 +455,19 @@ let ctx_normalize_trait_type_constraint (ctx : 'r norm_ctx)
   let ty = ctx_normalize_ty ctx ty in
   { T.trait_ref; generics; type_name; ty }
 
+let mk_snorm_ctx (ctx : C.eval_ctx) : T.RegionVarId.id T.region norm_ctx =
+  let get_ty_repr x = C.STraitTypeRefMap.find_opt x ctx.norm_trait_stypes in
+  {
+    ctx;
+    get_ty_repr;
+    convert_ety = TypesUtils.ety_no_regions_to_sty;
+    convert_etrait_ref = TypesUtils.etrait_ref_no_regions_to_gr_trait_ref;
+    ty_to_string = PA.sty_to_string ctx;
+    trait_ref_to_string = PA.strait_ref_to_string ctx;
+    trait_instance_id_to_string = PA.strait_instance_id_to_string ctx;
+    pp_r = T.pp_region T.pp_region_var_id;
+  }
+
 let mk_rnorm_ctx (ctx : C.eval_ctx) : T.RegionId.id T.region norm_ctx =
   let get_ty_repr x = C.RTraitTypeRefMap.find_opt x ctx.norm_trait_rtypes in
   {
@@ -424,6 +494,9 @@ let mk_enorm_ctx (ctx : C.eval_ctx) : T.erased_region norm_ctx =
     pp_r = T.pp_erased_region;
   }
 
+let ctx_normalize_sty (ctx : C.eval_ctx) (ty : T.sty) : T.sty =
+  ctx_normalize_ty (mk_snorm_ctx ctx) ty
+
 let ctx_normalize_rty (ctx : C.eval_ctx) (ty : T.rty) : T.rty =
   ctx_normalize_ty (mk_rnorm_ctx ctx) ty
 
diff --git a/compiler/Contexts.ml b/compiler/Contexts.ml
index a5bc7dc0..dac64a9a 100644
--- a/compiler/Contexts.ml
+++ b/compiler/Contexts.ml
@@ -288,6 +288,9 @@ type etrait_type_ref = erased_region trait_type_ref [@@deriving show, ord]
 type rtrait_type_ref = Types.RegionId.id Types.region trait_type_ref
 [@@deriving show, ord]
 
+type strait_type_ref = Types.RegionVarId.id Types.region trait_type_ref
+[@@deriving show, ord]
+
 (* TODO: correctly use the functors so as not to have a duplication below *)
 module ETraitTypeRefOrd = struct
   type t = etrait_type_ref
@@ -307,8 +310,18 @@ module RTraitTypeRefOrd = struct
   let show_t = show_rtrait_type_ref
 end
 
+module STraitTypeRefOrd = struct
+  type t = strait_type_ref
+
+  let compare = compare_strait_type_ref
+  let to_string = show_strait_type_ref
+  let pp_t = pp_strait_type_ref
+  let show_t = show_strait_type_ref
+end
+
 module ETraitTypeRefMap = Collections.MakeMap (ETraitTypeRefOrd)
 module RTraitTypeRefMap = Collections.MakeMap (RTraitTypeRefOrd)
+module STraitTypeRefMap = Collections.MakeMap (STraitTypeRefOrd)
 
 (** Evaluation context *)
 type eval_ctx = {
@@ -336,6 +349,13 @@ type eval_ctx = {
 
           TODO: how not to duplicate?
        *)
+  norm_trait_stypes : sty STraitTypeRefMap.t;
+      (** We sometimes need to normalize types in non-instantiated signatures.
+
+          Note that we either need to use the etypes/rtypes maps, or the stypes map.
+          This means that we either compute the maps for etypes and rtypes, or compute
+          the one for stypes (we don't always compute and carry all the maps).
+       *)
   env : env;
   ended_regions : RegionId.Set.t;
 }
diff --git a/compiler/Interpreter.ml b/compiler/Interpreter.ml
index 752d6f2f..09f25ca1 100644
--- a/compiler/Interpreter.ml
+++ b/compiler/Interpreter.ml
@@ -35,38 +35,6 @@ let compute_contexts (m : A.crate) : C.decls_ctx =
   let trait_impls_ctx = { C.trait_impls } in
   { C.type_ctx; fun_ctx; global_ctx; trait_decls_ctx; trait_impls_ctx }
 
-(** **WARNING**: this function doesn't compute the normalized types
-    (for the trait type aliases). This should be computed afterwards.
- *)
-let initialize_eval_context (ctx : C.decls_ctx)
-    (region_groups : T.RegionGroupId.id list) (type_vars : T.type_var list)
-    (const_generic_vars : T.const_generic_var list) : C.eval_ctx =
-  C.reset_global_counters ();
-  let const_generic_vars_map =
-    T.ConstGenericVarId.Map.of_list
-      (List.map
-         (fun (cg : T.const_generic_var) ->
-           let ty = TypesUtils.ety_no_regions_to_rty (T.Literal cg.ty) in
-           let cv = mk_fresh_symbolic_typed_value V.ConstGeneric ty in
-           (cg.index, cv))
-         const_generic_vars)
-  in
-  {
-    C.type_context = ctx.type_ctx;
-    C.fun_context = ctx.fun_ctx;
-    C.global_context = ctx.global_ctx;
-    C.trait_decls_context = ctx.trait_decls_ctx;
-    C.trait_impls_context = ctx.trait_impls_ctx;
-    C.region_groups;
-    C.type_vars;
-    C.const_generic_vars;
-    C.const_generic_vars_map;
-    C.norm_trait_etypes = C.ETraitTypeRefMap.empty (* Empty for now *);
-    C.norm_trait_rtypes = C.RTraitTypeRefMap.empty (* Empty for now *);
-    C.env = [ C.Frame ];
-    C.ended_regions = T.RegionId.Set.empty;
-  }
-
 (** Small helper.
 
     Normalize an instantiated function signature provided we used this signature
@@ -93,11 +61,10 @@ let normalize_inst_fun_sig (ctx : C.eval_ctx) (sg : A.inst_fun_sig) :
     clauses (we are not considering a function call, so we don't need to
     normalize because a trait clause was instantiated with a specific trait ref).
  *)
-let symbolic_instantiate_fun_sig (ctx : C.eval_ctx) (fdef : A.fun_decl) :
-    C.eval_ctx * A.inst_fun_sig =
-  let sg = fdef.signature in
+let symbolic_instantiate_fun_sig (ctx : C.eval_ctx) (sg : A.fun_sig)
+    (kind : A.fun_kind) : C.eval_ctx * A.inst_fun_sig =
   let tr_self =
-    match fdef.kind with
+    match kind with
     | RegularKind | TraitMethodImpl _ -> T.UnknownTrait __FUNCTION__
     | TraitMethodDecl _ | TraitMethodProvided _ -> T.Self
   in
@@ -185,7 +152,9 @@ let initialize_symbolic_context_for_fun (ctx : C.decls_ctx) (fdef : A.fun_decl)
   (* Instantiate the signature. This updates the context because we compute
      at the same time the normalization map for the associated types.
   *)
-  let ctx, inst_sg = symbolic_instantiate_fun_sig ctx fdef in
+  let ctx, inst_sg =
+    symbolic_instantiate_fun_sig ctx fdef.signature fdef.kind
+  in
   (* Create fresh symbolic values for the inputs *)
   let input_svs =
     List.map (fun ty -> mk_fresh_symbolic_value V.SynthInput ty) inst_sg.inputs
@@ -260,7 +229,9 @@ let evaluate_function_symbolic_synthesize_backward_from_return
    * an instantiation of the signature, so that we use fresh
    * region ids for the return abstractions. *)
   let sg = fdef.signature in
-  let _, ret_inst_sg = symbolic_instantiate_fun_sig ctx fdef in
+  let _, ret_inst_sg =
+    symbolic_instantiate_fun_sig ctx fdef.signature fdef.kind
+  in
   let ret_rty = ret_inst_sg.output in
   (* Move the return value out of the return variable *)
   let pop_return_value = is_regular_return in
diff --git a/compiler/InterpreterLoopsJoinCtxs.ml b/compiler/InterpreterLoopsJoinCtxs.ml
index fa44e20e..6d3ecb18 100644
--- a/compiler/InterpreterLoopsJoinCtxs.ml
+++ b/compiler/InterpreterLoopsJoinCtxs.ml
@@ -562,6 +562,7 @@ let join_ctxs (loop_id : V.LoopId.id) (fixed_ids : ids_sets) (ctx0 : C.eval_ctx)
       const_generic_vars_map;
       norm_trait_etypes;
       norm_trait_rtypes;
+      norm_trait_stypes;
       env = _;
       ended_regions = ended_regions0;
     } =
@@ -579,6 +580,7 @@ let join_ctxs (loop_id : V.LoopId.id) (fixed_ids : ids_sets) (ctx0 : C.eval_ctx)
       const_generic_vars_map = _;
       norm_trait_etypes = _;
       norm_trait_rtypes = _;
+      norm_trait_stypes = _;
       env = _;
       ended_regions = ended_regions1;
     } =
@@ -598,6 +600,7 @@ let join_ctxs (loop_id : V.LoopId.id) (fixed_ids : ids_sets) (ctx0 : C.eval_ctx)
         const_generic_vars_map;
         norm_trait_etypes;
         norm_trait_rtypes;
+        norm_trait_stypes;
         env;
         ended_regions;
       }
diff --git a/compiler/InterpreterStatements.ml b/compiler/InterpreterStatements.ml
index f54c5dbd..dc15c6ac 100644
--- a/compiler/InterpreterStatements.ml
+++ b/compiler/InterpreterStatements.ml
@@ -662,67 +662,6 @@ let eval_assumed_function_call_concrete (config : C.config)
       (* Compose and apply *)
       comp cf_eval_ops cf_eval_call
 
-let instantiate_fun_sig (ctx : C.eval_ctx) (generics : T.egeneric_args)
-    (tr_self : T.rtrait_instance_id) (sg : A.fun_sig) : A.inst_fun_sig =
-  log#ldebug
-    (lazy
-      ("instantiate_fun_sig:" ^ "\n- generics: "
-      ^ egeneric_args_to_string ctx generics
-      ^ "\n- tr_self: "
-      ^ rtrait_instance_id_to_string ctx tr_self
-      ^ "\n- sg: " ^ fun_sig_to_string ctx sg));
-  (* Generate fresh abstraction ids and create a substitution from region
-   * group ids to abstraction ids *)
-  let rg_abs_ids_bindings =
-    List.map
-      (fun rg ->
-        let abs_id = C.fresh_abstraction_id () in
-        (rg.T.id, abs_id))
-      sg.regions_hierarchy
-  in
-  let asubst_map : V.AbstractionId.id T.RegionGroupId.Map.t =
-    List.fold_left
-      (fun mp (rg_id, abs_id) -> T.RegionGroupId.Map.add rg_id abs_id mp)
-      T.RegionGroupId.Map.empty rg_abs_ids_bindings
-  in
-  let asubst (rg_id : T.RegionGroupId.id) : V.AbstractionId.id =
-    T.RegionGroupId.Map.find rg_id asubst_map
-  in
-  (* Generate fresh regions and their substitutions *)
-  let _, rsubst, _ = Subst.fresh_regions_with_substs sg.generics.regions in
-  (* Generate the type substitution
-   * Note that we need the substitution to map the type variables to
-   * {!rty} types (not {!ety}). In order to do that, we convert the
-   * type parameters to types with regions. This is possible only
-   * if those types don't contain any regions.
-   * This is a current limitation of the analysis: there is still some
-   * work to do to properly handle full type parametrization.
-   * *)
-  let rtype_params = List.map ety_no_regions_to_rty generics.types in
-  let tsubst = Subst.make_type_subst_from_vars sg.generics.types rtype_params in
-  let cgsubst =
-    Subst.make_const_generic_subst_from_vars sg.generics.const_generics
-      generics.const_generics
-  in
-  (* TODO: something annoying with the trait ref subst: we need to use region
-     types, but the arguments use erased regions. For now we use the fact
-     that no regions should appear inside. In the future: we should merge
-     ety and rty. *)
-  let trait_refs =
-    List.map TypesUtils.etrait_ref_no_regions_to_gr_trait_ref
-      generics.trait_refs
-  in
-  let tr_subst =
-    Subst.make_trait_subst_from_clauses sg.generics.trait_clauses trait_refs
-  in
-  (* Substitute the signature *)
-  let inst_sig =
-    Assoc.ctx_subst_norm_signature ctx asubst rsubst tsubst cgsubst tr_subst
-      tr_self sg
-  in
-  (* Return *)
-  inst_sig
-
 (** Helper
  
     Create abstractions (with no avalues, which have to be inserted afterwards)
diff --git a/compiler/InterpreterStatements.mli b/compiler/InterpreterStatements.mli
index 0a086fb2..e65758ae 100644
--- a/compiler/InterpreterStatements.mli
+++ b/compiler/InterpreterStatements.mli
@@ -25,19 +25,6 @@ open InterpreterExpressions
  *)
 val pop_frame : C.config -> bool -> (V.typed_value option -> m_fun) -> m_fun
 
-(** Instantiate a function signature, introducing **fresh** abstraction ids and
-    region ids. This is mostly used in preparation of function calls, when
-    evaluating in symbolic mode of course.
-
-    Note: there are no region parameters, because they should be erased.
- *)
-val instantiate_fun_sig :
-  C.eval_ctx ->
-  T.egeneric_args ->
-  T.rtrait_instance_id ->
-  LA.fun_sig ->
-  LA.inst_fun_sig
-
 (** Helper.
 
     Create a list of abstractions from a list of regions groups, and insert
diff --git a/compiler/InterpreterUtils.ml b/compiler/InterpreterUtils.ml
index 6fde8d68..7aaee6ff 100644
--- a/compiler/InterpreterUtils.ml
+++ b/compiler/InterpreterUtils.ml
@@ -10,6 +10,11 @@ open TypesUtils
 module PA = Print.EvalCtxLlbcAst
 open Cps
 
+(* TODO: we should probably rename the file to ContextsUtils *)
+
+(** The local logger *)
+let log = L.interpreter_log
+
 (** Some utilities *)
 
 (** Auxiliary function - call a function which requires a continuation,
@@ -413,3 +418,103 @@ let compute_contexts_ids (ctxl : C.eval_ctx list) : ids_sets * ids_to_values =
 (** Compute the sets of ids found in a context. *)
 let compute_context_ids (ctx : C.eval_ctx) : ids_sets * ids_to_values =
   compute_contexts_ids [ ctx ]
+
+(** **WARNING**: this function doesn't compute the normalized types
+    (for the trait type aliases). This should be computed afterwards.
+ *)
+let initialize_eval_context (ctx : C.decls_ctx)
+    (region_groups : T.RegionGroupId.id list) (type_vars : T.type_var list)
+    (const_generic_vars : T.const_generic_var list) : C.eval_ctx =
+  C.reset_global_counters ();
+  let const_generic_vars_map =
+    T.ConstGenericVarId.Map.of_list
+      (List.map
+         (fun (cg : T.const_generic_var) ->
+           let ty = TypesUtils.ety_no_regions_to_rty (T.Literal cg.ty) in
+           let cv = mk_fresh_symbolic_typed_value V.ConstGeneric ty in
+           (cg.index, cv))
+         const_generic_vars)
+  in
+  {
+    C.type_context = ctx.type_ctx;
+    C.fun_context = ctx.fun_ctx;
+    C.global_context = ctx.global_ctx;
+    C.trait_decls_context = ctx.trait_decls_ctx;
+    C.trait_impls_context = ctx.trait_impls_ctx;
+    C.region_groups;
+    C.type_vars;
+    C.const_generic_vars;
+    C.const_generic_vars_map;
+    C.norm_trait_etypes = C.ETraitTypeRefMap.empty (* Empty for now *);
+    C.norm_trait_rtypes = C.RTraitTypeRefMap.empty (* Empty for now *);
+    C.norm_trait_stypes = C.STraitTypeRefMap.empty (* Empty for now *);
+    C.env = [ C.Frame ];
+    C.ended_regions = T.RegionId.Set.empty;
+  }
+
+(** Instantiate a function signature, introducing **fresh** abstraction ids and
+    region ids. This is mostly used in preparation of function calls (when
+    evaluating in symbolic mode).
+
+    Note: there are no region parameters, because they should be erased.
+ *)
+let instantiate_fun_sig (ctx : C.eval_ctx) (generics : T.egeneric_args)
+    (tr_self : T.rtrait_instance_id) (sg : A.fun_sig) : A.inst_fun_sig =
+  log#ldebug
+    (lazy
+      ("instantiate_fun_sig:" ^ "\n- generics: "
+      ^ egeneric_args_to_string ctx generics
+      ^ "\n- tr_self: "
+      ^ rtrait_instance_id_to_string ctx tr_self
+      ^ "\n- sg: " ^ fun_sig_to_string ctx sg));
+  (* Generate fresh abstraction ids and create a substitution from region
+   * group ids to abstraction ids *)
+  let rg_abs_ids_bindings =
+    List.map
+      (fun rg ->
+        let abs_id = C.fresh_abstraction_id () in
+        (rg.T.id, abs_id))
+      sg.regions_hierarchy
+  in
+  let asubst_map : V.AbstractionId.id T.RegionGroupId.Map.t =
+    List.fold_left
+      (fun mp (rg_id, abs_id) -> T.RegionGroupId.Map.add rg_id abs_id mp)
+      T.RegionGroupId.Map.empty rg_abs_ids_bindings
+  in
+  let asubst (rg_id : T.RegionGroupId.id) : V.AbstractionId.id =
+    T.RegionGroupId.Map.find rg_id asubst_map
+  in
+  (* Generate fresh regions and their substitutions *)
+  let _, rsubst, _ = Subst.fresh_regions_with_substs sg.generics.regions in
+  (* Generate the type substitution
+   * Note that we need the substitution to map the type variables to
+   * {!rty} types (not {!ety}). In order to do that, we convert the
+   * type parameters to types with regions. This is possible only
+   * if those types don't contain any regions.
+   * This is a current limitation of the analysis: there is still some
+   * work to do to properly handle full type parametrization.
+   * *)
+  let rtype_params = List.map ety_no_regions_to_rty generics.types in
+  let tsubst = Subst.make_type_subst_from_vars sg.generics.types rtype_params in
+  let cgsubst =
+    Subst.make_const_generic_subst_from_vars sg.generics.const_generics
+      generics.const_generics
+  in
+  (* TODO: something annoying with the trait ref subst: we need to use region
+     types, but the arguments use erased regions. For now we use the fact
+     that no regions should appear inside. In the future: we should merge
+     ety and rty. *)
+  let trait_refs =
+    List.map TypesUtils.etrait_ref_no_regions_to_gr_trait_ref
+      generics.trait_refs
+  in
+  let tr_subst =
+    Subst.make_trait_subst_from_clauses sg.generics.trait_clauses trait_refs
+  in
+  (* Substitute the signature *)
+  let inst_sig =
+    AssociatedTypes.ctx_subst_norm_signature ctx asubst rsubst tsubst cgsubst
+      tr_subst tr_self sg
+  in
+  (* Return *)
+  inst_sig
diff --git a/compiler/Print.ml b/compiler/Print.ml
index 522d9fdd..5d5c16ee 100644
--- a/compiler/Print.ml
+++ b/compiler/Print.ml
@@ -478,6 +478,9 @@ module Contexts = struct
   let ctx_to_rtype_formatter (fmt : ctx_formatter) : PT.rtype_formatter =
     PV.value_to_rtype_formatter fmt
 
+  let ctx_to_stype_formatter (fmt : ctx_formatter) : PT.stype_formatter =
+    PV.value_to_stype_formatter fmt
+
   let eval_ctx_to_ctx_formatter (ctx : C.eval_ctx) : ctx_formatter =
     let rvar_to_string r =
       (* In theory we shouldn't use rvar_to_string, but it can happen
@@ -651,6 +654,11 @@ module EvalCtxLlbcAst = struct
     let fmt = PC.ctx_to_rtype_formatter fmt in
     PT.rty_to_string fmt t
 
+  let sty_to_string (ctx : C.eval_ctx) (t : T.sty) : string =
+    let fmt = PC.eval_ctx_to_ctx_formatter ctx in
+    let fmt = PC.ctx_to_stype_formatter fmt in
+    PT.sty_to_string fmt t
+
   let etrait_ref_to_string (ctx : C.eval_ctx) (x : T.etrait_ref) : string =
     let fmt = PC.eval_ctx_to_ctx_formatter ctx in
     let fmt = PC.ctx_to_etype_formatter fmt in
@@ -661,6 +669,11 @@ module EvalCtxLlbcAst = struct
     let fmt = PC.ctx_to_rtype_formatter fmt in
     PT.rtrait_ref_to_string fmt x
 
+  let strait_ref_to_string (ctx : C.eval_ctx) (x : T.strait_ref) : string =
+    let fmt = PC.eval_ctx_to_ctx_formatter ctx in
+    let fmt = PC.ctx_to_stype_formatter fmt in
+    PT.strait_ref_to_string fmt x
+
   let etrait_instance_id_to_string (ctx : C.eval_ctx) (x : T.etrait_instance_id)
       : string =
     let fmt = PC.eval_ctx_to_ctx_formatter ctx in
@@ -673,6 +686,12 @@ module EvalCtxLlbcAst = struct
     let fmt = PC.ctx_to_rtype_formatter fmt in
     PT.rtrait_instance_id_to_string fmt x
 
+  let strait_instance_id_to_string (ctx : C.eval_ctx) (x : T.strait_instance_id)
+      : string =
+    let fmt = PC.eval_ctx_to_ctx_formatter ctx in
+    let fmt = PC.ctx_to_stype_formatter fmt in
+    PT.strait_instance_id_to_string fmt x
+
   let egeneric_args_to_string (ctx : C.eval_ctx) (x : T.egeneric_args) : string
       =
     let fmt = PC.eval_ctx_to_ctx_formatter ctx in
diff --git a/compiler/SymbolicToPure.ml b/compiler/SymbolicToPure.ml
index be9b7261..429198ad 100644
--- a/compiler/SymbolicToPure.ml
+++ b/compiler/SymbolicToPure.ml
@@ -842,6 +842,33 @@ let translate_fun_sig (decls_ctx : C.decls_ctx) (fun_id : A.fun_id)
   (* Is the function stateful, and can it fail? *)
   let lid = None in
   let effect_info = get_fun_effect_info fun_infos (A.FunId fun_id) lid bid in
+  (* We need an evaluation context to normalize the types (to normalize the
+     associated types, etc. - for instance it may happen that the types
+     refer to the types associated to a trait ref, but where the trait ref
+     is a known impl). *)
+  (* Create the context *)
+  let ctx =
+    let region_groups =
+      List.map (fun (g : T.region_var_group) -> g.id) sg.regions_hierarchy
+    in
+    let ctx =
+      InterpreterUtils.initialize_eval_context decls_ctx region_groups
+        sg.generics.types sg.generics.const_generics
+    in
+    (* Compute the normalization map for the *sty* types and add it to the context *)
+    AssociatedTypes.ctx_add_norm_trait_stypes_from_preds ctx
+      sg.preds.trait_type_constraints
+  in
+
+  (* Normalize the signature *)
+  let sg =
+    let ({ A.inputs; output; _ } : A.fun_sig) = sg in
+    let norm = AssociatedTypes.ctx_normalize_sty ctx in
+    let inputs = List.map norm inputs in
+    let output = norm output in
+    { sg with A.inputs; output }
+  in
+
   (* List the inputs for:
    * - the fuel
    * - the forward function