Fix suggestions for match non-exhaustiveness

5 files changed, 54 insertions, 34 deletions

diff --git a/compiler/rustc_middle/src/thir.rs b/compiler/rustc_middle/src/thir.rs
index 05f6fbbbfa3..95e8509c2d3 100644
--- a/compiler/rustc_middle/src/thir.rs
+++ b/compiler/rustc_middle/src/thir.rs
@@ -12,7 +12,7 @@ use rustc_ast::{InlineAsmOptions, InlineAsmTemplatePiece};
 use rustc_errors::{DiagArgValue, IntoDiagArg};
 use rustc_hir as hir;
 use rustc_hir::def_id::DefId;
-use rustc_hir::{BindingAnnotation, ByRef, RangeEnd};
+use rustc_hir::{BindingAnnotation, ByRef, MatchSource, RangeEnd};
 use rustc_index::newtype_index;
 use rustc_index::IndexVec;
 use rustc_middle::middle::region;
@@ -358,6 +358,7 @@ pub enum ExprKind<'tcx> {
         scrutinee: ExprId,
         scrutinee_hir_id: hir::HirId,
         arms: Box<[ArmId]>,
+        match_source: MatchSource,
     },
     /// A block.
     Block {
diff --git a/compiler/rustc_mir_build/src/errors.rs b/compiler/rustc_mir_build/src/errors.rs
index 848da56f981..26f10fdd333 100644
--- a/compiler/rustc_mir_build/src/errors.rs
+++ b/compiler/rustc_mir_build/src/errors.rs
@@ -456,8 +456,8 @@ pub enum UnusedUnsafeEnclosing {
 
 pub(crate) struct NonExhaustivePatternsTypeNotEmpty<'p, 'tcx, 'm> {
     pub cx: &'m RustcPatCtxt<'p, 'tcx>,
-    pub expr_span: Span,
-    pub span: Span,
+    pub scrut_span: Span,
+    pub braces_span: Option<Span>,
     pub ty: Ty<'tcx>,
 }
 
@@ -465,7 +465,7 @@ impl<'a, G: EmissionGuarantee> Diagnostic<'a, G> for NonExhaustivePatternsTypeNo
     fn into_diag(self, dcx: &'a DiagCtxt, level: Level) -> Diag<'_, G> {
         let mut diag =
             Diag::new(dcx, level, fluent::mir_build_non_exhaustive_patterns_type_not_empty);
-        diag.span(self.span);
+        diag.span(self.scrut_span);
         diag.code(E0004);
         let peeled_ty = self.ty.peel_refs();
         diag.arg("ty", self.ty);
@@ -502,26 +502,19 @@ impl<'a, G: EmissionGuarantee> Diagnostic<'a, G> for NonExhaustivePatternsTypeNo
             }
         }
 
-        let mut suggestion = None;
         let sm = self.cx.tcx.sess.source_map();
-        if self.span.eq_ctxt(self.expr_span) {
+        if let Some(braces_span) = self.braces_span {
             // Get the span for the empty match body `{}`.
-            let (indentation, more) = if let Some(snippet) = sm.indentation_before(self.span) {
+            let (indentation, more) = if let Some(snippet) = sm.indentation_before(self.scrut_span)
+            {
                 (format!("\n{snippet}"), "    ")
             } else {
                 (" ".to_string(), "")
             };
-            suggestion = Some((
-                self.span.shrink_to_hi().with_hi(self.expr_span.hi()),
-                format!(" {{{indentation}{more}_ => todo!(),{indentation}}}",),
-            ));
-        }
-
-        if let Some((span, sugg)) = suggestion {
             diag.span_suggestion_verbose(
-                span,
+                braces_span,
                 fluent::mir_build_suggestion,
-                sugg,
+                format!(" {{{indentation}{more}_ => todo!(),{indentation}}}"),
                 Applicability::HasPlaceholders,
             );
         } else {
diff --git a/compiler/rustc_mir_build/src/thir/cx/expr.rs b/compiler/rustc_mir_build/src/thir/cx/expr.rs
index 1e508ffc1e7..c697e16217b 100644
--- a/compiler/rustc_mir_build/src/thir/cx/expr.rs
+++ b/compiler/rustc_mir_build/src/thir/cx/expr.rs
@@ -716,10 +716,11 @@ impl<'tcx> Cx<'tcx> {
                 then: self.mirror_expr(then),
                 else_opt: else_opt.map(|el| self.mirror_expr(el)),
             },
-            hir::ExprKind::Match(discr, arms, _) => ExprKind::Match {
+            hir::ExprKind::Match(discr, arms, match_source) => ExprKind::Match {
                 scrutinee: self.mirror_expr(discr),
                 scrutinee_hir_id: discr.hir_id,
                 arms: arms.iter().map(|a| self.convert_arm(a)).collect(),
+                match_source,
             },
             hir::ExprKind::Loop(body, ..) => {
                 let block_ty = self.typeck_results().node_type(body.hir_id);
diff --git a/compiler/rustc_mir_build/src/thir/pattern/check_match.rs b/compiler/rustc_mir_build/src/thir/pattern/check_match.rs
index 3a688a14dd5..a3e6c2abc51 100644
--- a/compiler/rustc_mir_build/src/thir/pattern/check_match.rs
+++ b/compiler/rustc_mir_build/src/thir/pattern/check_match.rs
@@ -144,16 +144,8 @@ impl<'p, 'tcx> Visitor<'p, 'tcx> for MatchVisitor<'p, 'tcx> {
                 });
                 return;
             }
-            ExprKind::Match { scrutinee, scrutinee_hir_id, box ref arms } => {
-                let source = match ex.span.desugaring_kind() {
-                    Some(DesugaringKind::ForLoop) => hir::MatchSource::ForLoopDesugar,
-                    Some(DesugaringKind::QuestionMark) => {
-                        hir::MatchSource::TryDesugar(scrutinee_hir_id)
-                    }
-                    Some(DesugaringKind::Await) => hir::MatchSource::AwaitDesugar,
-                    _ => hir::MatchSource::Normal,
-                };
-                self.check_match(scrutinee, arms, source, ex.span);
+            ExprKind::Match { scrutinee, scrutinee_hir_id: _, box ref arms, match_source } => {
+                self.check_match(scrutinee, arms, match_source, ex.span);
             }
             ExprKind::Let { box ref pat, expr } => {
                 self.check_let(pat, Some(expr), ex.span);
@@ -505,8 +497,41 @@ impl<'p, 'tcx> MatchVisitor<'p, 'tcx> {
                     None,
                 );
             } else {
+                // span after scrutinee, or after `.match`. That is, the braces, arms,
+                // and any whitespace preceding the braces.
+                let braces_span = match source {
+                    hir::MatchSource::Normal => scrut
+                        .span
+                        .find_ancestor_in_same_ctxt(expr_span)
+                        .map(|scrut_span| scrut_span.shrink_to_hi().with_hi(expr_span.hi())),
+                    hir::MatchSource::Postfix => {
+                        // This is horrendous, and we should deal with it by just
+                        // stashing the span of the braces somewhere (like in the match source).
+                        scrut.span.find_ancestor_in_same_ctxt(expr_span).and_then(|scrut_span| {
+                            let sm = self.tcx.sess.source_map();
+                            let brace_span = sm.span_extend_to_next_char(scrut_span, '{', true);
+                            if sm.span_to_snippet(sm.next_point(brace_span)).as_deref() == Ok("{") {
+                                let sp = brace_span.shrink_to_hi().with_hi(expr_span.hi());
+                                // We also need to extend backwards for whitespace
+                                sm.span_extend_prev_while(sp, |c| c.is_whitespace()).ok()
+                            } else {
+                                None
+                            }
+                        })
+                    }
+                    hir::MatchSource::ForLoopDesugar
+                    | hir::MatchSource::TryDesugar(_)
+                    | hir::MatchSource::AwaitDesugar
+                    | hir::MatchSource::FormatArgs => None,
+                };
                 self.error = Err(report_non_exhaustive_match(
-                    &cx, self.thir, scrut.ty, scrut.span, witnesses, arms, expr_span,
+                    &cx,
+                    self.thir,
+                    scrut.ty,
+                    scrut.span,
+                    witnesses,
+                    arms,
+                    braces_span,
                 ));
             }
         }
@@ -929,7 +954,7 @@ fn report_non_exhaustive_match<'p, 'tcx>(
     sp: Span,
     witnesses: Vec<WitnessPat<'p, 'tcx>>,
     arms: &[ArmId],
-    expr_span: Span,
+    braces_span: Option<Span>,
 ) -> ErrorGuaranteed {
     let is_empty_match = arms.is_empty();
     let non_empty_enum = match scrut_ty.kind() {
@@ -941,8 +966,8 @@ fn report_non_exhaustive_match<'p, 'tcx>(
     if is_empty_match && !non_empty_enum {
         return cx.tcx.dcx().emit_err(NonExhaustivePatternsTypeNotEmpty {
             cx,
-            expr_span,
-            span: sp,
+            scrut_span: sp,
+            braces_span,
             ty: scrut_ty,
         });
     }
@@ -1028,7 +1053,7 @@ fn report_non_exhaustive_match<'p, 'tcx>(
     let mut suggestion = None;
     let sm = cx.tcx.sess.source_map();
     match arms {
-        [] if sp.eq_ctxt(expr_span) => {
+        [] if let Some(braces_span) = braces_span => {
             // Get the span for the empty match body `{}`.
             let (indentation, more) = if let Some(snippet) = sm.indentation_before(sp) {
                 (format!("\n{snippet}"), "    ")
@@ -1036,7 +1061,7 @@ fn report_non_exhaustive_match<'p, 'tcx>(
                 (" ".to_string(), "")
             };
             suggestion = Some((
-                sp.shrink_to_hi().with_hi(expr_span.hi()),
+                braces_span,
                 format!(" {{{indentation}{more}{suggested_arm},{indentation}}}",),
             ));
         }
diff --git a/tests/ui/match/postfix-match/match-after-as.stderr b/tests/ui/match/postfix-match/match-after-as.stderr
index 99e42991c52..68e4762b8d9 100644
--- a/tests/ui/match/postfix-match/match-after-as.stderr
+++ b/tests/ui/match/postfix-match/match-after-as.stderr
@@ -18,7 +18,7 @@ LL |     1 as i32.match {};
    = note: the matched value is of type `i32`
 help: ensure that all possible cases are being handled by adding a match arm with a wildcard pattern as shown
    |
-LL ~     1 as i32 {
+LL ~     1 as i32.match {
 LL +         _ => todo!(),
 LL ~     };
    |