Install bidirectional outlives predicates for RPITITs (and RPITs) correctly

4 files changed, 121 insertions, 88 deletions

diff --git a/compiler/rustc_hir_analysis/src/collect/predicates_of.rs b/compiler/rustc_hir_analysis/src/collect/predicates_of.rs
index f0a426722c3..495e663666c 100644
--- a/compiler/rustc_hir_analysis/src/collect/predicates_of.rs
+++ b/compiler/rustc_hir_analysis/src/collect/predicates_of.rs
@@ -2,16 +2,16 @@ use crate::astconv::{AstConv, OnlySelfBounds, PredicateFilter};
 use crate::bounds::Bounds;
 use crate::collect::ItemCtxt;
 use crate::constrained_generic_params as cgp;
-use hir::{HirId, Lifetime, Node};
+use hir::{HirId, Node};
 use rustc_data_structures::fx::FxIndexSet;
 use rustc_hir as hir;
 use rustc_hir::def::DefKind;
 use rustc_hir::def_id::{DefId, LocalDefId};
 use rustc_hir::intravisit::{self, Visitor};
 use rustc_middle::ty::{self, Ty, TyCtxt};
-use rustc_middle::ty::{GenericPredicates, Generics, ImplTraitInTraitData, ToPredicate};
+use rustc_middle::ty::{GenericPredicates, ImplTraitInTraitData, ToPredicate};
 use rustc_span::symbol::Ident;
-use rustc_span::{Span, Symbol, DUMMY_SP};
+use rustc_span::{Span, DUMMY_SP};
 
 /// Returns a list of all type predicates (explicit and implicit) for the definition with
 /// ID `def_id`. This includes all predicates returned by `predicates_defined_on`, plus
