Error for RPITIT hidden tys that capture more than their trait defn

4 files changed, 130 insertions, 39 deletions

diff --git a/compiler/rustc_hir_analysis/src/check/compare_impl_item.rs b/compiler/rustc_hir_analysis/src/check/compare_impl_item.rs
index 3048c175e1e..6873fa0bf27 100644
--- a/compiler/rustc_hir_analysis/src/check/compare_impl_item.rs
+++ b/compiler/rustc_hir_analysis/src/check/compare_impl_item.rs
@@ -19,7 +19,7 @@ use rustc_middle::ty::{
     self, InternalSubsts, Ty, TypeFoldable, TypeFolder, TypeSuperFoldable, TypeVisitableExt,
 };
 use rustc_middle::ty::{GenericParamDefKind, ToPredicate, TyCtxt};
-use rustc_span::Span;
+use rustc_span::{Span, DUMMY_SP};
 use rustc_trait_selection::traits::error_reporting::TypeErrCtxtExt;
 use rustc_trait_selection::traits::outlives_bounds::InferCtxtExt as _;
 use rustc_trait_selection::traits::{
@@ -767,8 +767,10 @@ pub(super) fn collect_return_position_impl_trait_in_trait_tys<'tcx>(
                 // contains `def_id`'s early-bound regions.
                 let id_substs = InternalSubsts::identity_for_item(tcx, def_id);
                 debug!(?id_substs, ?substs);
-                let map: FxHashMap<ty::GenericArg<'tcx>, ty::GenericArg<'tcx>> =
-                    std::iter::zip(substs, id_substs).collect();
+                let map: FxHashMap<_, _> = std::iter::zip(substs, id_substs)
+                    .skip(tcx.generics_of(trait_m.def_id).count())
+                    .filter_map(|(a, b)| Some((a.as_region()?, b.as_region()?)))
+                    .collect();
                 debug!(?map);
 
                 // NOTE(compiler-errors): RPITITs, like all other RPITs, have early-bound
@@ -793,25 +795,19 @@ pub(super) fn collect_return_position_impl_trait_in_trait_tys<'tcx>(
                 // same generics.
                 let num_trait_substs = trait_to_impl_substs.len();
                 let num_impl_substs = tcx.generics_of(impl_m.container_id(tcx)).params.len();
-                let ty = tcx.fold_regions(ty, |region, _| {
-                    match region.kind() {
-                        // Remap all free regions, which correspond to late-bound regions in the function.
-                        ty::ReFree(_) => {}
-                        // Remap early-bound regions as long as they don't come from the `impl` itself.
-                        ty::ReEarlyBound(ebr) if tcx.parent(ebr.def_id) != impl_m.container_id(tcx) => {}
-                        _ => return region,
-                    }
-                    let Some(ty::ReEarlyBound(e)) = map.get(&region.into()).map(|r| r.expect_region().kind())
-                    else {
-                        return ty::Region::new_error_with_message(tcx, return_span, "expected ReFree to map to ReEarlyBound")
-                    };
-                    ty::Region::new_early_bound(tcx, ty::EarlyBoundRegion {
-                        def_id: e.def_id,
-                        name: e.name,
-                        index: (e.index as usize - num_trait_substs + num_impl_substs) as u32,
-                    })
-                });
-                debug!(%ty);
+                let ty = match ty.try_fold_with(&mut RemapHiddenTyRegions {
+                    tcx,
+                    map,
+                    num_trait_substs,
+                    num_impl_substs,
+                    def_id,
+                    impl_def_id: impl_m.container_id(tcx),
+                    ty,
+                    return_span,
+                }) {
+                    Ok(ty) => ty,
+                    Err(guar) => tcx.ty_error(guar),
+                };
                 collected_tys.insert(def_id, ty::EarlyBinder::bind(ty));
             }
             Err(err) => {
@@ -895,6 +891,97 @@ impl<'tcx> TypeFolder<TyCtxt<'tcx>> for ImplTraitInTraitCollector<'_, 'tcx> {
     }
 }
 
