Address review comments

Co-authored-by: Léo Lanteri Thauvin <leseulartichaut@gmail.com>

3 files changed, 23 insertions, 14 deletions

diff --git a/compiler/rustc_typeck/src/check/demand.rs b/compiler/rustc_typeck/src/check/demand.rs
index 8107cfbf32d..2d84e7830c2 100644
--- a/compiler/rustc_typeck/src/check/demand.rs
+++ b/compiler/rustc_typeck/src/check/demand.rs
@@ -366,16 +366,10 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
         false
     }
 
-    crate fn hir_id_sole_block_element(
-        &self,
-        hir_id: hir::HirId,
-    ) -> Option<&'tcx rustc_hir::Expr<'tcx>> {
-        let node: Option<Node<'_>> = self.tcx.hir().find(hir_id);
-        match node {
-            Some(Node::Expr(rustc_hir::Expr {
-                kind: rustc_hir::ExprKind::Block(block, ..),
-                ..
-            })) if block.stmts.len() == 0 => block.expr,
+    /// If the given `HirId` corresponds to a block with a trailing expression, return that expression
+    crate fn maybe_get_block_expr(&self, hir_id: hir::HirId) -> Option<&'tcx hir::Expr<'tcx>> {
+        match self.tcx.hir().find(hir_id)? {
+            Node::Expr(hir::Expr { kind: hir::ExprKind::Block(block, ..), .. }) => block.expr,
             _ => None,
         }
     }
@@ -666,9 +660,7 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
                                 };
                                 let suggestion = if is_struct_pat_shorthand_field {
                                     format!("{}: *{}", code, code)
-                                } else if let Some(expr) =
-                                    self.hir_id_sole_block_element(expr.hir_id)
-                                {
+                                } else if let Some(expr) = self.maybe_get_block_expr(expr.hir_id) {
                                     if let Ok(inner_code) = sm.span_to_snippet(expr.span) {
                                         format!("*{}", inner_code)
                                     } else {
diff --git a/src/test/ui/deref-suggestion.rs b/src/test/ui/deref-suggestion.rs
index 5cb680541dd..e62e8467bd2 100644
--- a/src/test/ui/deref-suggestion.rs
+++ b/src/test/ui/deref-suggestion.rs
@@ -55,4 +55,12 @@ fn main() {
         b
         //~^ ERROR mismatched types
     };
+    let val: i32 = if true {
+        let _ = 2;
+        a + 1
+    } else {
+        let _ = 2;
+        b
+        //~^ ERROR mismatched types
+    };
 }
diff --git a/src/test/ui/deref-suggestion.stderr b/src/test/ui/deref-suggestion.stderr
index 6eef06b68cb..0de125695e6 100644
--- a/src/test/ui/deref-suggestion.stderr
+++ b/src/test/ui/deref-suggestion.stderr
@@ -98,6 +98,15 @@ LL |         b
    |         expected `i32`, found `&{integer}`
    |         help: consider dereferencing the borrow: `*b`
 
-error: aborting due to 11 previous errors
+error[E0308]: mismatched types
+  --> $DIR/deref-suggestion.rs:63:9
+   |
+LL |         b
+   |         ^
+   |         |
+   |         expected `i32`, found `&{integer}`
+   |         help: consider dereferencing the borrow: `*b`
+
+error: aborting due to 12 previous errors
 
 For more information about this error, try `rustc --explain E0308`.