Rollup merge of #145338 - lcnr:coroutine-witness-yikes, r=compiler-errors

actually provide the correct args to coroutine witnesses

rust-lang/rust#145194 accidentally provided all arguments of the closure to the witness, but the witness only takes the generic parameters of the defining scope: https://github.com/rust-lang/rust/blob/216cdb7b22b637cef75b7225c642cb7587192643/compiler/rustc_hir_typeck/src/closure.rs#L164

Fixes rust-lang/rust#145288

5 files changed, 85 insertions, 63 deletions

diff --git a/compiler/rustc_middle/src/ty/sty.rs b/compiler/rustc_middle/src/ty/sty.rs
index 72474a60566..755fc68d86f 100644
--- a/compiler/rustc_middle/src/ty/sty.rs
+++ b/compiler/rustc_middle/src/ty/sty.rs
@@ -821,10 +821,38 @@ impl<'tcx> Ty<'tcx> {
     #[inline]
     pub fn new_coroutine_witness(
         tcx: TyCtxt<'tcx>,
-        id: DefId,
+        def_id: DefId,
         args: GenericArgsRef<'tcx>,
     ) -> Ty<'tcx> {
-        Ty::new(tcx, CoroutineWitness(id, args))
+        if cfg!(debug_assertions) {
+            tcx.debug_assert_args_compatible(tcx.typeck_root_def_id(def_id), args);
+        }
+        Ty::new(tcx, CoroutineWitness(def_id, args))
+    }
+
+    pub fn new_coroutine_witness_for_coroutine(
+        tcx: TyCtxt<'tcx>,
+        def_id: DefId,
+        coroutine_args: GenericArgsRef<'tcx>,
+    ) -> Ty<'tcx> {
+        tcx.debug_assert_args_compatible(def_id, coroutine_args);
+        // HACK: Coroutine witness types are lifetime erased, so they
+        // never reference any lifetime args from the coroutine. We erase
+        // the regions here since we may get into situations where a
+        // coroutine is recursively contained within itself, leading to
+        // witness types that differ by region args. This means that
+        // cycle detection in fulfillment will not kick in, which leads
+        // to unnecessary overflows in async code. See the issue:
+        // <https://github.com/rust-lang/rust/issues/145151>.
+        let args =
+            ty::GenericArgs::for_item(tcx, tcx.typeck_root_def_id(def_id), |def, _| {
+                match def.kind {
+                    ty::GenericParamDefKind::Lifetime => tcx.lifetimes.re_erased.into(),
+                    ty::GenericParamDefKind::Type { .. }
+                    | ty::GenericParamDefKind::Const { .. } => coroutine_args[def.index as usize],
+                }
+            });
+        Ty::new_coroutine_witness(tcx, def_id, args)
     }
 
     // misc
@@ -985,6 +1013,14 @@ impl<'tcx> rustc_type_ir::inherent::Ty<TyCtxt<'tcx>> for Ty<'tcx> {
         Ty::new_coroutine_witness(interner, def_id, args)
     }
 
+    fn new_coroutine_witness_for_coroutine(
+        interner: TyCtxt<'tcx>,
+        def_id: DefId,
+        coroutine_args: ty::GenericArgsRef<'tcx>,
+    ) -> Self {
+        Ty::new_coroutine_witness_for_coroutine(interner, def_id, coroutine_args)
+    }
+
     fn new_ptr(interner: TyCtxt<'tcx>, ty: Self, mutbl: hir::Mutability) -> Self {
         Ty::new_ptr(interner, ty, mutbl)
     }
diff --git a/compiler/rustc_next_trait_solver/src/solve/assembly/structural_traits.rs b/compiler/rustc_next_trait_solver/src/solve/assembly/structural_traits.rs
index faa86734d08..d25e74e7335 100644
--- a/compiler/rustc_next_trait_solver/src/solve/assembly/structural_traits.rs
+++ b/compiler/rustc_next_trait_solver/src/solve/assembly/structural_traits.rs
@@ -75,17 +75,10 @@ where
             Ok(ty::Binder::dummy(vec![args.as_coroutine_closure().tupled_upvars_ty()]))
         }
 
