Extract borrowck coroutine drop-liveness hack

Author: compiler-errors

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: Locals can be stored across await points in the coroutine,
             // called interior/witness types. Since we do not compute these witnesses until after
             // building MIR, we consider all coroutines to unconditionally require a drop during
@@ -344,8 +344,9 @@ pub fn dtorck_constraint_for_ty_inner<'tcx>(
             // redundant. There is no storage for the resume type, so if it is actually stored
             // in the interior, we'll already detect the need for a drop by checking the witness.
             let typing_env = tcx.erase_regions(typing_env);
-            let needs_drop = args.witness().needs_drop(tcx, typing_env)
-                || args.upvar_tys().iter().any(|ty| ty.needs_drop(tcx, typing_env));
+            let needs_drop = tcx.mir_coroutine_witnesses(def_id).is_some_and(|witness| {
+                witness.field_tys.iter().any(|field| field.ty.needs_drop(tcx, typing_env))
+            }) || args.upvar_tys().iter().any(|ty| ty.needs_drop(tcx, typing_env));
             if needs_drop {
                 constraints.outlives.extend(args.upvar_tys().iter().map(ty::GenericArg::from));
                 constraints.outlives.push(args.resume_ty().into());

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
@@ -131,7 +128,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 +191,26 @@ 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) => {
+                        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) => {