Resolve sorts of associated types in trait definitions (#903)

flux-rs · Nov 22, 2024 · cfaf7eb · cfaf7eb
1 parent 93e22bb
commit cfaf7eb
Show file tree

Hide file tree

Showing 11 changed files with 730 additions and 526 deletions.
diff --git a/crates/flux-desugar/src/resolver/refinement_resolver.rs b/crates/flux-desugar/src/resolver/refinement_resolver.rs
@@ -614,19 +614,24 @@ impl<'a, 'genv, 'tcx> RefinementResolver<'a, 'genv, 'tcx> {
     }
 
     fn try_resolve_sort_with_ribs(&mut self, path: &surface::SortPath) -> Option<fhir::SortRes> {
-        let res = self
+        let partial_res = self
             .resolver
-            .resolve_path_with_ribs(&path.segments, TypeNS)?
-            .full_res()?;
-        match res {
-            fhir::Res::Def(DefKind::Struct | DefKind::Enum, def_id) => {
+            .resolve_path_with_ribs(&path.segments, TypeNS)?;
+        match (partial_res.base_res(), partial_res.unresolved_segments()) {
+            (fhir::Res::Def(DefKind::Struct | DefKind::Enum, def_id), 0) => {
                 Some(fhir::SortRes::Adt(def_id))
             }
-            fhir::Res::Def(DefKind::TyParam, def_id) => Some(fhir::SortRes::TyParam(def_id)),
-            fhir::Res::SelfTyParam { trait_ } => {
+            (fhir::Res::Def(DefKind::TyParam, def_id), 0) => Some(fhir::SortRes::TyParam(def_id)),
+            (fhir::Res::SelfTyParam { trait_ }, 0) => {
                 Some(fhir::SortRes::SelfParam { trait_id: trait_ })
             }
-            fhir::Res::SelfTyAlias { alias_to, .. } => Some(fhir::SortRes::SelfAlias { alias_to }),
+            (fhir::Res::SelfTyParam { trait_ }, 1) => {
+                let ident = *path.segments.last().unwrap();
+                Some(fhir::SortRes::SelfParamAssoc { trait_id: trait_, ident })
+            }
+            (fhir::Res::SelfTyAlias { alias_to, .. }, 0) => {
+                Some(fhir::SortRes::SelfAlias { alias_to })
+            }
             _ => None,
         }
     }

diff --git a/crates/flux-fhir-analysis/src/compare_impl_item.rs b/crates/flux-fhir-analysis/src/compare_impl_item.rs
@@ -1,6 +1,9 @@
 use flux_common::result::ResultExt;
-use flux_middle::{def_id_to_string, global_env::GlobalEnv, MaybeExternId};
+use flux_middle::{
+    def_id_to_string, global_env::GlobalEnv, rty::fold::TypeFoldable, MaybeExternId,
+};
 use rustc_hash::FxHashSet;
+use rustc_infer::infer::TyCtxtInferExt;
 use rustc_span::{def_id::DefId, ErrorGuaranteed, Symbol};
 type Result<T = ()> = std::result::Result<T, ErrorGuaranteed>;
 
@@ -55,6 +58,8 @@ fn check_assoc_reft(
     trait_id: DefId,
     name: Symbol,
 ) -> Result {
+    let infcx = genv.tcx().infer_ctxt().build();
+
     let impl_span = genv
         .map()
         .expect_item(impl_id.local_id())
@@ -78,7 +83,10 @@ fn check_assoc_reft(
         )));
     };
 
-    let impl_sort = impl_sort.instantiate_identity();
+    let impl_sort = impl_sort
+        .instantiate_identity()
+        .normalize_projections(genv, &infcx, impl_id.resolved_id())
+        .emit(&genv)?;
 
     let Some(trait_sort) = genv.sort_of_assoc_reft(trait_id, name).emit(genv.sess())? else {
         return Err(genv.sess().emit_err(errors::InvalidAssocReft::new(
@@ -87,7 +95,10 @@ fn check_assoc_reft(
             def_id_to_string(trait_id),
         )));
     };
-    let trait_sort = trait_sort.instantiate(genv.tcx(), &impl_trait_ref.args, &[]);
+    let trait_sort = trait_sort
+        .instantiate(genv.tcx(), &impl_trait_ref.args, &[])
+        .normalize_projections(genv, &infcx, impl_id.resolved_id())
+        .emit(&genv)?;
 
     if impl_sort != trait_sort {
         return Err(genv