Normalize the types when computing the regions hierarchies

4 files changed, 186 insertions, 77 deletions

diff --git a/compiler/AssociatedTypes.ml b/compiler/AssociatedTypes.ml
index 27e08495..2f14f0f2 100644
--- a/compiler/AssociatedTypes.ml
+++ b/compiler/AssociatedTypes.ml
@@ -104,32 +104,91 @@ let rec trait_instance_id_is_local_clause (id : T.trait_instance_id) : bool =
 (** About the conversion functions: for now we need them (TODO: merge ety, rty, etc.),
     but they should be applied to types without regions.
  *)
-type norm_ctx = { ctx : C.eval_ctx }
-
-let ctx_get_ty_repr (ctx : norm_ctx) (x : C.trait_type_ref) : T.ty option =
-  C.TraitTypeRefMap.find_opt x ctx.ctx.norm_trait_types
+type norm_ctx = {
+  norm_trait_types : T.ty C.TraitTypeRefMap.t;
+  type_decls : T.type_decl T.TypeDeclId.Map.t;
+  fun_decls : A.fun_decl A.FunDeclId.Map.t;
+  global_decls : A.global_decl A.GlobalDeclId.Map.t;
+  trait_decls : A.trait_decl A.TraitDeclId.Map.t;
+  trait_impls : A.trait_impl A.TraitImplId.Map.t;
+  type_vars : T.type_var list;
+  const_generic_vars : T.const_generic_var list;
+}
+
+let norm_ctx_to_type_formatter (ctx : norm_ctx) : Print.Types.type_formatter =
+  let open Print in
+  let region_id_to_string r = PT.region_id_to_string r in
+
+  let type_var_id_to_string vid =
+    (* The context may be invalid *)
+    match T.TypeVarId.nth_opt ctx.type_vars vid with
+    | None -> T.TypeVarId.to_string vid
+    | Some v -> v.name
+  in
+  let const_generic_var_id_to_string vid =
+    match T.ConstGenericVarId.nth_opt ctx.const_generic_vars vid with
+    | None -> T.ConstGenericVarId.to_string vid
+    | Some v -> v.name
+  in
+  let type_decl_id_to_string def_id =
+    let def = T.TypeDeclId.Map.find def_id ctx.type_decls in
+    name_to_string def.name
+  in
+  let global_decl_id_to_string def_id =
+    let def = A.GlobalDeclId.Map.find def_id ctx.global_decls in
+    name_to_string def.name
+  in
+  let trait_decl_id_to_string def_id =
+    let def = A.TraitDeclId.Map.find def_id ctx.trait_decls in
+    name_to_string def.name
+  in
+  let trait_impl_id_to_string def_id =
+    let def = A.TraitImplId.Map.find def_id ctx.trait_impls in
+    name_to_string def.name
+  in
+  let trait_clause_id_to_string id = PT.trait_clause_id_to_pretty_string id in
+  {
+    region_id_to_string;
+    type_var_id_to_string;
+    type_decl_id_to_string;
+    const_generic_var_id_to_string;
+    global_decl_id_to_string;
+    trait_decl_id_to_string;
+    trait_impl_id_to_string;
+    trait_clause_id_to_string;
+  }
+
+let norm_ctx_get_ty_repr (ctx : norm_ctx) (x : C.trait_type_ref) : T.ty option =
+  C.TraitTypeRefMap.find_opt x ctx.norm_trait_types
 
 let ty_to_string (ctx : norm_ctx) (ty : T.ty) : string =
-  PA.ty_to_string ctx.ctx ty
+  let ctx = norm_ctx_to_type_formatter ctx in
+  Print.Types.ty_to_string ctx ty
 
 let trait_ref_to_string (ctx : norm_ctx) (x : T.trait_ref) : string =
-  PA.trait_ref_to_string ctx.ctx x
+  let ctx = norm_ctx_to_type_formatter ctx in
+  Print.Types.trait_ref_to_string ctx x
 
 let trait_instance_id_to_string (ctx : norm_ctx) (x : T.trait_instance_id) :
     string =
-  PA.trait_instance_id_to_string ctx.ctx x
+  let ctx = norm_ctx_to_type_formatter ctx in
+  Print.Types.trait_instance_id_to_string ctx x
 
 let generic_args_to_string (ctx : norm_ctx) (x : T.generic_args) : string =