@@ -55,17 +55,7 @@ fn gather_explicit_predicates_of(tcx: TyCtxt<'_>, def_id: LocalDefId) -> ty::Gen
     use rustc_hir::*;
 
     match tcx.opt_rpitit_info(def_id.to_def_id()) {
-        Some(ImplTraitInTraitData::Trait { opaque_def_id, fn_def_id }) => {
-            let opaque_ty_id = tcx.hir().local_def_id_to_hir_id(opaque_def_id.expect_local());
-            let opaque_ty_node = tcx.hir().get(opaque_ty_id);
-            let Node::Item(&Item {
-                kind: ItemKind::OpaqueTy(OpaqueTy { lifetime_mapping, .. }),
-                ..
-            }) = opaque_ty_node
-            else {
-                bug!("unexpected {opaque_ty_node:?}")
-            };
-
+        Some(ImplTraitInTraitData::Trait { fn_def_id, .. }) => {
             let mut predicates = Vec::new();
 
             // RPITITs should inherit the predicates of their parent. This is
@@ -78,13 +68,12 @@ fn gather_explicit_predicates_of(tcx: TyCtxt<'_>, def_id: LocalDefId) -> ty::Gen
 
             // We also install bidirectional outlives predicates for the RPITIT
             // to keep the duplicates lifetimes from opaque lowering in sync.
+            // We only need to compute bidirectional outlives for the duplicated
+            // opaque lifetimes, which explains the slicing below.
             compute_bidirectional_outlives_predicates(
                 tcx,
-                def_id,
-                lifetime_mapping.iter().map(|(lifetime, def_id)| {
-                    (**lifetime, (*def_id, lifetime.ident.name, lifetime.ident.span))
-                }),
-                tcx.generics_of(def_id.to_def_id()),
+                &tcx.generics_of(def_id.to_def_id()).params
+                    [tcx.generics_of(fn_def_id).params.len()..],
                 &mut predicates,
             );
 
@@ -351,21 +340,7 @@ fn gather_explicit_predicates_of(tcx: TyCtxt<'_>, def_id: LocalDefId) -> ty::Gen
         };
         debug!(?lifetimes);
 
-        let lifetime_mapping = std::iter::zip(lifetimes, ast_generics.params)
-            .map(|(arg, dup)| {
-                let hir::GenericArg::Lifetime(arg) = arg else { bug!() };
-                (**arg, dup)
-            })
-            .filter(|(_, dup)| matches!(dup.kind, hir::GenericParamKind::Lifetime { .. }))
-            .map(|(lifetime, dup)| (lifetime, (dup.def_id, dup.name.ident().name, dup.span)));
-
-        compute_bidirectional_outlives_predicates(
-            tcx,
-            def_id,
-            lifetime_mapping,
-            generics,
-            &mut predicates,
-        );
+        compute_bidirectional_outlives_predicates(tcx, &generics.params, &mut predicates);
         debug!(?predicates);
     }
 
@@ -379,41 +354,28 @@ fn gather_explicit_predicates_of(tcx: TyCtxt<'_>, def_id: LocalDefId) -> ty::Gen
 /// enforce that these lifetimes stay in sync.
 fn compute_bidirectional_outlives_predicates<'tcx>(
     tcx: TyCtxt<'tcx>,
-    item_def_id: LocalDefId,
-    lifetime_mapping: impl Iterator<Item = (Lifetime, (LocalDefId, Symbol, Span))>,
-    generics: &Generics,
+    opaque_own_params: &[ty::GenericParamDef],
     predicates: &mut Vec<(ty::Clause<'tcx>, Span)>,
 ) {
-    let icx = ItemCtxt::new(tcx, item_def_id);
-
-    for (arg, (dup_def, name, span)) in lifetime_mapping {
-        let orig_region = icx.astconv().ast_region_to_region(&arg, None);
-        if !matches!(orig_region.kind(), ty::ReEarlyBound(..)) {
-            // There is no late-bound lifetime to actually match up here, since the lifetime doesn't
-            // show up in the opaque's parent's args.
-            continue;
+    for param in opaque_own_params {
+        let orig_lifetime = tcx.map_rpit_lifetime_to_fn_lifetime(param.def_id.expect_local());
+        if let ty::ReEarlyBound(..) = *orig_lifetime {
+            let dup_lifetime = ty::Region::new_early_bound(
+                tcx,
+                ty::EarlyBoundRegion { def_id: param.def_id, index: param.index, name: param.name },
+            );
+            let span = tcx.def_span(param.def_id);
+            predicates.push((
+                ty::ClauseKind::RegionOutlives(ty::OutlivesPredicate(orig_lifetime, dup_lifetime))
+                    .to_predicate(tcx),
+                span,
+            ));
+            predicates.push((
+                ty::ClauseKind::RegionOutlives(ty::OutlivesPredicate(dup_lifetime, orig_lifetime))
+                    .to_predicate(tcx),
+                span,
+            ));
         }
-
-        let Some(dup_index) = generics.param_def_id_to_index(icx.tcx, dup_def.to_def_id()) else {
-            bug!()
-        };
-
-        let dup_region = ty::Region::new_early_bound(
-            tcx,
-            ty::EarlyBoundRegion { def_id: dup_def.to_def_id(), index: dup_index, name },
-        );
-
-        predicates.push((
-            ty::ClauseKind::RegionOutlives(ty::OutlivesPredicate(orig_region, dup_region))
-                .to_predicate(tcx),
-            span,
-        ));
-
-        predicates.push((
-            ty::ClauseKind::RegionOutlives(ty::OutlivesPredicate(dup_region, orig_region))
-                .to_predicate(tcx),
-            span,
-        ));
     }
 }
 
diff --git a/compiler/rustc_middle/src/ty/context.rs b/compiler/rustc_middle/src/ty/context.rs
index 5eef3e45f32..75a2c6cbd20 100644
--- a/compiler/rustc_middle/src/ty/context.rs
+++ b/compiler/rustc_middle/src/ty/context.rs
@@ -1931,6 +1931,75 @@ impl<'tcx> TyCtxt<'tcx> {
         )
     }
 
+    /// Given the def-id of an early-bound lifetime on an RPIT corresponding to
+    /// a duplicated captured lifetime, map it back to the early- or late-bound
+    /// lifetime of the function from which it originally as captured. If it is
+    /// a late-bound lifetime, this will represent the liberated (`ReFree`) lifetime
+    /// of the signature.
+    // FIXME(RPITIT): if we ever synthesize new lifetimes for RPITITs and not just
+    // re-use the generics of the opaque, this function will need to be tweaked slightly.
+    pub fn map_rpit_lifetime_to_fn_lifetime(
+        self,
+        mut rpit_lifetime_param_def_id: LocalDefId,
+    ) -> ty::Region<'tcx> {
+        debug_assert!(
+            matches!(self.def_kind(rpit_lifetime_param_def_id), DefKind::LifetimeParam),
+            "{rpit_lifetime_param_def_id:?} is a {}",
+            self.def_descr(rpit_lifetime_param_def_id.to_def_id())
+        );
+
+        loop {
+            let parent = self.local_parent(rpit_lifetime_param_def_id);
+            let hir::OpaqueTy { lifetime_mapping, .. } =
+                self.hir().get_by_def_id(parent).expect_item().expect_opaque_ty();
+
+            let Some((lifetime, _)) = lifetime_mapping
+                .iter()
+                .find(|(_, duplicated_param)| *duplicated_param == rpit_lifetime_param_def_id)
+            else {
+                bug!("duplicated lifetime param should be present");
+            };
+
+            match self.named_bound_var(lifetime.hir_id) {
+                Some(resolve_bound_vars::ResolvedArg::EarlyBound(ebv)) => {
+                    let new_parent = self.parent(ebv);
+
+                    // If we map to another opaque, then it should be a parent
+                    // of the opaque we mapped from. Continue mapping.
+                    if matches!(self.def_kind(new_parent), DefKind::OpaqueTy) {
+                        debug_assert_eq!(self.parent(parent.to_def_id()), new_parent);
+                        rpit_lifetime_param_def_id = ebv.expect_local();
+                        continue;
+                    }
+
+                    let generics = self.generics_of(new_parent);
+                    return ty::Region::new_early_bound(
+                        self,
+                        ty::EarlyBoundRegion {
+                            def_id: ebv,
+                            index: generics
+                                .param_def_id_to_index(self, ebv)
+                                .expect("early-bound var should be present in fn generics"),
+                            name: self.hir().name(self.local_def_id_to_hir_id(ebv.expect_local())),
+                        },
+                    );
+                }
+                Some(resolve_bound_vars::ResolvedArg::LateBound(_, _, lbv)) => {
+                    let new_parent = self.parent(lbv);
+                    return ty::Region::new_free(
+                        self,
+                        new_parent,
+                        ty::BoundRegionKind::BrNamed(
+                            lbv,
+                            self.hir().name(self.local_def_id_to_hir_id(lbv.expect_local())),
+                        ),
+                    );
+                }
+                _ => bug!(),
+            }
+        }
+    }
+
     /// Whether the `def_id` counts as const fn in the current crate, considering all active
     /// feature gates
     pub fn is_const_fn(self, def_id: DefId) -> bool {
diff --git a/compiler/rustc_ty_utils/src/implied_bounds.rs b/compiler/rustc_ty_utils/src/implied_bounds.rs
index 01fe3a408e8..ec2e0daaf88 100644
--- a/compiler/rustc_ty_utils/src/implied_bounds.rs
+++ b/compiler/rustc_ty_utils/src/implied_bounds.rs
@@ -2,7 +2,6 @@ use rustc_data_structures::fx::FxHashMap;
 use rustc_hir as hir;
 use rustc_hir::def::DefKind;
 use rustc_hir::def_id::LocalDefId;
-use rustc_middle::middle::resolve_bound_vars as rbv;
 use rustc_middle::query::Providers;
 use rustc_middle::ty::{self, Ty, TyCtxt};
 use rustc_span::Span;
@@ -52,9 +51,7 @@ fn assumed_wf_types<'tcx>(tcx: TyCtxt<'tcx>, def_id: LocalDefId) -> &'tcx [(Ty<'
             tcx.arena.alloc_from_iter(tys.into_iter().map(|ty| (ty, impl_spans.next().unwrap())))
         }
         DefKind::AssocTy if let Some(data) = tcx.opt_rpitit_info(def_id.to_def_id()) => match data {
-            ty::ImplTraitInTraitData::Trait { fn_def_id, opaque_def_id } => {
-                let hir::OpaqueTy { lifetime_mapping, .. } =
-                    *tcx.hir().expect_item(opaque_def_id.expect_local()).expect_opaque_ty();
+            ty::ImplTraitInTraitData::Trait { fn_def_id, .. } => {
                 // We need to remap all of the late-bound lifetimes in theassumed wf types
                 // of the fn (which are represented as ReFree) to the early-bound lifetimes
                 // of the RPITIT (which are represented by ReEarlyBound owned by the opaque).
@@ -66,28 +63,22 @@ fn assumed_wf_types<'tcx>(tcx: TyCtxt<'tcx>, def_id: LocalDefId) -> &'tcx [(Ty<'
                 // predicates we insert in the `explicit_predicates_of` query for RPITITs.
                 let mut mapping = FxHashMap::default();
                 let generics = tcx.generics_of(def_id);
-                for &(lifetime, new_early_bound_def_id) in
-                    lifetime_mapping
-                {
-                    if let Some(rbv::ResolvedArg::LateBound(_, _, def_id)) =
-                        tcx.named_bound_var(lifetime.hir_id)
-                    {
-                        let name = tcx.hir().name(lifetime.hir_id);
-                        let index = generics
-                            .param_def_id_to_index(tcx, new_early_bound_def_id.to_def_id())
-                            .unwrap();
+
+                // For each captured opaque lifetime, if it's late-bound (`ReFree` in this case,
+                // since it has been liberated), map it back to the early-bound lifetime of
+                // the GAT. Since RPITITs also have all of the fn's generics, we slice only
+                // the end of the list corresponding to the opaque's generics.
+                for param in &generics.params[tcx.generics_of(fn_def_id).params.len()..] {
+                    let orig_lt = tcx.map_rpit_lifetime_to_fn_lifetime(param.def_id.expect_local());
+                    if matches!(*orig_lt, ty::ReFree(..)) {
                         mapping.insert(
-                            ty::Region::new_free(
-                                tcx,
-                                fn_def_id,
-                                ty::BoundRegionKind::BrNamed(def_id, name),
-                            ),
+                            orig_lt,
                             ty::Region::new_early_bound(
                                 tcx,
                                 ty::EarlyBoundRegion {
-                                    def_id: new_early_bound_def_id.to_def_id(),
-                                    index,
-                                    name,
+                                    def_id: param.def_id,
+                                    index: param.index,
+                                    name: param.name,
                                 },
                             ),
                         );
diff --git a/tests/ui/impl-trait/in-trait/outlives-in-nested-rpit.rs b/tests/ui/impl-trait/in-trait/outlives-in-nested-rpit.rs
new file mode 100644
index 00000000000..6330242ceeb
--- /dev/null
+++ b/tests/ui/impl-trait/in-trait/outlives-in-nested-rpit.rs
@@ -0,0 +1,11 @@
+// check-pass
+
+#![feature(return_position_impl_trait_in_trait)]
+
+trait Foo {
+    fn early<'a, T: 'a>(x: &'a T) -> impl Iterator<Item = impl Into<&'a T>>;
+
+    fn late<'a, T>(x: &'a T) -> impl Iterator<Item = impl Into<&'a T>>;
+}
+
+fn main() {}