+struct RemapHiddenTyRegions<'tcx> {
+    tcx: TyCtxt<'tcx>,
+    map: FxHashMap<ty::Region<'tcx>, ty::Region<'tcx>>,
+    num_trait_substs: usize,
+    num_impl_substs: usize,
+    def_id: DefId,
+    impl_def_id: DefId,
+    ty: Ty<'tcx>,
+    return_span: Span,
+}
+
+impl<'tcx> ty::FallibleTypeFolder<TyCtxt<'tcx>> for RemapHiddenTyRegions<'tcx> {
+    type Error = ErrorGuaranteed;
+
+    fn interner(&self) -> TyCtxt<'tcx> {
+        self.tcx
+    }
+
+    fn try_fold_ty(&mut self, t: Ty<'tcx>) -> Result<Ty<'tcx>, Self::Error> {
+        if let ty::Alias(ty::Opaque, ty::AliasTy { substs, def_id, .. }) = *t.kind() {
+            let mut mapped_substs = Vec::with_capacity(substs.len());
+            for (arg, v) in std::iter::zip(substs, self.tcx.variances_of(def_id)) {
+                mapped_substs.push(match (arg.unpack(), v) {
+                    // Skip uncaptured opaque substs
+                    (ty::GenericArgKind::Lifetime(_), ty::Bivariant) => arg,
+                    _ => arg.try_fold_with(self)?,
+                });
+            }
+            Ok(self.tcx.mk_opaque(def_id, self.tcx.mk_substs(&mapped_substs)))
+        } else {
+            t.try_super_fold_with(self)
+        }
+    }
+
+    fn try_fold_region(
+        &mut self,
+        region: ty::Region<'tcx>,
+    ) -> Result<ty::Region<'tcx>, Self::Error> {
+        match region.kind() {
+            // Remap all free regions, which correspond to late-bound regions in the function.
+            ty::ReFree(_) => {}
+            // Remap early-bound regions as long as they don't come from the `impl` itself,
+            // in which case we don't really need to renumber them.
+            ty::ReEarlyBound(ebr) if self.tcx.parent(ebr.def_id) != self.impl_def_id => {}
+            _ => return Ok(region),
+        }
+
+        let e = if let Some(region) = self.map.get(&region) {
+            if let ty::ReEarlyBound(e) = region.kind() { e } else { bug!() }
+        } else {
+            let guar = match region.kind() {
+                ty::ReEarlyBound(ty::EarlyBoundRegion { def_id, .. })
+                | ty::ReFree(ty::FreeRegion {
+                    bound_region: ty::BoundRegionKind::BrNamed(def_id, _),
+                    ..
+                }) => {
+                    let return_span = if let ty::Alias(ty::Opaque, opaque_ty) = self.ty.kind() {
+                        self.tcx.def_span(opaque_ty.def_id)
+                    } else {
+                        self.return_span
+                    };
+                    self.tcx
+                        .sess
+                        .struct_span_err(
+                            return_span,
+                            "return type captures more lifetimes than trait definition",
+                        )
+                        .span_label(self.tcx.def_span(def_id), "this lifetime was captured")
+                        .span_note(
+                            self.tcx.def_span(self.def_id),
+                            "hidden type must only reference lifetimes captured by this impl trait",
+                        )
+                        .note(format!("hidden type inferred to be `{}`", self.ty))
+                        .emit()
+                }
+                _ => self.tcx.sess.delay_span_bug(DUMMY_SP, "should've been able to remap region"),
+            };
+            return Err(guar);
+        };
+
+        Ok(ty::Region::new_early_bound(
+            self.tcx,
+            ty::EarlyBoundRegion {
+                def_id: e.def_id,
+                name: e.name,
+                index: (e.index as usize - self.num_trait_substs + self.num_impl_substs) as u32,
+            },
+        ))
+    }
+}
+
 fn report_trait_method_mismatch<'tcx>(
     infcx: &InferCtxt<'tcx>,
     mut cause: ObligationCause<'tcx>,
diff --git a/tests/ui/impl-trait/in-trait/signature-mismatch.current.stderr b/tests/ui/impl-trait/in-trait/signature-mismatch.current.stderr
index eba270af7f0..2db7dd3421c 100644
--- a/tests/ui/impl-trait/in-trait/signature-mismatch.current.stderr
+++ b/tests/ui/impl-trait/in-trait/signature-mismatch.current.stderr
@@ -1,16 +1,18 @@
-error: `impl` item signature doesn't match `trait` item signature
-  --> $DIR/signature-mismatch.rs:17:5
+error: return type captures more lifetimes than trait definition
+  --> $DIR/signature-mismatch.rs:17:47
    |
 LL |     fn async_fn(&self, buff: &[u8]) -> impl Future<Output = Vec<u8>>;