-  PA.generic_args_to_string ctx.ctx x
+  let ctx = norm_ctx_to_type_formatter ctx in
+  Print.Types.generic_args_to_string ctx x
 
 let generic_params_to_string (ctx : norm_ctx) (x : T.generic_params) : string =
-  "<" ^ String.concat ", " (fst (PA.generic_params_to_strings ctx.ctx x)) ^ ">"
+  let ctx = norm_ctx_to_type_formatter ctx in
+  "<"
+  ^ String.concat ", " (fst (Print.Types.generic_params_to_strings ctx x))
+  ^ ">"
 
 (** Small utility to lookup trait impls, together with a substitution. *)
-let ctx_lookup_trait_impl (ctx : norm_ctx) (impl_id : T.TraitImplId.id)
+let norm_ctx_lookup_trait_impl (ctx : norm_ctx) (impl_id : T.TraitImplId.id)
     (generics : T.generic_args) : A.trait_impl * Subst.subst =
   (* Lookup the implementation *)
-  let trait_impl = C.ctx_lookup_trait_impl ctx.ctx impl_id in
+  let trait_impl = A.TraitImplId.Map.find impl_id ctx.trait_impls in
   (* The substitution *)
   let tr_self = T.UnknownTrait __FUNCTION__ in
   let subst =
@@ -138,20 +197,20 @@ let ctx_lookup_trait_impl (ctx : norm_ctx) (impl_id : T.TraitImplId.id)
   (* Return *)
   (trait_impl, subst)
 
-let ctx_lookup_trait_impl_ty (ctx : norm_ctx) (impl_id : T.TraitImplId.id)
+let norm_ctx_lookup_trait_impl_ty (ctx : norm_ctx) (impl_id : T.TraitImplId.id)
     (generics : T.generic_args) (type_name : string) : T.ty =
   (* Lookup the implementation *)
-  let trait_impl, subst = ctx_lookup_trait_impl ctx impl_id generics in
+  let trait_impl, subst = norm_ctx_lookup_trait_impl ctx impl_id generics in
   (* Lookup the type *)
   let ty = snd (List.assoc type_name trait_impl.types) in
   (* Substitute *)
   Subst.ty_substitute subst ty
 
-let ctx_lookup_trait_impl_parent_clause (ctx : norm_ctx)
+let norm_ctx_lookup_trait_impl_parent_clause (ctx : norm_ctx)
     (impl_id : T.TraitImplId.id) (generics : T.generic_args)
     (clause_id : T.TraitClauseId.id) : T.trait_ref =
   (* Lookup the implementation *)
-  let trait_impl, subst = ctx_lookup_trait_impl ctx impl_id generics in
+  let trait_impl, subst = norm_ctx_lookup_trait_impl ctx impl_id generics in
   (* Lookup the clause *)
   let clause = T.TraitClauseId.nth trait_impl.parent_trait_refs clause_id in
   (* Sanity check: the clause necessarily refers to an impl *)
@@ -159,11 +218,11 @@ let ctx_lookup_trait_impl_parent_clause (ctx : norm_ctx)
   (* Substitute *)
   Subst.trait_ref_substitute subst clause
 
-let ctx_lookup_trait_impl_item_clause (ctx : norm_ctx)
+let norm_ctx_lookup_trait_impl_item_clause (ctx : norm_ctx)
     (impl_id : T.TraitImplId.id) (generics : T.generic_args)
     (item_name : string) (clause_id : T.TraitClauseId.id) : T.trait_ref =
   (* Lookup the implementation *)
-  let trait_impl, subst = ctx_lookup_trait_impl ctx impl_id generics in
+  let trait_impl, subst = norm_ctx_lookup_trait_impl ctx impl_id generics in
   (* Lookup the item then its clause *)
   let item = List.assoc item_name trait_impl.types in
   let clause = T.TraitClauseId.nth (fst item) clause_id in
@@ -176,35 +235,36 @@ let ctx_lookup_trait_impl_item_clause (ctx : norm_ctx)
     and choosing a representative when there are equalities between types
     enforced by local clauses (i.e., `where Trait1::T = Trait2::U`.
 
-    See the comments for {!ctx_normalize_trait_instance_id}.
+    See the comments for {!norm_ctx_normalize_trait_instance_id}.
   *)
