Rollup merge of #111533 - clubby789:drop-tracking-error, r=oli-obk

Handle error body in generator layout

Fixes #111468

I feel like making this query return `Option<GeneratorLayout>` might be better but had some issues with that approach

8 files changed, 49 insertions, 11 deletions

diff --git a/compiler/rustc_hir_analysis/src/check/check.rs b/compiler/rustc_hir_analysis/src/check/check.rs
index 78ffe59679a..d3495d3dbd7 100644
--- a/compiler/rustc_hir_analysis/src/check/check.rs
+++ b/compiler/rustc_hir_analysis/src/check/check.rs
@@ -1514,8 +1514,8 @@ fn opaque_type_cycle_error(
                     }
                     if tcx.sess.opts.unstable_opts.drop_tracking_mir
                         && let DefKind::Generator = tcx.def_kind(closure_def_id)
+                        && let Some(generator_layout) = tcx.mir_generator_witnesses(closure_def_id)
                     {
-                        let generator_layout = tcx.mir_generator_witnesses(closure_def_id);
                         for interior_ty in &generator_layout.field_tys {
                             label_match(interior_ty.ty, interior_ty.source_info.span);
                         }
diff --git a/compiler/rustc_metadata/src/rmeta/encoder.rs b/compiler/rustc_metadata/src/rmeta/encoder.rs
index c1815ae3851..79eb48a1a31 100644
--- a/compiler/rustc_metadata/src/rmeta/encoder.rs
+++ b/compiler/rustc_metadata/src/rmeta/encoder.rs
@@ -1516,8 +1516,11 @@ impl<'a, 'tcx> EncodeContext<'a, 'tcx> {
             if encode_opt {
                 record!(self.tables.optimized_mir[def_id.to_def_id()] <- tcx.optimized_mir(def_id));
 
-                if tcx.sess.opts.unstable_opts.drop_tracking_mir && let DefKind::Generator = self.tcx.def_kind(def_id) {
-                    record!(self.tables.mir_generator_witnesses[def_id.to_def_id()] <- tcx.mir_generator_witnesses(def_id));
+                if tcx.sess.opts.unstable_opts.drop_tracking_mir
+                    && let DefKind::Generator = self.tcx.def_kind(def_id)
+                    && let Some(witnesses) = tcx.mir_generator_witnesses(def_id)
+                {
+                    record!(self.tables.mir_generator_witnesses[def_id.to_def_id()] <- witnesses);
                 }
             }
             if encode_const {
diff --git a/compiler/rustc_middle/src/query/mod.rs b/compiler/rustc_middle/src/query/mod.rs
index ef13a277207..f564f5e99e8 100644
--- a/compiler/rustc_middle/src/query/mod.rs
+++ b/compiler/rustc_middle/src/query/mod.rs
@@ -527,7 +527,7 @@ rustc_queries! {
         }
     }
 
-    query mir_generator_witnesses(key: DefId) -> &'tcx mir::GeneratorLayout<'tcx> {
+    query mir_generator_witnesses(key: DefId) -> &'tcx Option<mir::GeneratorLayout<'tcx>> {
         arena_cache
         desc { |tcx| "generator witness types for `{}`", tcx.def_path_str(key) }
         cache_on_disk_if { key.is_local() }
diff --git a/compiler/rustc_middle/src/ty/util.rs b/compiler/rustc_middle/src/ty/util.rs
index c32ba3edd5c..eb903ebfd99 100644
--- a/compiler/rustc_middle/src/ty/util.rs
+++ b/compiler/rustc_middle/src/ty/util.rs
@@ -668,10 +668,10 @@ impl<'tcx> TyCtxt<'tcx> {
         self,
         def_id: DefId,
     ) -> impl Iterator<Item = ty::EarlyBinder<Ty<'tcx>>> {
-        let generator_layout = &self.mir_generator_witnesses(def_id);
+        let generator_layout = self.mir_generator_witnesses(def_id);
         generator_layout
-            .field_tys
-            .iter()
+            .as_ref()
+            .map_or_else(|| [].iter(), |l| l.field_tys.iter())
             .filter(|decl| !decl.ignore_for_traits)
             .map(|decl| ty::EarlyBinder(decl.ty))
     }
diff --git a/compiler/rustc_mir_transform/src/generator.rs b/compiler/rustc_mir_transform/src/generator.rs
index c9144729145..891e446942e 100644
--- a/compiler/rustc_mir_transform/src/generator.rs
+++ b/compiler/rustc_mir_transform/src/generator.rs
@@ -1397,7 +1397,7 @@ fn create_cases<'tcx>(
 pub(crate) fn mir_generator_witnesses<'tcx>(
     tcx: TyCtxt<'tcx>,
     def_id: LocalDefId,
-) -> GeneratorLayout<'tcx> {
+) -> Option<GeneratorLayout<'tcx>> {
     assert!(tcx.sess.opts.unstable_opts.drop_tracking_mir);
 
     let (body, _) = tcx.mir_promoted(def_id);
@@ -1410,6 +1410,7 @@ pub(crate) fn mir_generator_witnesses<'tcx>(
     // Get the interior types and substs which typeck computed
     let movable = match *gen_ty.kind() {
         ty::Generator(_, _, movability) => movability == hir::Movability::Movable,
+        ty::Error(_) => return None,
         _ => span_bug!(body.span, "unexpected generator type {}", gen_ty),
     };
 
@@ -1425,7 +1426,7 @@ pub(crate) fn mir_generator_witnesses<'tcx>(
 
     check_suspend_tys(tcx, &generator_layout, &body);
 
-    generator_layout
+    Some(generator_layout)
 }
 
 impl<'tcx> MirPass<'tcx> for StateTransform {
diff --git a/compiler/rustc_trait_selection/src/traits/error_reporting/suggestions.rs b/compiler/rustc_trait_selection/src/traits/error_reporting/suggestions.rs
index 49b309abcda..ea17f23434b 100644
--- a/compiler/rustc_trait_selection/src/traits/error_reporting/suggestions.rs
+++ b/compiler/rustc_trait_selection/src/traits/error_reporting/suggestions.rs
@@ -2447,10 +2447,9 @@ impl<'tcx> TypeErrCtxtExt<'tcx> for TypeErrCtxt<'_, 'tcx> {
             && generator_did.is_local()
             // Try to avoid cycles.
             && !generator_within_in_progress_typeck
+            && let Some(generator_info) = self.tcx.mir_generator_witnesses(generator_did)
         {
-            let generator_info = &self.tcx.mir_generator_witnesses(generator_did);
             debug!(?generator_info);
-
             'find_source: for (variant, source_info) in
                 generator_info.variant_fields.iter().zip(&generator_info.variant_source_info)
             {
diff --git a/tests/ui/generator/drop-tracking-error-body.rs b/tests/ui/generator/drop-tracking-error-body.rs
new file mode 100644
index 00000000000..f99d9ab6bf8
--- /dev/null
+++ b/tests/ui/generator/drop-tracking-error-body.rs
@@ -0,0 +1,18 @@
+// compile-flags: -Zdrop-tracking-mir --edition=2021
+
+#![feature(generators)]
+
+pub async fn async_bad_body() {
+    match true {} //~ ERROR non-exhaustive patterns: type `bool` is non-empty
+}
+
+pub fn generator_bad_body() {
+    || {
+        // 'non-exhaustive pattern' only seems to be reported once, so this annotation doesn't work
+        // keep the function around so we can make sure it doesn't ICE
+        match true {}; // ERROR non-exhaustive patterns: type `bool` is non-empty
+        yield ();
+    };
+}
+
+fn main() {}
diff --git a/tests/ui/generator/drop-tracking-error-body.stderr b/tests/ui/generator/drop-tracking-error-body.stderr
new file mode 100644
index 00000000000..28a6892336f
--- /dev/null
+++ b/tests/ui/generator/drop-tracking-error-body.stderr
@@ -0,0 +1,17 @@
+error[E0004]: non-exhaustive patterns: type `bool` is non-empty
+  --> $DIR/drop-tracking-error-body.rs:6:11
+   |
+LL |     match true {}
+   |           ^^^^
+   |
+   = note: the matched value is of type `bool`
+help: ensure that all possible cases are being handled by adding a match arm with a wildcard pattern as shown
+   |
+LL ~     match true {
+LL +         _ => todo!(),
+LL ~     }
+   |
+
+error: aborting due to previous error
+
+For more information about this error, try `rustc --explain E0004`.