-   |     ----------------------------------------------------------------- expected `fn(&'1 Struct, &'2 [u8]) -> impl Future<Output = Vec<u8>> + '3`
+   |                              - this lifetime was captured
 ...
 LL |     fn async_fn<'a>(&self, buff: &'a [u8]) -> impl Future<Output = Vec<u8>> + 'a {
-   |     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ found `fn(&'1 Struct, &'2 [u8]) -> impl Future<Output = Vec<u8>> + '2`
+   |                                               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
    |
-   = note: expected signature `fn(&'1 Struct, &'2 [u8]) -> impl Future<Output = Vec<u8>> + '3`
-              found signature `fn(&'1 Struct, &'2 [u8]) -> impl Future<Output = Vec<u8>> + '2`
-   = help: the lifetime requirements from the `impl` do not correspond to the requirements in the `trait`
-   = help: verify the lifetime relationships in the `trait` and `impl` between the `self` argument, the other inputs and its output
+note: hidden type must only reference lifetimes captured by this impl trait
+  --> $DIR/signature-mismatch.rs:11:40
+   |
+LL |     fn async_fn(&self, buff: &[u8]) -> impl Future<Output = Vec<u8>>;
+   |                                        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+   = note: hidden type inferred to be `impl Future<Output = Vec<u8>> + '_`
 
 error: aborting due to previous error
 
diff --git a/tests/ui/impl-trait/in-trait/signature-mismatch.next.stderr b/tests/ui/impl-trait/in-trait/signature-mismatch.next.stderr
index eba270af7f0..2db7dd3421c 100644
--- a/tests/ui/impl-trait/in-trait/signature-mismatch.next.stderr
+++ b/tests/ui/impl-trait/in-trait/signature-mismatch.next.stderr
@@ -1,16 +1,18 @@
-error: `impl` item signature doesn't match `trait` item signature
-  --> $DIR/signature-mismatch.rs:17:5
+error: return type captures more lifetimes than trait definition
+  --> $DIR/signature-mismatch.rs:17:47
    |
 LL |     fn async_fn(&self, buff: &[u8]) -> impl Future<Output = Vec<u8>>;
-   |     ----------------------------------------------------------------- expected `fn(&'1 Struct, &'2 [u8]) -> impl Future<Output = Vec<u8>> + '3`
+   |                              - this lifetime was captured
 ...
 LL |     fn async_fn<'a>(&self, buff: &'a [u8]) -> impl Future<Output = Vec<u8>> + 'a {
-   |     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ found `fn(&'1 Struct, &'2 [u8]) -> impl Future<Output = Vec<u8>> + '2`
+   |                                               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
    |
-   = note: expected signature `fn(&'1 Struct, &'2 [u8]) -> impl Future<Output = Vec<u8>> + '3`
-              found signature `fn(&'1 Struct, &'2 [u8]) -> impl Future<Output = Vec<u8>> + '2`
-   = help: the lifetime requirements from the `impl` do not correspond to the requirements in the `trait`
-   = help: verify the lifetime relationships in the `trait` and `impl` between the `self` argument, the other inputs and its output
+note: hidden type must only reference lifetimes captured by this impl trait
+  --> $DIR/signature-mismatch.rs:11:40
+   |
+LL |     fn async_fn(&self, buff: &[u8]) -> impl Future<Output = Vec<u8>>;
+   |                                        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+   = note: hidden type inferred to be `impl Future<Output = Vec<u8>> + '_`
 
 error: aborting due to previous error
 
diff --git a/tests/ui/impl-trait/in-trait/signature-mismatch.rs b/tests/ui/impl-trait/in-trait/signature-mismatch.rs
index 38c902a97a9..c562441774e 100644
--- a/tests/ui/impl-trait/in-trait/signature-mismatch.rs
+++ b/tests/ui/impl-trait/in-trait/signature-mismatch.rs
@@ -15,7 +15,7 @@ pub struct Struct;
 
 impl AsyncTrait for Struct {
     fn async_fn<'a>(&self, buff: &'a [u8]) -> impl Future<Output = Vec<u8>> + 'a {
-        //~^ ERROR `impl` item signature doesn't match `trait` item signature
+        //~^ ERROR return type captures more lifetimes than trait definition
         async move { buff.to_vec() }
     }
 }