-let rec ctx_normalize_ty (ctx : norm_ctx) (ty : T.ty) : T.ty =
-  log#ldebug (lazy ("ctx_normalize_ty: " ^ ty_to_string ctx ty));
+let rec norm_ctx_normalize_ty (ctx : norm_ctx) (ty : T.ty) : T.ty =
+  log#ldebug (lazy ("norm_ctx_normalize_ty: " ^ ty_to_string ctx ty));
   match ty with
-  | T.TAdt (id, generics) -> TAdt (id, ctx_normalize_generic_args ctx generics)
+  | T.TAdt (id, generics) ->
+      TAdt (id, norm_ctx_normalize_generic_args ctx generics)
   | TVar _ | TLiteral _ | TNever -> ty
   | TRef (r, ty, rkind) ->
-      let ty = ctx_normalize_ty ctx ty in
+      let ty = norm_ctx_normalize_ty ctx ty in
       T.TRef (r, ty, rkind)
   | TRawPtr (ty, rkind) ->
-      let ty = ctx_normalize_ty ctx ty in
+      let ty = norm_ctx_normalize_ty ctx ty in
       TRawPtr (ty, rkind)
   | TArrow (inputs, output) ->
-      let inputs = List.map (ctx_normalize_ty ctx) inputs in
-      let output = ctx_normalize_ty ctx output in
+      let inputs = List.map (norm_ctx_normalize_ty ctx) inputs in
+      let output = norm_ctx_normalize_ty ctx output in
       TArrow (inputs, output)
   | TTraitType (trait_ref, generics, type_name) -> (
       log#ldebug
         (lazy
-          ("ctx_normalize_ty:\n- trait type: " ^ ty_to_string ctx ty
+          ("norm_ctx_normalize_ty:\n- trait type: " ^ ty_to_string ctx ty
          ^ "\n- trait_ref: "
           ^ trait_ref_to_string ctx trait_ref
           ^ "\n- raw trait ref:\n" ^ T.show_trait_ref trait_ref
           ^ "\n- generics:\n"
           ^ generic_args_to_string ctx generics));
       (* Normalize and attempt to project the type from the trait ref *)
-      let trait_ref = ctx_normalize_trait_ref ctx trait_ref in
-      let generics = ctx_normalize_generic_args ctx generics in
+      let trait_ref = norm_ctx_normalize_trait_ref ctx trait_ref in
+      let generics = norm_ctx_normalize_generic_args ctx generics in
       (*  For now, we don't support higher order types *)
       assert (generics = TypesUtils.mk_empty_generic_args);
       let ty : T.ty =
@@ -214,18 +274,19 @@ let rec ctx_normalize_ty (ctx : norm_ctx) (ty : T.ty) : T.ty =
             assert (ref_generics = TypesUtils.mk_empty_generic_args);
             log#ldebug
               (lazy
-                ("ctx_normalize_ty: trait type: trait ref: "
+                ("norm_ctx_normalize_ty: trait type: trait ref: "
                ^ ty_to_string ctx ty));
             (* Lookup the type *)
             let ty =
-              ctx_lookup_trait_impl_ty ctx impl_id trait_ref.generics type_name
+              norm_ctx_lookup_trait_impl_ty ctx impl_id trait_ref.generics
+                type_name
             in
             (* Normalize *)
-            ctx_normalize_ty ctx ty
+            norm_ctx_normalize_ty ctx ty
         | T.TraitImpl impl_id ->
             log#ldebug
               (lazy
-                ("ctx_normalize_ty (trait impl):\n- trait type: "
+                ("norm_ctx_normalize_ty (trait impl):\n- trait type: "
                ^ ty_to_string ctx ty ^ "\n- trait_ref: "
                 ^ trait_ref_to_string ctx trait_ref
                 ^ "\n- raw trait ref:\n" ^ T.show_trait_ref trait_ref));
@@ -237,14 +298,15 @@ let rec ctx_normalize_ty (ctx : norm_ctx) (ty : T.ty) : T.ty =
             *)
             (* Lookup the type *)
             let ty =
-              ctx_lookup_trait_impl_ty ctx impl_id trait_ref.generics type_name
+              norm_ctx_lookup_trait_impl_ty ctx impl_id trait_ref.generics
+                type_name
             in
             (* Normalize *)
-            ctx_normalize_ty ctx ty
+            norm_ctx_normalize_ty ctx ty
         | _ ->
             log#ldebug
               (lazy
-                ("ctx_normalize_ty: trait type: not a trait ref: "
+                ("norm_ctx_normalize_ty: trait type: not a trait ref: "
                ^ ty_to_string ctx ty ^ "\n- trait_ref: "
                 ^ trait_ref_to_string ctx trait_ref
                 ^ "\n- raw trait ref:\n" ^ T.show_trait_ref trait_ref));
@@ -254,7 +316,7 @@ let rec ctx_normalize_ty (ctx : norm_ctx) (ty : T.ty) : T.ty =
       in
       let tr : C.trait_type_ref = { C.trait_ref; type_name } in
       (* Lookup the representative, if there is *)
-      match ctx_get_ty_repr ctx tr with None -> ty | Some ty -> ty)
+      match norm_ctx_get_ty_repr ctx tr with None -> ty | Some ty -> ty)
 
 (** This returns the normalized trait instance id together with an optional
     reference to a trait **implementation** (the `trait_ref` we return has
@@ -300,8 +362,8 @@ let rec ctx_normalize_ty (ctx : norm_ctx) (ty : T.ty) : T.ty =
         In this case we can lookup the trait implementation and recursively project
         over it.
  *)
-and ctx_normalize_trait_instance_id (ctx : norm_ctx) (id : T.trait_instance_id)
-    : T.trait_instance_id * T.trait_ref option =
+and norm_ctx_normalize_trait_instance_id (ctx : norm_ctx)
+    (id : T.trait_instance_id) : T.trait_instance_id * T.trait_ref option =
   match id with
   | Self -> (id, None)
   | TraitImpl _ ->
@@ -312,7 +374,7 @@ and ctx_normalize_trait_instance_id (ctx : norm_ctx) (id : T.trait_instance_id)
   | Clause _ -> (id, None)
   | BuiltinOrAuto _ -> (id, None)
   | ParentClause (inst_id, decl_id, clause_id) -> (
-      let inst_id, impl = ctx_normalize_trait_instance_id ctx inst_id in
+      let inst_id, impl = norm_ctx_normalize_trait_instance_id ctx inst_id in
       (* Check if the inst_id refers to a specific implementation, if yes project *)
       match impl with
       | None ->
@@ -336,14 +398,14 @@ and ctx_normalize_trait_instance_id (ctx : norm_ctx) (id : T.trait_instance_id)
             TypesUtils.trait_instance_id_as_trait_impl impl.trait_id
           in
           let clause =
-            ctx_lookup_trait_impl_parent_clause ctx impl_id impl.generics
+            norm_ctx_lookup_trait_impl_parent_clause ctx impl_id impl.generics
               clause_id
           in
           (* Normalize the clause *)
-          let clause = ctx_normalize_trait_ref ctx clause in
+          let clause = norm_ctx_normalize_trait_ref ctx clause in
           (TraitRef clause, Some clause))
   | ItemClause (inst_id, decl_id, item_name, clause_id) -> (
-      let inst_id, impl = ctx_normalize_trait_instance_id ctx inst_id in
+      let inst_id, impl = norm_ctx_normalize_trait_instance_id ctx inst_id in
       (* Check if the inst_id refers to a specific implementation, if yes project *)
       match impl with
       | None ->
@@ -367,18 +429,20 @@ and ctx_normalize_trait_instance_id (ctx : norm_ctx) (id : T.trait_instance_id)
             TypesUtils.trait_instance_id_as_trait_impl impl.trait_id
           in
           let clause =
-            ctx_lookup_trait_impl_item_clause ctx impl_id impl.generics
+            norm_ctx_lookup_trait_impl_item_clause ctx impl_id impl.generics
               item_name clause_id
           in
           (* Normalize the clause *)
-          let clause = ctx_normalize_trait_ref ctx clause in
+          let clause = norm_ctx_normalize_trait_ref ctx clause in
           (TraitRef clause, Some clause))
   | TraitRef { T.trait_id = T.TraitImpl trait_id; generics; trait_decl_ref } ->
       (* We can't simplify the id *yet* : we will simplify it when projecting.
          However, we have an implementation to return *)
       (* Normalize the generics *)
-      let generics = ctx_normalize_generic_args ctx generics in
-      let trait_decl_ref = ctx_normalize_trait_decl_ref ctx trait_decl_ref in
+      let generics = norm_ctx_normalize_generic_args ctx generics in
+      let trait_decl_ref =
+        norm_ctx_normalize_trait_decl_ref ctx trait_decl_ref
+      in
       let trait_ref : T.trait_ref =
         { T.trait_id = T.TraitImpl trait_id; generics; trait_decl_ref }
       in
@@ -390,7 +454,7 @@ and ctx_normalize_trait_instance_id (ctx : norm_ctx) (id : T.trait_instance_id)
       assert (trait_ref.generics = TypesUtils.mk_empty_generic_args);
       (trait_ref.trait_id, None)
   | FnPointer ty ->
-      let ty = ctx_normalize_ty ctx ty in
+      let ty = norm_ctx_normalize_ty ctx ty in
       (* TODO: we might want to return the ref to the function pointer,
          in order to later normalize a call to this function pointer *)
       (FnPointer ty, None)
@@ -398,59 +462,73 @@ and ctx_normalize_trait_instance_id (ctx : norm_ctx) (id : T.trait_instance_id)
       (* This is actually an error case *)
       (id, None)
 
-and ctx_normalize_generic_args (ctx : norm_ctx) (generics : T.generic_args) :
-    T.generic_args =
+and norm_ctx_normalize_generic_args (ctx : norm_ctx) (generics : T.generic_args)
+    : T.generic_args =
   let { T.regions; types; const_generics; trait_refs } = generics in
-  let types = List.map (ctx_normalize_ty ctx) types in
-  let trait_refs = List.map (ctx_normalize_trait_ref ctx) trait_refs in
+  let types = List.map (norm_ctx_normalize_ty ctx) types in
+  let trait_refs = List.map (norm_ctx_normalize_trait_ref ctx) trait_refs in
   { T.regions; types; const_generics; trait_refs }
 
-and ctx_normalize_trait_ref (ctx : norm_ctx) (trait_ref : T.trait_ref) :
+and norm_ctx_normalize_trait_ref (ctx : norm_ctx) (trait_ref : T.trait_ref) :
     T.trait_ref =
   log#ldebug
     (lazy
-      ("ctx_normalize_trait_ref: "
+      ("norm_ctx_normalize_trait_ref: "
       ^ trait_ref_to_string ctx trait_ref
       ^ "\n- raw trait ref:\n" ^ T.show_trait_ref trait_ref));
   let { T.trait_id; generics; trait_decl_ref } = trait_ref in
   (* Check if the id is an impl, otherwise normalize it *)
-  let trait_id, norm_trait_ref = ctx_normalize_trait_instance_id ctx trait_id in
+  let trait_id, norm_trait_ref =
+    norm_ctx_normalize_trait_instance_id ctx trait_id
+  in
   match norm_trait_ref with
   | None ->
       log#ldebug
         (lazy
-          ("ctx_normalize_trait_ref: no norm: "
+          ("norm_ctx_normalize_trait_ref: no norm: "
           ^ trait_instance_id_to_string ctx trait_id));
-      let generics = ctx_normalize_generic_args ctx generics in
-      let trait_decl_ref = ctx_normalize_trait_decl_ref ctx trait_decl_ref in
+      let generics = norm_ctx_normalize_generic_args ctx generics in
+      let trait_decl_ref =
+        norm_ctx_normalize_trait_decl_ref ctx trait_decl_ref
+      in
       { T.trait_id; generics; trait_decl_ref }
   | Some trait_ref ->
       log#ldebug
         (lazy
-          ("ctx_normalize_trait_ref: normalized to: "
+          ("norm_ctx_normalize_trait_ref: normalized to: "
           ^ trait_ref_to_string ctx trait_ref));
       assert (generics = TypesUtils.mk_empty_generic_args);
       trait_ref
 
 (* Not sure this one is really necessary *)
-and ctx_normalize_trait_decl_ref (ctx : norm_ctx)
+and norm_ctx_normalize_trait_decl_ref (ctx : norm_ctx)
     (trait_decl_ref : T.trait_decl_ref) : T.trait_decl_ref =
   let { T.trait_decl_id; decl_generics } = trait_decl_ref in
-  let decl_generics = ctx_normalize_generic_args ctx decl_generics in
+  let decl_generics = norm_ctx_normalize_generic_args ctx decl_generics in
   { T.trait_decl_id; decl_generics }
 
-let ctx_normalize_trait_type_constraint (ctx : norm_ctx)
+let norm_ctx_normalize_trait_type_constraint (ctx : norm_ctx)
     (ttc : T.trait_type_constraint) : T.trait_type_constraint =
   let { T.trait_ref; generics; type_name; ty } = ttc in
-  let trait_ref = ctx_normalize_trait_ref ctx trait_ref in
-  let generics = ctx_normalize_generic_args ctx generics in
-  let ty = ctx_normalize_ty ctx ty in
+  let trait_ref = norm_ctx_normalize_trait_ref ctx trait_ref in
+  let generics = norm_ctx_normalize_generic_args ctx generics in
+  let ty = norm_ctx_normalize_ty ctx ty in
   { T.trait_ref; generics; type_name; ty }
 
-let mk_norm_ctx (ctx : C.eval_ctx) : norm_ctx = { ctx }
+let mk_norm_ctx (ctx : C.eval_ctx) : norm_ctx =
+  {
+    norm_trait_types = ctx.norm_trait_types;
+    type_decls = ctx.type_context.type_decls;
+    fun_decls = ctx.fun_context.fun_decls;
+    global_decls = ctx.global_context.global_decls;
+    trait_decls = ctx.trait_decls_context.trait_decls;
+    trait_impls = ctx.trait_impls_context.trait_impls;
+    type_vars = ctx.type_vars;
+    const_generic_vars = ctx.const_generic_vars;
+  }
 
 let ctx_normalize_ty (ctx : C.eval_ctx) (ty : T.ty) : T.ty =
-  ctx_normalize_ty (mk_norm_ctx ctx) ty
+  norm_ctx_normalize_ty (mk_norm_ctx ctx) ty
 
 (** Normalize a type and erase the regions at the same time *)
 let ctx_normalize_erase_ty (ctx : C.eval_ctx) (ty : T.ty) : T.ty =
@@ -459,7 +537,7 @@ let ctx_normalize_erase_ty (ctx : C.eval_ctx) (ty : T.ty) : T.ty =
 
 let ctx_normalize_trait_type_constraint (ctx : C.eval_ctx)
     (ttc : T.trait_type_constraint) : T.trait_type_constraint =
-  ctx_normalize_trait_type_constraint (mk_norm_ctx ctx) ttc
+  norm_ctx_normalize_trait_type_constraint (mk_norm_ctx ctx) ttc
 
 (** Same as [type_decl_get_instantiated_variants_fields_types] but normalizes the types *)
 let type_decl_get_inst_norm_variants_fields_rtypes (ctx : C.eval_ctx)
diff --git a/compiler/Contexts.ml b/compiler/Contexts.ml
index 12927aab..17ebd315 100644
--- a/compiler/Contexts.ml
+++ b/compiler/Contexts.ml
@@ -254,18 +254,14 @@ type eval_ctx = {
 
 let lookup_type_var_opt (ctx : eval_ctx) (vid : TypeVarId.id) : type_var option
     =
-  if TypeVarId.to_int vid < List.length ctx.type_vars then
-    Some (TypeVarId.nth ctx.type_vars vid)
-  else None
+  TypeVarId.nth_opt ctx.type_vars vid
 
 let lookup_type_var (ctx : eval_ctx) (vid : TypeVarId.id) : type_var =
   TypeVarId.nth ctx.type_vars vid
 
 let lookup_const_generic_var_opt (ctx : eval_ctx) (vid : ConstGenericVarId.id) :
     const_generic_var option =
-  if ConstGenericVarId.to_int vid < List.length ctx.const_generic_vars then
-    Some (ConstGenericVarId.nth ctx.const_generic_vars vid)
-  else None
+  ConstGenericVarId.nth_opt ctx.const_generic_vars vid
 
 let lookup_const_generic_var (ctx : eval_ctx) (vid : ConstGenericVarId.id) :
     const_generic_var =
diff --git a/compiler/Interpreter.ml b/compiler/Interpreter.ml
index b94825cc..e9c0f17f 100644
--- a/compiler/Interpreter.ml
+++ b/compiler/Interpreter.ml
@@ -31,6 +31,7 @@ let compute_contexts (m : A.crate) : C.decls_ctx =
   in
   let regions_hierarchies =
     RegionsHierarchy.compute_regions_hierarchies type_decls fun_decls
+      global_decls trait_decls trait_impls
   in
   let fun_ctx = { C.fun_decls; fun_infos; regions_hierarchies } in
   let global_ctx = { C.global_decls } in
diff --git a/compiler/RegionsHierarchy.ml b/compiler/RegionsHierarchy.ml
index ce5880bf..8227e1fa 100644
--- a/compiler/RegionsHierarchy.ml
+++ b/compiler/RegionsHierarchy.ml
@@ -39,8 +39,31 @@ module Subst = Substitute
 let log = Logging.regions_hierarchy_log
 
 let compute_regions_hierarchy_for_sig (type_decls : type_decl TypeDeclId.Map.t)
-    (fun_name : name) (sg : fun_sig) : region_groups =
+    (fun_decls : fun_decl FunDeclId.Map.t)
+    (global_decls : global_decl GlobalDeclId.Map.t)
+    (trait_decls : trait_decl TraitDeclId.Map.t)
+    (trait_impls : trait_impl TraitImplId.Map.t) (fun_name : name)
+    (sg : fun_sig) : region_groups =
   log#ldebug (lazy (__FUNCTION__ ^ ": " ^ name_to_string fun_name));
+  (* Initialize a normalization context (we may need to normalize some
+     associated types) *)
+  let norm_ctx : AssociatedTypes.norm_ctx =
+    let norm_trait_types =
+      AssociatedTypes.compute_norm_trait_types_from_preds
+        sg.preds.trait_type_constraints
+    in
+    {
+      norm_trait_types;
+      type_decls;
+      fun_decls;
+      global_decls;
+      trait_decls;
+      trait_impls;
+      type_vars = sg.generics.types;
+      const_generic_vars = sg.generics.const_generics;
+    }
+  in
+
   (* Create the dependency graph.
 
      An edge from 'short to 'long means that 'long outlives 'short (that is
@@ -139,7 +162,13 @@ let compute_regions_hierarchy_for_sig (type_decls : type_decl TypeDeclId.Map.t)
     List.iter (explore_ty outer) types
   in
 
-  List.iter (explore_ty []) (sg.output :: sg.inputs);
+  (* Normalize the types then explore *)
+  let tys =
+    List.map
+      (AssociatedTypes.norm_ctx_normalize_ty norm_ctx)
+      (sg.output :: sg.inputs)
+  in
+  List.iter (explore_ty []) tys;
 
   (* Compute the ordered SCCs *)
   let module Scc = SCC.Make (RegionOrderedType) in
@@ -231,7 +260,10 @@ let compute_regions_hierarchy_for_sig (type_decls : type_decl TypeDeclId.Map.t)
     (SccId.Map.bindings sccs.sccs)
 
 let compute_regions_hierarchies (type_decls : type_decl TypeDeclId.Map.t)
-    (fun_decls : fun_decl FunDeclId.Map.t) : region_groups FunIdMap.t =
+    (fun_decls : fun_decl FunDeclId.Map.t)
+    (global_decls : global_decl GlobalDeclId.Map.t)
+    (trait_decls : trait_decl TraitDeclId.Map.t)
+    (trait_impls : trait_impl TraitImplId.Map.t) : region_groups FunIdMap.t =
   let regular =
     List.map
       (fun ((fid, d) : FunDeclId.id * fun_decl) ->
@@ -247,5 +279,7 @@ let compute_regions_hierarchies (type_decls : type_decl TypeDeclId.Map.t)
   FunIdMap.of_list
     (List.map
        (fun (fid, (name, sg)) ->
-         (fid, compute_regions_hierarchy_for_sig type_decls name sg))
+         ( fid,
+           compute_regions_hierarchy_for_sig type_decls fun_decls global_decls
+             trait_decls trait_impls name sg ))
        (regular @ assumed))