-        ty::Coroutine(def_id, args) => {
-            let coroutine_args = args.as_coroutine();
-            Ok(ty::Binder::dummy(vec![
-                coroutine_args.tupled_upvars_ty(),
-                Ty::new_coroutine_witness(
-                    ecx.cx(),
-                    def_id,
-                    ecx.cx().mk_args(coroutine_args.parent_args().as_slice()),
-                ),
-            ]))
-        }
+        ty::Coroutine(def_id, args) => Ok(ty::Binder::dummy(vec![
+            args.as_coroutine().tupled_upvars_ty(),
+            Ty::new_coroutine_witness_for_coroutine(ecx.cx(), def_id, args),
+        ])),
 
         ty::CoroutineWitness(def_id, args) => Ok(ecx
             .cx()
@@ -251,14 +244,9 @@ where
             Movability::Static => Err(NoSolution),
             Movability::Movable => {
                 if ecx.cx().features().coroutine_clone() {
-                    let coroutine = args.as_coroutine();
                     Ok(ty::Binder::dummy(vec![
-                        coroutine.tupled_upvars_ty(),
-                        Ty::new_coroutine_witness(
-                            ecx.cx(),
-                            def_id,
-                            ecx.cx().mk_args(coroutine.parent_args().as_slice()),
-                        ),
+                        args.as_coroutine().tupled_upvars_ty(),
+                        Ty::new_coroutine_witness_for_coroutine(ecx.cx(), def_id, args),
                     ]))
                 } else {
                     Err(NoSolution)
diff --git a/compiler/rustc_trait_selection/src/traits/select/mod.rs b/compiler/rustc_trait_selection/src/traits/select/mod.rs
index 468c42abf48..1dd31990ab7 100644
--- a/compiler/rustc_trait_selection/src/traits/select/mod.rs
+++ b/compiler/rustc_trait_selection/src/traits/select/mod.rs
@@ -4,9 +4,9 @@
 
 use std::assert_matches::assert_matches;
 use std::cell::{Cell, RefCell};
+use std::cmp;
 use std::fmt::{self, Display};
 use std::ops::ControlFlow;
-use std::{cmp, iter};
 
 use hir::def::DefKind;
 use rustc_data_structures::fx::{FxIndexMap, FxIndexSet};
@@ -2185,32 +2185,23 @@ impl<'tcx> SelectionContext<'_, 'tcx> {
                 ty::Binder::dummy(vec![ty])
             }
 
-            ty::Coroutine(coroutine_def_id, args) => {
-                match self.tcx().coroutine_movability(coroutine_def_id) {
-                    hir::Movability::Static => {
-                        unreachable!("tried to assemble `Sized` for static coroutine")
-                    }
-                    hir::Movability::Movable => {
-                        if self.tcx().features().coroutine_clone() {
-                            ty::Binder::dummy(
-                                args.as_coroutine()
-                                    .upvar_tys()
-                                    .iter()
-                                    .chain([Ty::new_coroutine_witness(
-                                        self.tcx(),
-                                        coroutine_def_id,
-                                        self.tcx().mk_args(args.as_coroutine().parent_args()),
-                                    )])
-                                    .collect::<Vec<_>>(),
-                            )
-                        } else {
-                            unreachable!(
-                                "tried to assemble `Sized` for coroutine without enabled feature"
-                            )
-                        }
+            ty::Coroutine(def_id, args) => match self.tcx().coroutine_movability(def_id) {
+                hir::Movability::Static => {
+                    unreachable!("tried to assemble `Clone` for static coroutine")
+                }
+                hir::Movability::Movable => {
+                    if self.tcx().features().coroutine_clone() {
+                        ty::Binder::dummy(vec![
+                            args.as_coroutine().tupled_upvars_ty(),
+                            Ty::new_coroutine_witness_for_coroutine(self.tcx(), def_id, args),
+                        ])
+                    } else {
+                        unreachable!(
+                            "tried to assemble `Clone` for coroutine without enabled feature"
+                        )
                     }
                 }
-            }
+            },
 
             ty::CoroutineWitness(def_id, args) => self
                 .infcx
@@ -2334,25 +2325,9 @@ impl<'tcx> SelectionContext<'_, 'tcx> {
             ty::Coroutine(def_id, args) => {
                 let ty = self.infcx.shallow_resolve(args.as_coroutine().tupled_upvars_ty());
                 let tcx = self.tcx();
-                let witness = Ty::new_coroutine_witness(
-                    tcx,
-                    def_id,
-                    ty::GenericArgs::for_item(tcx, def_id, |def, _| match def.kind {
-                        // HACK: Coroutine witnesse types are lifetime erased, so they
-                        // never reference any lifetime args from the coroutine. We erase
-                        // the regions here since we may get into situations where a
-                        // coroutine is recursively contained within itself, leading to
-                        // witness types that differ by region args. This means that
-                        // cycle detection in fulfillment will not kick in, which leads
-                        // to unnecessary overflows in async code. See the issue:
-                        // <https://github.com/rust-lang/rust/issues/145151>.
-                        ty::GenericParamDefKind::Lifetime => tcx.lifetimes.re_erased.into(),
-                        ty::GenericParamDefKind::Type { .. }
-                        | ty::GenericParamDefKind::Const { .. } => args[def.index as usize],
-                    }),
-                );
+                let witness = Ty::new_coroutine_witness_for_coroutine(tcx, def_id, args);
                 ty::Binder::dummy(AutoImplConstituents {
-                    types: [ty].into_iter().chain(iter::once(witness)).collect(),
+                    types: vec![ty, witness],
                     assumptions: vec![],
                 })
             }
diff --git a/compiler/rustc_type_ir/src/inherent.rs b/compiler/rustc_type_ir/src/inherent.rs
index 1a6c99ce7de..569570b5783 100644
--- a/compiler/rustc_type_ir/src/inherent.rs
+++ b/compiler/rustc_type_ir/src/inherent.rs
@@ -91,6 +91,12 @@ pub trait Ty<I: Interner<Ty = Self>>:
 
     fn new_coroutine_witness(interner: I, def_id: I::DefId, args: I::GenericArgs) -> Self;
 
+    fn new_coroutine_witness_for_coroutine(
+        interner: I,
+        def_id: I::DefId,
+        coroutine_args: I::GenericArgs,
+    ) -> Self;
+
     fn new_ptr(interner: I, ty: Self, mutbl: Mutability) -> Self;
 
     fn new_ref(interner: I, region: I::Region, ty: Self, mutbl: Mutability) -> Self;
diff --git a/tests/ui/async-await/recursive-async-auto-trait-overflow-only-parent-args.rs b/tests/ui/async-await/recursive-async-auto-trait-overflow-only-parent-args.rs
new file mode 100644
index 00000000000..9681f66412a
--- /dev/null
+++ b/tests/ui/async-await/recursive-async-auto-trait-overflow-only-parent-args.rs
@@ -0,0 +1,17 @@
+// Regression test for #145288. This is the same issue as #145151
+// which we fixed in #145194. However in that PR we accidentally created
+// a `CoroutineWitness` which referenced all generic arguments of the
+// coroutine, including upvars and the signature.
+
+//@ edition: 2024
+//@ check-pass
+
+async fn process<'a>(x: &'a u32) {
+    Box::pin(process(x)).await;
+}
+
+fn require_send(_: impl Send) {}
+
+fn main() {
+    require_send(process(&1));
+}