And for patterns too

5 files changed, 184 insertions, 119 deletions

diff --git a/compiler/rustc_infer/src/infer/error_reporting/mod.rs b/compiler/rustc_infer/src/infer/error_reporting/mod.rs
index a8fc306b286..5c1ab2bf588 100644
--- a/compiler/rustc_infer/src/infer/error_reporting/mod.rs
+++ b/compiler/rustc_infer/src/infer/error_reporting/mod.rs
@@ -614,13 +614,16 @@ impl<'a, 'tcx> InferCtxt<'a, 'tcx> {
                 err.span_label(span, "expected due to this");
             }
             ObligationCauseCode::MatchExpressionArm(box MatchExpressionArmCause {
-                semi_span,
+                arm_block_id,
+                arm_span,
+                arm_ty,
+                prior_arm_block_id,
+                prior_arm_span,
+                prior_arm_ty,
                 source,
                 ref prior_arms,
-                last_ty,
                 scrut_hir_id,
                 opt_suggest_box_span,
-                arm_span,
                 scrut_span,
                 ..
             }) => match source {
@@ -651,10 +654,10 @@ impl<'a, 'tcx> InferCtxt<'a, 'tcx> {
                     }
                 }
                 _ => {
-                    // `last_ty` can be `!`, `expected` will have better info when present.
+                    // `prior_arm_ty` can be `!`, `expected` will have better info when present.
                     let t = self.resolve_vars_if_possible(match exp_found {
                         Some(ty::error::ExpectedFound { expected, .. }) => expected,
-                        _ => last_ty,
+                        _ => prior_arm_ty,
                     });
                     let source_map = self.tcx.sess.source_map();
                     let mut any_multiline_arm = source_map.is_multiline(arm_span);
@@ -679,37 +682,15 @@ impl<'a, 'tcx> InferCtxt<'a, 'tcx> {
                     };
                     let msg = "`match` arms have incompatible types";
                     err.span_label(outer_error_span, msg);
-                    if let Some((sp, boxed)) = semi_span {
-                        if let (StatementAsExpression::NeedsBoxing, [.., prior_arm]) =
-                            (boxed, &prior_arms[..])
-                        {
-                            err.multipart_suggestion(
-                                "consider removing this semicolon and boxing the expressions",
-                                vec![
-                                    (prior_arm.shrink_to_lo(), "Box::new(".to_string()),
-                                    (prior_arm.shrink_to_hi(), ")".to_string()),
-                                    (arm_span.shrink_to_lo(), "Box::new(".to_string()),
-                                    (arm_span.shrink_to_hi(), ")".to_string()),
-                                    (sp, String::new()),
-                                ],
-                                Applicability::HasPlaceholders,
-                            );
-                        } else if matches!(boxed, StatementAsExpression::NeedsBoxing) {
-                            err.span_suggestion_short(
-                                sp,
-                                "consider removing this semicolon and boxing the expressions",
-                                "",
-                                Applicability::MachineApplicable,
-                            );
-                        } else {
-                            err.span_suggestion_short(
-                                sp,
-                                "consider removing this semicolon",
-                                "",
-                                Applicability::MachineApplicable,
-                            );
-                        }
-                    }
+                    self.suggest_remove_semi_or_return_binding(
+                        err,
+                        prior_arm_block_id,
+                        prior_arm_ty,
+                        prior_arm_span,
+                        arm_block_id,
+                        arm_ty,
+                        arm_span,
+                    );
                     if let Some(ret_sp) = opt_suggest_box_span {
                         // Get return type span and point to it.
                         self.suggest_boxing_for_return_impl_trait(
@@ -734,48 +715,15 @@ impl<'a, 'tcx> InferCtxt<'a, 'tcx> {
                 if let Some(sp) = outer_span {
                     err.span_label(sp, "`if` and `else` have incompatible types");
                 }
-                let semicolon = if let hir::Node::Block(blk) = self.tcx.hir().get(then_id)
-                    && let Some(remove_semicolon) = self.could_remove_semicolon(blk, else_ty)
-                {
-                    Some(remove_semicolon)
-                } else if let hir::Node::Block(blk) = self.tcx.hir().get(else_id)
-                    && let Some(remove_semicolon) = self.could_remove_semicolon(blk, then_ty)
-                {
-                    Some(remove_semicolon)
-                } else {
-                    None
-                };
-                if let Some((sp, boxed)) = semicolon {
-                    if matches!(boxed, StatementAsExpression::NeedsBoxing) {
-                        err.multipart_suggestion(
-                            "consider removing this semicolon and boxing the expression",
-                            vec![
-                                (then_span.shrink_to_lo(), "Box::new(".to_string()),
-                                (then_span.shrink_to_hi(), ")".to_string()),
-                                (else_span.shrink_to_lo(), "Box::new(".to_string()),
-                                (else_span.shrink_to_hi(), ")".to_string()),
-                                (sp, String::new()),
-                            ],
-                            Applicability::MachineApplicable,
-                        );
-                    } else {
-                        err.span_suggestion_short(
-                            sp,
-                            "consider removing this semicolon",
-                            "",
-                            Applicability::MachineApplicable,
-                        );
-                    }
-                } else {
-                    let suggested = if let hir::Node::Block(blk) = self.tcx.hir().get(then_id) {
-                        self.consider_returning_binding(blk, else_ty, err)
-                    } else {
-                        false
-                    };
-                    if !suggested && let hir::Node::Block(blk) = self.tcx.hir().get(else_id) {
-                        self.consider_returning_binding(blk, then_ty, err);
-                    }
-                }
+                self.suggest_remove_semi_or_return_binding(
+                    err,
+                    Some(then_id),
+                    then_ty,
+                    then_span,
+                    Some(else_id),
+                    else_ty,
+                    else_span,
+                );
                 if let Some(ret_sp) = opt_suggest_box_span {
                     self.suggest_boxing_for_return_impl_trait(
                         err,
@@ -800,6 +748,69 @@ impl<'a, 'tcx> InferCtxt<'a, 'tcx> {
         }
     }
 
+    fn suggest_remove_semi_or_return_binding(
+        &self,
+        err: &mut Diagnostic,
+        first_id: Option<hir::HirId>,
+        first_ty: Ty<'tcx>,
+        first_span: Span,
+        second_id: Option<hir::HirId>,
+        second_ty: Ty<'tcx>,
+        second_span: Span,
+    ) {
+        let semicolon =
+            if let Some(first_id) = first_id
+                && let hir::Node::Block(blk) = self.tcx.hir().get(first_id)
+                && let Some(remove_semicolon) = self.could_remove_semicolon(blk, second_ty)
+            {
+                Some(remove_semicolon)
+            } else if let Some(second_id) = second_id
+                && let hir::Node::Block(blk) = self.tcx.hir().get(second_id)
+                && let Some(remove_semicolon) = self.could_remove_semicolon(blk, first_ty)
+            {
+                Some(remove_semicolon)
+            } else {
+                None
+            };
+        if let Some((sp, boxed)) = semicolon {
+            if matches!(boxed, StatementAsExpression::NeedsBoxing) {
+                err.multipart_suggestion(
+                    "consider removing this semicolon and boxing the expressions",
+                    vec![
+                        (first_span.shrink_to_lo(), "Box::new(".to_string()),
+                        (first_span.shrink_to_hi(), ")".to_string()),
+                        (second_span.shrink_to_lo(), "Box::new(".to_string()),
+                        (second_span.shrink_to_hi(), ")".to_string()),
+                        (sp, String::new()),
+                    ],
+                    Applicability::MachineApplicable,
+                );
+            } else {
+                err.span_suggestion_short(
+                    sp,
+                    "consider removing this semicolon",
+                    "",
+                    Applicability::MachineApplicable,
+                );
+            }
+        } else {
+            let suggested =
+                if let Some(first_id) = first_id
+                    && let hir::Node::Block(blk) = self.tcx.hir().get(first_id)
+                {
+                    self.consider_returning_binding(blk, second_ty, err)
+                } else {
+                    false
+                };
+            if !suggested
+                && let Some(second_id) = second_id
+                && let hir::Node::Block(blk) = self.tcx.hir().get(second_id)
+            {
+                self.consider_returning_binding(blk, first_ty, err);
+            }
+        }
+    }
+
     fn suggest_boxing_for_return_impl_trait(
         &self,
         err: &mut Diagnostic,
diff --git a/compiler/rustc_middle/src/traits/mod.rs b/compiler/rustc_middle/src/traits/mod.rs
index 7a38c800ccf..c55971557fa 100644
--- a/compiler/rustc_middle/src/traits/mod.rs
+++ b/compiler/rustc_middle/src/traits/mod.rs
@@ -488,12 +488,15 @@ impl<'tcx> ty::Lift<'tcx> for StatementAsExpression {
 
 #[derive(Clone, Debug, PartialEq, Eq, Hash, Lift)]
 pub struct MatchExpressionArmCause<'tcx> {
+    pub arm_block_id: Option<hir::HirId>,
+    pub arm_ty: Ty<'tcx>,
     pub arm_span: Span,
+    pub prior_arm_block_id: Option<hir::HirId>,
+    pub prior_arm_ty: Ty<'tcx>,
+    pub prior_arm_span: Span,
     pub scrut_span: Span,
-    pub semi_span: Option<(Span, StatementAsExpression)>,
     pub source: hir::MatchSource,
     pub prior_arms: Vec<Span>,
-    pub last_ty: Ty<'tcx>,
     pub scrut_hir_id: hir::HirId,
     pub opt_suggest_box_span: Option<Span>,
 }
diff --git a/compiler/rustc_typeck/src/check/_match.rs b/compiler/rustc_typeck/src/check/_match.rs
index 00547a6d827..2671f2f4f89 100644
--- a/compiler/rustc_typeck/src/check/_match.rs
+++ b/compiler/rustc_typeck/src/check/_match.rs
@@ -9,7 +9,6 @@ use rustc_span::Span;
 use rustc_trait_selection::traits::query::evaluate_obligation::InferCtxtExt;
 use rustc_trait_selection::traits::{
     IfExpressionCause, MatchExpressionArmCause, ObligationCause, ObligationCauseCode,
-    StatementAsExpression,
 };
 
 impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
@@ -75,8 +74,8 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
         };
 
         let mut other_arms = vec![]; // Used only for diagnostics.
-        let mut prior_arm_ty = None;
-        for (i, arm) in arms.iter().enumerate() {
+        let mut prior_arm = None;
+        for arm in arms {
             if let Some(g) = &arm.guard {
                 self.diverges.set(Diverges::Maybe);
                 match g {
@@ -96,21 +95,28 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
 
             let opt_suggest_box_span = self.opt_suggest_box_span(arm_ty, orig_expected);
 
-            let (arm_span, semi_span) =
-                self.get_appropriate_arm_semicolon_removal_span(&arms, i, prior_arm_ty, arm_ty);
-            let (span, code) = match i {
+            let (arm_block_id, arm_span) = if let hir::ExprKind::Block(blk, _) = arm.body.kind {
+                (Some(blk.hir_id), self.find_block_span(blk))
+            } else {
+                (None, arm.body.span)
+            };
+
+            let (span, code) = match prior_arm {
                 // The reason for the first arm to fail is not that the match arms diverge,
                 // but rather that there's a prior obligation that doesn't hold.
-                0 => (arm_span, ObligationCauseCode::BlockTailExpression(arm.body.hir_id)),
-                _ => (
+                None => (arm_span, ObligationCauseCode::BlockTailExpression(arm.body.hir_id)),
+                Some((prior_arm_block_id, prior_arm_ty, prior_arm_span)) => (
                     expr.span,
                     ObligationCauseCode::MatchExpressionArm(Box::new(MatchExpressionArmCause {
+                        arm_block_id,
                         arm_span,
+                        arm_ty,
+                        prior_arm_block_id,
+                        prior_arm_ty,
+                        prior_arm_span,
                         scrut_span: scrut.span,
-                        semi_span,
                         source: match_src,
                         prior_arms: other_arms.clone(),
-                        last_ty: prior_arm_ty.unwrap(),
                         scrut_hir_id: scrut.hir_id,
                         opt_suggest_box_span,
                     })),
@@ -139,7 +145,7 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
                             let ret_ty = ret_coercion.borrow().expected_ty();
                             let ret_ty = self.inh.infcx.shallow_resolve(ret_ty);
                             self.can_coerce(arm_ty, ret_ty)
-                                && prior_arm_ty.map_or(true, |t| self.can_coerce(t, ret_ty))
+                                && prior_arm.map_or(true, |(_, t, _)| self.can_coerce(t, ret_ty))
                                 // The match arms need to unify for the case of `impl Trait`.
                                 && !matches!(ret_ty.kind(), ty::Opaque(..))
                         }
@@ -181,7 +187,8 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
             if other_arms.len() > 5 {
                 other_arms.remove(0);
             }
-            prior_arm_ty = Some(arm_ty);
+
+            prior_arm = Some((arm_block_id, arm_ty, arm_span));
         }
 
         // If all of the arms in the `match` diverge,
@@ -207,32 +214,6 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
         match_ty
     }
 
-    fn get_appropriate_arm_semicolon_removal_span(
-        &self,
-        arms: &'tcx [hir::Arm<'tcx>],
-        i: usize,
-        prior_arm_ty: Option<Ty<'tcx>>,
-        arm_ty: Ty<'tcx>,
-    ) -> (Span, Option<(Span, StatementAsExpression)>) {
-        let arm = &arms[i];
-        let (arm_span, mut semi_span) = if let hir::ExprKind::Block(blk, _) = &arm.body.kind {
-            (
-                self.find_block_span(blk),
-                prior_arm_ty
-                    .and_then(|prior_arm_ty| self.could_remove_semicolon(blk, prior_arm_ty)),
-            )
-        } else {
-            (arm.body.span, None)
-        };
-        if semi_span.is_none() && i > 0 {
-            if let hir::ExprKind::Block(blk, _) = &arms[i - 1].body.kind {
-                let semi_span_prev = self.could_remove_semicolon(blk, arm_ty);
-                semi_span = semi_span_prev;
-            }
-        }
-        (arm_span, semi_span)
-    }
-
     /// When the previously checked expression (the scrutinee) diverges,
     /// warn the user about the match arms being unreachable.
     fn warn_arms_when_scrutinee_diverges(&self, arms: &'tcx [hir::Arm<'tcx>]) {
diff --git a/src/test/ui/suggestions/return-bindings.rs b/src/test/ui/suggestions/return-bindings.rs
index 80c83a70d50..fa1bad37699 100644
--- a/src/test/ui/suggestions/return-bindings.rs
+++ b/src/test/ui/suggestions/return-bindings.rs
@@ -24,4 +24,28 @@ fn d(opt_str: Option<String>) {
     };
 }
 
+fn d2(opt_str: Option<String>) {
+    let s = if let Some(s) = opt_str {
+    } else {
+        String::new()
+        //~^ ERROR `if` and `else` have incompatible types
+    };
+}
+
+fn e(opt_str: Option<String>) {
+    let s: String = match opt_str {
+        Some(s) => {}
+        //~^ ERROR mismatched types
+        None => String::new(),
+    };
+}
+
+fn e2(opt_str: Option<String>) {
+    let s = match opt_str {
+        Some(s) => {}
+        None => String::new(),
+        //~^ ERROR `match` arms have incompatible types
+    };
+}
+
 fn main() {}
diff --git a/src/test/ui/suggestions/return-bindings.stderr b/src/test/ui/suggestions/return-bindings.stderr
index 53ef7106fa8..c14fb336773 100644
--- a/src/test/ui/suggestions/return-bindings.stderr
+++ b/src/test/ui/suggestions/return-bindings.stderr
@@ -59,6 +59,52 @@ LL +         s
 LL ~
    |
 
-error: aborting due to 4 previous errors
+error[E0308]: `if` and `else` have incompatible types
+  --> $DIR/return-bindings.rs:30:9
+   |
+LL |       let s = if let Some(s) = opt_str {
+   |  ______________________________________-
+LL | |     } else {
+   | |_____- expected because of this
+LL |           String::new()
+   |           ^^^^^^^^^^^^^ expected `()`, found struct `String`
+   |
+help: consider returning the local binding `s`
+   |
+LL ~     let s = if let Some(s) = opt_str {
+LL +         s
+LL ~     } else {
+   |
+
+error[E0308]: mismatched types
+  --> $DIR/return-bindings.rs:37:20
+   |
+LL |         Some(s) => {}
+   |                    ^^ expected struct `String`, found `()`
+   |
+help: consider returning the local binding `s`
+   |
+LL |         Some(s) => { s }
+   |                      +
+
+error[E0308]: `match` arms have incompatible types
+  --> $DIR/return-bindings.rs:46:17
+   |
+LL |       let s = match opt_str {
+   |  _____________-
+LL | |         Some(s) => {}
+   | |                    -- this is found to be of type `()`
+LL | |         None => String::new(),
+   | |                 ^^^^^^^^^^^^^ expected `()`, found struct `String`
+LL | |
+LL | |     };
+   | |_____- `match` arms have incompatible types
+   |
+help: consider returning the local binding `s`
+   |
+LL |         Some(s) => { s }
+   |                      +
+
+error: aborting due to 7 previous errors
 
 For more information about this error, try `rustc --explain E0308`.