Extract borrowck coroutine drop-liveness hack

2 files changed, 31 insertions, 19 deletions

diff --git a/compiler/rustc_trait_selection/src/traits/query/dropck_outlives.rs b/compiler/rustc_trait_selection/src/traits/query/dropck_outlives.rs
index 38cfdcdc22d..b1b331d1b61 100644
--- a/compiler/rustc_trait_selection/src/traits/query/dropck_outlives.rs
+++ b/compiler/rustc_trait_selection/src/traits/query/dropck_outlives.rs
@@ -318,7 +318,7 @@ pub fn dtorck_constraint_for_ty_inner<'tcx>(
             })
         }
 
-        ty::Coroutine(_, args) => {
+        ty::Coroutine(def_id, args) => {
             // rust-lang/rust#49918: types can be constructed, stored
             // in the interior, and sit idle when coroutine yields
             // (and is subsequently dropped).
@@ -346,7 +346,10 @@ pub fn dtorck_constraint_for_ty_inner<'tcx>(
             // While we conservatively assume that all coroutines require drop
             // to avoid query cycles during MIR building, we can check the actual
             // witness during borrowck to avoid unnecessary liveness constraints.
-            if args.witness().needs_drop(tcx, tcx.erase_regions(typing_env)) {
+            let typing_env = tcx.erase_regions(typing_env);
+            if tcx.mir_coroutine_witnesses(def_id).is_some_and(|witness| {
+                witness.field_tys.iter().any(|field| field.ty.needs_drop(tcx, typing_env))
+            }) {
                 constraints.outlives.extend(args.upvar_tys().iter().map(ty::GenericArg::from));
                 constraints.outlives.push(args.resume_ty().into());
             }
diff --git a/compiler/rustc_ty_utils/src/needs_drop.rs b/compiler/rustc_ty_utils/src/needs_drop.rs
index c3b04c20f4b..13d56889bd1 100644
--- a/compiler/rustc_ty_utils/src/needs_drop.rs
+++ b/compiler/rustc_ty_utils/src/needs_drop.rs
@@ -101,9 +101,6 @@ fn has_significant_drop_raw<'tcx>(
 struct NeedsDropTypes<'tcx, F> {
     tcx: TyCtxt<'tcx>,
     typing_env: ty::TypingEnv<'tcx>,
-    /// Whether to reveal coroutine witnesses, this is set
-    /// to `false` unless we compute `needs_drop` for a coroutine witness.
-    reveal_coroutine_witnesses: bool,
     query_ty: Ty<'tcx>,
     seen_tys: FxHashSet<Ty<'tcx>>,
     /// A stack of types left to process, and the recursion depth when we
@@ -115,6 +112,15 @@ struct NeedsDropTypes<'tcx, F> {
     adt_components: F,
     /// Set this to true if an exhaustive list of types involved in
     /// drop obligation is requested.
+    // FIXME: Calling this bool `exhaustive` is confusing and possibly a footgun,
+    // since it does two things: It makes the iterator yield *all* of the types
+    // that need drop, and it also affects the computation of the drop components
+    // on `Coroutine`s. The latter is somewhat confusing, and probably should be
+    // a function of `typing_env`. See the HACK comment below for why this is
+    // necessary. If this isn't possible, then we probably should turn this into
+    // a `NeedsDropMode` so that we can have a variant like `CollectAllSignificantDrops`,
+    // which will more accurately indicate that we want *all* of the *significant*
+    // drops, which are the two important behavioral changes toggled by this bool.
     exhaustive: bool,
 }
 
@@ -131,7 +137,6 @@ impl<'tcx, F> NeedsDropTypes<'tcx, F> {
         Self {
             tcx,
             typing_env,
-            reveal_coroutine_witnesses: exhaustive,
             seen_tys,
             query_ty: ty,
             unchecked_tys: vec![(ty, 0)],
@@ -195,23 +200,27 @@ where
                     // for the coroutine witness and check whether any of the contained types
                     // need to be dropped, and only require the captured types to be live
                     // if they do.
-                    ty::Coroutine(_, args) => {
-                        if self.reveal_coroutine_witnesses {
-                            queue_type(self, args.as_coroutine().witness());
+                    ty::Coroutine(def_id, args) => {
+                        // FIXME: See FIXME on `exhaustive` field above.
+                        if self.exhaustive {
+                            for upvar in args.as_coroutine().upvar_tys() {
+                                queue_type(self, upvar);
+                            }
+                            queue_type(self, args.as_coroutine().resume_ty());
+                            if let Some(witness) = tcx.mir_coroutine_witnesses(def_id) {
+                                for field_ty in &witness.field_tys {
+                                    queue_type(
+                                        self,
+                                        EarlyBinder::bind(field_ty.ty).instantiate(tcx, args),
+                                    );
+                                }
+                            }
                         } else {
                             return Some(self.always_drop_component(ty));
                         }
                     }
-                    ty::CoroutineWitness(def_id, args) => {
-                        if let Some(witness) = tcx.mir_coroutine_witnesses(def_id) {
-                            self.reveal_coroutine_witnesses = true;
-                            for field_ty in &witness.field_tys {
-                                queue_type(
-                                    self,
-                                    EarlyBinder::bind(field_ty.ty).instantiate(tcx, args),
-                                );
-                            }
-                        }
+                    ty::CoroutineWitness(..) => {
+                        unreachable!("witness should be handled in parent");
                     }
 
                     ty::UnsafeBinder(bound_ty) => {