From 7242bda3dc6ad638392577fd43d1e30c3671ae74 Mon Sep 17 00:00:00 2001
From: Jonas Schievink <jonasschievink@gmail.com>
Date: Sat, 30 May 2020 03:04:15 +0200
Subject: [PATCH] Be more careful around ty::Error in generators

---
 src/librustc_mir/transform/generator.rs | 69 +++++++++++++++----------
 src/librustc_ty/needs_drop.rs           | 10 +++-
 2 files changed, 51 insertions(+), 28 deletions(-)

diff --git a/src/librustc_mir/transform/generator.rs b/src/librustc_mir/transform/generator.rs
index 461b13c4f6382..60c040c25f97c 100644
--- a/src/librustc_mir/transform/generator.rs
+++ b/src/librustc_mir/transform/generator.rs
@@ -669,40 +669,33 @@ fn compute_storage_conflicts(
     storage_conflicts
 }
 
-fn compute_layout<'tcx>(
+/// Validates the typeck view of the generator against the actual set of types retained between
+/// yield points.
+fn sanitize_witness<'tcx>(
     tcx: TyCtxt<'tcx>,
-    source: MirSource<'tcx>,
+    body: &Body<'tcx>,
+    did: DefId,
+    witness: Ty<'tcx>,
     upvars: &Vec<Ty<'tcx>>,
-    interior: Ty<'tcx>,
-    always_live_locals: &storage::AlwaysLiveLocals,
-    movable: bool,
-    body: &mut Body<'tcx>,
-) -> (
-    FxHashMap<Local, (Ty<'tcx>, VariantIdx, usize)>,
-    GeneratorLayout<'tcx>,
-    IndexVec<BasicBlock, Option<BitSet<Local>>>,
+    retained: &BitSet<Local>,
 ) {
-    // Use a liveness analysis to compute locals which are live across a suspension point
-    let LivenessInfo {
-        live_locals,
-        live_locals_at_suspension_points,
-        storage_conflicts,
-        storage_liveness,
-    } = locals_live_across_suspend_points(tcx, body, source, always_live_locals, movable);
-
-    // Erase regions from the types passed in from typeck so we can compare them with
-    // MIR types
     let allowed_upvars = tcx.erase_regions(upvars);
-    let allowed = match interior.kind {
+    let allowed = match witness.kind {
         ty::GeneratorWitness(s) => tcx.erase_late_bound_regions(&s),
-        _ => bug!(),
+        _ => {
+            tcx.sess.delay_span_bug(
+                body.span,
+                &format!("unexpected generator witness type {:?}", witness.kind),
+            );
+            return;
+        }
     };
 
-    let param_env = tcx.param_env(source.def_id());
+    let param_env = tcx.param_env(did);
 
     for (local, decl) in body.local_decls.iter_enumerated() {
-        // Ignore locals which are internal or not live
-        if !live_locals.contains(local) || decl.internal {
+        // Ignore locals which are internal or not retained between yields.
+        if !retained.contains(local) || decl.internal {
             continue;
         }
         let decl_ty = tcx.normalize_erasing_regions(param_env, decl.ty);
@@ -715,10 +708,34 @@ fn compute_layout<'tcx>(
                 "Broken MIR: generator contains type {} in MIR, \
                        but typeck only knows about {}",
                 decl.ty,
-                interior
+                witness,
             );
         }
     }
+}
+
+fn compute_layout<'tcx>(
+    tcx: TyCtxt<'tcx>,
+    source: MirSource<'tcx>,
+    upvars: &Vec<Ty<'tcx>>,
+    interior: Ty<'tcx>,
+    always_live_locals: &storage::AlwaysLiveLocals,
+    movable: bool,
+    body: &mut Body<'tcx>,
+) -> (
+    FxHashMap<Local, (Ty<'tcx>, VariantIdx, usize)>,
+    GeneratorLayout<'tcx>,
+    IndexVec<BasicBlock, Option<BitSet<Local>>>,
+) {
+    // Use a liveness analysis to compute locals which are live across a suspension point
+    let LivenessInfo {
+        live_locals,
+        live_locals_at_suspension_points,
+        storage_conflicts,
+        storage_liveness,
+    } = locals_live_across_suspend_points(tcx, body, source, always_live_locals, movable);
+
+    sanitize_witness(tcx, body, source.def_id(), interior, upvars, &live_locals);
 
     // Gather live local types and their indices.
     let mut locals = IndexVec::<GeneratorSavedLocal, _>::new();
diff --git a/src/librustc_ty/needs_drop.rs b/src/librustc_ty/needs_drop.rs
index 1b059fa3dbdf0..ca5477ef4a986 100644
--- a/src/librustc_ty/needs_drop.rs
+++ b/src/librustc_ty/needs_drop.rs
@@ -98,7 +98,7 @@ where
                         }
                     }
 
-                    ty::Generator(_, substs, _) => {
+                    ty::Generator(def_id, substs, _) => {
                         let substs = substs.as_generator();
                         for upvar_ty in substs.upvar_tys() {
                             queue_type(self, upvar_ty);
@@ -107,7 +107,13 @@ where
                         let witness = substs.witness();
                         let interior_tys = match &witness.kind {
                             ty::GeneratorWitness(tys) => tcx.erase_late_bound_regions(tys),
-                            _ => bug!(),
+                            _ => {
+                                tcx.sess.delay_span_bug(
+                                    tcx.hir().span_if_local(def_id).unwrap_or(DUMMY_SP),
+                                    &format!("unexpected generator witness type {:?}", witness),
+                                );
+                                return Some(Err(AlwaysRequiresDrop));
+                            }
                         };
 
                         for interior_ty in interior_tys {