Correctly resolve type relative paths in methods (#1489)

Author: nilehmann

crates/flux-desugar/src/desugar.rs

@@ -1698,10 +1698,7 @@ trait DesugarCtxt<'genv, 'tcx: 'genv>: ErrorEmitter + ErrorCollector<ErrorGuaran
 
 impl<'genv, 'tcx> DesugarCtxt<'genv, 'tcx> for RustItemCtxt<'_, 'genv, 'tcx> {
     fn next_fhir_id(&self) -> FhirId {
-        FhirId {
-            owner: FluxOwnerId::Rust(self.owner.local_id()),
-            local_id: self.local_id_gen.fresh(),
-        }
+        FhirId { owner: FluxOwnerId::Rust(self.owner), local_id: self.local_id_gen.fresh() }
     }
 
     fn genv(&self) -> GlobalEnv<'genv, 'tcx> {

crates/flux-desugar/src/desugar/lift.rs

@@ -533,10 +533,7 @@ impl<'genv> RustItemCtxt<'_, 'genv, '_> {
     }
 
     fn next_fhir_id(&self) -> FhirId {
-        FhirId {
-            owner: FluxOwnerId::Rust(self.owner.local_id()),
-            local_id: self.local_id_gen.fresh(),
-        }
+        FhirId { owner: FluxOwnerId::Rust(self.owner), local_id: self.local_id_gen.fresh() }
     }
 
     fn local_id(&self) -> LocalDefId {

crates/flux-fhir-analysis/src/conv/mod.rs

@@ -295,7 +295,7 @@ pub(crate) fn conv_adt_sort_def(
     def_id: MaybeExternId,
     kind: &fhir::RefinementKind,
 ) -> QueryResult<rty::AdtSortDef> {
-    let wfckresults = &WfckResults::new(OwnerId { def_id: def_id.local_id() });
+    let wfckresults = &WfckResults::new(def_id.map(|def_id| OwnerId { def_id }));
     let mut cx = AfterSortck::new(genv, wfckresults).into_conv_ctxt();
     match kind {
         fhir::RefinementKind::Refined(refined_by) => {
@@ -1502,7 +1502,8 @@ impl<'genv, 'tcx: 'genv, P: ConvPhase<'genv, 'tcx>> ConvCtxt<P> {
             }
             fhir::Res::Def(DefKind::TyParam, param_id)
             | fhir::Res::SelfTyParam { trait_: param_id } => {
-                let predicates = type_param_predicates(tcx, param_id);
+                let item_def_id = self.owner().resolved_id().unwrap();
+                let predicates = type_param_predicates(tcx, item_def_id, param_id);
                 self.probe_single_bound_for_assoc_item(
                     || {
                         tag.transitive_bounds_that_define_assoc_item(
@@ -1598,8 +1599,7 @@ impl<'genv, 'tcx: 'genv, P: ConvPhase<'genv, 'tcx>> ConvCtxt<P> {
     fn refiner(&self) -> QueryResult<Refiner<'genv, 'tcx>> {
         match self.owner() {
             FluxOwnerId::Rust(owner_id) => {
-                let def_id = self.genv().maybe_extern_id(owner_id.def_id);
-                Refiner::default_for_item(self.genv(), def_id.resolved_id())
+                Refiner::default_for_item(self.genv(), owner_id.resolved_id())
             }
             FluxOwnerId::Flux(_) => Err(query_bug!("cannot refine types insicde flux item")),
         }
@@ -2843,19 +2843,14 @@ impl AssocItemTag for AssocReftTag {
 /// problem for us so we can use it instead of [`TyCtxt::type_param_predicates`].
 fn type_param_predicates<'tcx>(
     tcx: TyCtxt<'tcx>,
+    item_def_id: DefId,
     param_id: DefId,
 ) -> impl Iterator<Item = ty::PolyTraitPredicate<'tcx>> {
-    let parent = if tcx.def_kind(param_id) == DefKind::Trait {
-        // If the param_id is a trait then this is the `Self` parameter and the parent is the trait itself
-        param_id
-    } else {
-        tcx.parent(param_id)
-    };
     let param_index = tcx
-        .generics_of(parent)
+        .generics_of(item_def_id)
         .param_def_id_to_index(tcx, param_id)
         .unwrap();
-    let predicates = tcx.predicates_of(parent).instantiate_identity(tcx);
+    let predicates = tcx.predicates_of(item_def_id).instantiate_identity(tcx);
     predicates.into_iter().filter_map(move |(clause, _)| {
         clause
             .as_trait_clause()

crates/flux-fhir-analysis/src/lib.rs

@@ -202,14 +202,14 @@ fn adt_def(genv: GlobalEnv, def_id: MaybeExternId) -> QueryResult<rty::AdtDef> {
 
 fn constant_info(genv: GlobalEnv, def_id: MaybeExternId) -> QueryResult<rty::ConstantInfo> {
     let node = genv.fhir_node(def_id.local_id())?;
-    let owner = rustc_hir::OwnerId { def_id: def_id.local_id() };
     let Some(sort) = genv.sort_of_def_id(def_id.resolved_id()).emit(&genv)? else {
         return Ok(rty::ConstantInfo::Uninterpreted);
     };
     let tcx = genv.tcx();
     match node {
         fhir::Node::Item(fhir::Item { kind: fhir::ItemKind::Const(Some(expr)), .. }) => {
             // If the constant has a `#[consant(expr)]` annotation we use that
+            let owner = def_id.map(|def_id| rustc_hir::OwnerId { def_id });
             let wfckresults = wf::check_constant_expr(genv, owner, expr, &sort)?;
             let expr = AfterSortck::new(genv, &wfckresults)
                 .into_conv_ctxt()
@@ -377,7 +377,7 @@ fn sort_of_assoc_reft(
     match &genv.fhir_expect_item(container_id.local_id())?.kind {
         fhir::ItemKind::Trait(trait_) => {
             let assoc_reft = trait_.find_assoc_reft(assoc_id.name()).unwrap();
-            let wfckresults = WfckResults::new(OwnerId { def_id: container_id.local_id() });
+            let wfckresults = WfckResults::new(container_id.map(|def_id| OwnerId { def_id }));
             let mut cx = AfterSortck::new(genv, &wfckresults).into_conv_ctxt();
             let inputs = assoc_reft
                 .params
@@ -389,7 +389,7 @@ fn sort_of_assoc_reft(
         }
         fhir::ItemKind::Impl(impl_) => {
             let assoc_reft = impl_.find_assoc_reft(assoc_id.name()).unwrap();
-            let wfckresults = WfckResults::new(OwnerId { def_id: container_id.local_id() });
+            let wfckresults = WfckResults::new(container_id.map(|def_id| OwnerId { def_id }));
             let mut cx = AfterSortck::new(genv, &wfckresults).into_conv_ctxt();
             let inputs = assoc_reft
                 .params

crates/flux-fhir-analysis/src/wf/mod.rs

@@ -48,7 +48,7 @@ pub(crate) fn check_flux_item<'genv>(
 
 pub(crate) fn check_constant_expr<'genv>(
     genv: GlobalEnv<'genv, '_>,
-    owner: OwnerId,
+    owner: MaybeExternId<OwnerId>,
     expr: &fhir::Expr<'genv>,
     sort: &rty::Sort,
 ) -> Result<WfckResults> {
@@ -70,7 +70,7 @@ pub(crate) fn check_invariants<'genv>(
     params: &[fhir::RefineParam<'genv>],
     invariants: &[fhir::Expr<'genv>],
 ) -> Result<WfckResults> {
-    let owner = FluxOwnerId::Rust(adt_def_id.local_id());
+    let owner = FluxOwnerId::Rust(adt_def_id);
     let mut infcx = InferCtxt::new(genv, owner);
     Wf::with(&mut infcx, |wf| {
         wf.declare_params_for_invariants(params, invariants)?;
@@ -94,7 +94,7 @@ pub(crate) fn check_node<'genv>(
     genv: GlobalEnv<'genv, '_>,
     node: &fhir::OwnerNode<'genv>,
 ) -> Result<WfckResults> {
-    let mut infcx = InferCtxt::new(genv, node.owner_id().local_id().into());
+    let mut infcx = InferCtxt::new(genv, node.owner_id().into());
     Wf::with(&mut infcx, |wf| {
         wf.init_infcx_for_node(node)
             .map_err(|err| err.at(genv.tcx().def_span(node.owner_id().local_id())))

crates/flux-fhir-analysis/src/wf/sortck.rs

@@ -581,8 +581,7 @@ impl<'genv, 'tcx> InferCtxt<'genv, 'tcx> {
             .infer_ctxt()
             .with_next_trait_solver(true)
             .build(TypingMode::non_body_analysis());
-        if let Some(def_id) = owner.def_id() {
-            let def_id = genv.maybe_extern_id(def_id).resolved_id();
+        if let Some(def_id) = owner.resolved_id() {
             t.deeply_normalize_sorts(def_id, genv, &infcx)
         } else {
             Ok(t.clone())
@@ -801,12 +800,14 @@ impl<'genv> InferCtxt<'genv, '_> {
                 .insert(node.fhir_id, sort);
         }
 
-        let allow_uninterpreted_cast = self
-            .owner
-            .def_id()
-            .map_or_else(flux_config::allow_uninterpreted_cast, |def_id| {
-                self.genv.infer_opts(def_id).allow_uninterpreted_cast
-            });
+        let allow_uninterpreted_cast =
+            self.owner
+                .as_rust()
+                .map_or_else(flux_config::allow_uninterpreted_cast, |def_id| {
+                    self.genv
+                        .infer_opts(def_id.local_id().def_id)
+                        .allow_uninterpreted_cast
+                });
 
         // Make sure that function applications are fully resolved
         for (node, sort_args) in std::mem::take(&mut self.sort_args_of_app) {

crates/flux-middle/src/def_id.rs

@@ -110,7 +110,7 @@ impl rustc_middle::query::IntoQueryParam<FluxDefId> for FluxLocalDefId {
 /// [_local id_]: MaybeExternId::local_id
 /// [_resolved id_]: MaybeExternId::resolved_id
 /// [`GlobalEnv::maybe_extern_id`]: crate::global_env::GlobalEnv::maybe_extern_id
-#[derive(Clone, Copy, Debug)]
+#[derive(Debug, Copy, Clone, Hash, PartialEq, Eq, Encodable, Decodable)]
 pub enum MaybeExternId<Id = LocalDefId> {
     /// An id for a local spec.
     Local(Id),

crates/flux-middle/src/fhir.rs

@@ -591,16 +591,20 @@ pub enum Lifetime {
 #[derive(Debug, Copy, Clone, Hash, PartialEq, Eq, Encodable, Decodable)]
 pub enum FluxOwnerId {
     Flux(FluxLocalDefId),
-    Rust(OwnerId),
+    Rust(MaybeExternId<OwnerId>),
 }
 
 impl FluxOwnerId {
-    pub fn def_id(self) -> Option<LocalDefId> {
+    pub fn as_rust(self) -> Option<MaybeExternId<OwnerId>> {
         match self {
             FluxOwnerId::Flux(_) => None,
-            FluxOwnerId::Rust(owner_id) => Some(owner_id.def_id),
+            FluxOwnerId::Rust(owner_id) => Some(owner_id),
         }
     }
+
+    pub fn resolved_id(self) -> Option<DefId> {
+        self.as_rust().map(MaybeExternId::resolved_id)
+    }
 }
 
 /// A unique identifier for a node in the AST. Like [`HirId`] it is composed of an `owner` and a
@@ -1107,8 +1111,8 @@ impl PolyTraitRef<'_> {
     }
 }
 
-impl From<OwnerId> for FluxOwnerId {
-    fn from(owner_id: OwnerId) -> Self {
+impl From<MaybeExternId<OwnerId>> for FluxOwnerId {
+    fn from(owner_id: MaybeExternId<OwnerId>) -> Self {
         FluxOwnerId::Rust(owner_id)
     }
 }

tests/tests/pos/surface/associated_type01.rs

@@ -0,0 +1,18 @@
+// Test that we can resolve sasociated type in methods
+
+struct S<T> {
+    x: T,
+}
+
+trait MyTrait {
+    type Assoc;
+}
+
+impl<T> S<T> {
+    fn test(&self) -> &T::Assoc
+    where
+        T: MyTrait,
+    {
+        todo!()
+    }
+}